Choosing Sample Splits and Cross-Fitting Folds

Why Fold Design Matters

Cross-fitting lets each observation contribute to an AIPW estimate while keeping that observation out of the nuisance-model fits used to calculate its contribution. The fold assignment is therefore part of the estimator’s design, not an administrative detail. A poor choice can leave too little data to fit the nuisance functions, produce evaluation folds with very few treated or untreated observations, or allow information from an evaluation observation to enter its own predictions.

The goal is not to find one universally optimal number of folds. It is to construct training and evaluation sets that support credible nuisance fits and preserve the separation required by cross-fitting. The asymptotic validity of AIPW also depends on the assumptions and nuisance-estimation conditions established in the earlier material on asymptotic linearity, Neyman orthogonality, and product-rate conditions. A fold design cannot repair failures of those conditions.

Fold Construction for Cross-Fitted AIPW

Let the observed data be \(O_i=(Y_i,A_i,W_i)\), \(i=1,\ldots,n\), where \(Y_i\) is the outcome, \(A_i\in\{0,1\}\) is treatment, and \(W_i\) is a vector of measured covariates. For the average treatment effect, write \(e(w)=P(A=1\mid W=w)\) for the propensity score and \(m_a(w)=E(Y\mid A=a,W=w)\) for the treatment-specific outcome regression. Denote their estimates by \(\hat e\) and \(\hat m_a\).

Write \(k(i)\) for the fold containing observation \(i\), and let \(\hat e^{(-k)}\) and \(\hat m_a^{(-k)}\) denote nuisance estimates fitted without observations in \(I_k\). A cross-fitted AIPW estimate is

Each observation is evaluated exactly once in this version of the estimator. If folds have unequal sizes, the outer sum still gives each observation weight \(1/n\); it is not appropriate to average the \(K\) fold-specific means equally unless all folds have the same size.

Choosing the Number and Size of Folds

With \(K\) folds of approximately equal size, each evaluation fold has about \(n/K\) observations and each nuisance fit has about \(n(K-1)/K\) observations. Increasing \(K\) generally gives each nuisance fit more training data, but makes each evaluation fold smaller. Small evaluation folds can have sparse treatment groups, so their AIPW contributions may be highly variable. Conversely, very small \(K\) leaves a larger fraction of the sample out of each training fit.

For a fixed \(K\) as \(n\) grows, the training fraction is \((K-1)/K\) and the evaluation fraction is \(1/K\). Values such as 5 or 10 are common practical choices, not universal requirements. The appropriate choice depends on sample size, treatment prevalence, covariate structure, and how much data the nuisance learners need. The key is to use folds large enough to fit the intended models and to check that each training set includes adequate treated and untreated observations.

Fold assignment should ordinarily be randomized independently of the observed outcomes and covariates, with a recorded random seed so the analysis can be reproduced. If treatment-stratified assignment is used to avoid empty treatment groups, it should be specified in advance and reported. In any case, every observation must be excluded from the nuisance fits used for its own evaluation contribution. Small samples may make it impossible to achieve both substantial training sets and well-populated evaluation folds; this is a limitation to report, not a reason to silently reuse the evaluation outcomes in training.

Why Separation Helps

The basic justification is a conditional-independence argument. Suppose the observations are independent and identically distributed, and the fold assignment is independent of them. For a particular fold \(k\), condition on the training data outside \(I_k\). The fitted nuisance functions are then fixed, and the observations in \(I_k\) remain independent of those fitted functions. Thus, the evaluation observations are not being scored using nuisance fits that have adapted to their own outcomes.

This separation is useful when nuisance functions are estimated flexibly: it limits the empirical dependence between a training observation’s outcome and its own fitted prediction in the score. It does not make the score exactly unbiased in finite samples, nor does it guarantee valid inference. The required identification, overlap, and nuisance-rate conditions still matter.

A short calculation shows why nuisance errors enter the AIPW score through a product under the usual identification conditions. Treat the fitted functions as fixed, and write \(\tilde e\) and \(\tilde m_a\) for generic nuisance functions. Conditional on \(W=w\), the expectation of the treatment-one part of the score is

The treatment-zero part has the analogous product form, with denominator \(1-\tilde e(w)\). Hence the conditional bias is governed by products of propensity-score and outcome-regression errors, rather than either error alone. Cross-fitting supports the empirical-process argument used with this structure; it does not change the product-rate requirements or protect against severe positivity violations.

Repeated Splitting

A single random partition can make an estimate sensitive to which observations happen to share a fold. Repeated cross-fitting addresses this source of split variability by generating \(R\) partitions and averaging the resulting estimates:

Every repetition must independently follow the stated fold rule, and within each repetition each observation must be evaluated using nuisance fits that exclude its fold. Repetition can reduce sensitivity to a particular partition, but it does not create \(R\) independent data sets: the same observations appear in every repetition. Therefore, the variation among the \(R\) estimates is not, by itself, a valid standard error for the treatment effect. Inference must account for the sampling variation of the original observations and the estimator’s influence-function structure; the standard-error construction is not obtained by pretending the repetitions are independent.

Preventing Leakage and Diagnosing Failures

Leakage occurs when information from an evaluation observation enters the nuisance prediction used to score that observation. The obvious case is fitting \(\hat e\) or \(\hat m_a\) on all observations and then labeling the resulting predictions “cross-fitted.” Less obvious cases occur when data-adaptive preprocessing or model selection is performed using the full data before folds are applied.

• Fit preprocessing within each training set. Imputation rules, scaling constants, feature screening, and dimensionality reduction must be learned without the evaluation fold, then applied to it.
• Keep tuning inside the training data. Hyperparameter selection and learner selection should not use evaluation outcomes. Internal validation within the training set is compatible with cross-fitting.
• Track observation membership explicitly. For every fold-specific prediction, verify that its observation identifier is absent from the data used to fit and tune the nuisance models.
• Inspect fold composition and predictions. Check treatment counts, convergence, missing predictions, extreme propensity estimates, and implausible outcome predictions in every fold. A numerical prediction does not prove that the fit was adequate.
• Record the design. Report \(K\), the fold allocation rule, whether splits were repeated, and how tuning and preprocessing were handled.

Alternatives and Limitations

A single holdout split keeps training and evaluation data separate, but discards the evaluation observations from nuisance training and may yield an estimate sensitive to one split. Cross-fitting uses all observations for evaluation contributions while retaining fold-specific separation. Fitting nuisance models on the full sample and evaluating the score on that same sample avoids partitioning, but loses the separation that cross-fitting is designed to provide and can be problematic with highly adaptive learners.

No fold choice guarantees stable estimation. If a treatment group is very small, splitting can leave some training sets with too few observations for reliable nuisance fits. Increasing \(K\) can enlarge training sets but shrink evaluation folds; decreasing \(K\) does the opposite. Repeated splits can show whether the point estimate is sensitive to the partition, but cannot solve a lack of overlap, inadequate sample size, poor nuisance predictions, or violations of identification assumptions. For dependent observations, the independent-observation argument above does not directly apply; fold assignment must respect the sampling structure, and the corresponding inference requires separate treatment.

A common mistake is to choose the largest possible \(K\) on the grounds that more training data must always improve the estimate. That reasoning ignores smaller evaluation folds, sparse treatment groups, and the fact that large \(K\) does not fix poor nuisance learning. Choose \(K\) based on the actual training and evaluation counts, audit the full fitting pipeline, and report the design so its consequences can be assessed.