Why Cross-Fit AIPW?
AIPW combines treatment-specific outcome predictions with propensity-score weighting to estimate the average treatment effect (ATE). Flexible learners can capture nonlinearities and interactions, but their predictions are adapted to the data used to fit them. If an observation helps train a learner and is then used to evaluate that learner’s AIPW contribution, overfitting can make the contribution misleading. Cross-fitting separates nuisance-model training from score evaluation for each observation.
The goal is not to make a weak learner strong or to fix confounding. It is to reduce the direct dependence between an observation and the nuisance predictions used in its score, while allowing data-adaptive learners to be used. The usual causal assumptions are still needed: consistency and no interference, conditional exchangeability given measured covariates, and positivity. In particular, positivity requires the treatment probability to be bounded away from 0 and 1 in the target population. Observations are also taken to be independent and identically distributed for the basic setup here.
Constructing Cross-Fit Predictions
Partition the \(n\) observations into \(K\) folds, indexed by \(k=1,\ldots,K\). For each fold, fit the propensity learner and both treatment-specific outcome learners using observations outside that fold. Then predict \(e(X_i)\), \(m_1(X_i)\), and \(m_0(X_i)\) for every observation \(i\) in the held-out fold. Write these predictions as \(\hat e_{-k}(X_i)\) and \(\hat m_{a,-k}(X_i)\): the subscript \(-k\) indicates that fold \(k\) was excluded from training.
For each observation, substitute its held-out predictions into the AIPW score. If \(I_k\) is the set of observation indices in fold \(k\), the cross-fit estimate is
Each observation appears once in the evaluation sum, even though each fitted nuisance model is used to predict an entire held-out fold. The same folds are commonly used for all three nuisance functions so that the predictions combined in an observation’s score are all based on training data that excluded that observation. Flexible learners may include regression trees, regularized models, or ensembles; the defining feature is the separation between training and evaluation, not a particular algorithm.
Averaging all \(n\) contributions is essential. If fold sizes differ, average the scores with weights proportional to the number of observations in each fold, rather than giving every fold mean equal weight. For standard-error calculations, the empirical variability of the cross-fit score contributions is commonly used; as with other influence-function-based inference, its validity requires appropriate regularity and nuisance-estimation conditions. Cross-fitting by itself does not establish those conditions.
Assign each observation to exactly one held-out fold. Record the fold assignment so predictions can be linked to the correct observations.
For fold \(k\), fit \(\hat e_{-k}\), \(\hat m_{1,-k}\), and \(\hat m_{0,-k}\) using only observations not in \(I_k\).
Use those fits to predict nuisance values for observations in \(I_k\), then calculate one AIPW score contribution per observation.
Sum all held-out contributions and divide by \(n\). Inspect prediction quality and score behavior before interpreting the estimate.
Worked Calculations
Worked Example: Calculating Scores Across Two Folds
Consider four hypothetical observations divided into two folds. The nuisance predictions shown below were produced by fits trained on the other fold. Outcomes are measured in points, so the AIPW estimate is also in points.
| Observation | Fold | \(A_i\) | \(Y_i\) | \(\hat e_{-k}(X_i)\) | \(\hat m_{1,-k}(X_i)\) | \(\hat m_{0,-k}(X_i)\) |
|---|---|---|---|---|---|---|
| 1 | 1 | 1 | 8 | 0.50 | 7 | 5 |
| 2 | 1 | 0 | 4 | 0.40 | 6 | 4.5 |
| 3 | 2 | 1 | 10 | 0.60 | 8 | 6 |
| 4 | 2 | 0 | 5 | 0.50 | 7 | 5.5 |
For observation 1, \(A_1=1\), so the untreated residual term is zero:
For observation 2, \(A_2=0\), so the treated residual term is zero:
The other two contributions are
The cross-fit estimate is the mean of the four held-out contributions:
This is an estimate of the ATE under the stated causal assumptions; it is not a claim that the treatment caused a 3.542-point change in this small hypothetical data set. The calculations show how predictions from the appropriate training fold enter each observation’s score.
Worked Example: Weighting Fold Means Correctly
Suppose fold 1 contains two observations with scores 1 and 3, while fold 2 contains three observations with scores 2, 4, and 6. The fold means are 2 and 4, respectively. The overall estimate is the mean over all five observations:
Simply averaging the fold means would give \((2+4)/2=3\), which is wrong here because it gives a two-observation fold the same weight as a three-observation fold. Equivalently, the correct calculation weights each fold mean by its fold size.
Why Excluding the Evaluation Fold Matters
The practical justification is direct: when an observation is excluded from nuisance-model training, its outcome cannot directly determine the fitted values used in its own AIPW contribution. This matters especially for flexible learners that can closely adapt to training outcomes. If an observation is used both to fit a learner and to evaluate its residual, an unusually small in-sample residual may reflect overfitting rather than a genuinely informative prediction.
This separation does not make held-out predictions independent of all the data: the training sets for different folds overlap. Nor does it make score contributions across folds independent. Cross-fitting is a construction of the estimator, not a claim that every contribution is statistically independent of every other contribution. Its benefit is the observation-level separation between each evaluation outcome and the nuisance fits used for that evaluation.
Worked Example: How In-Sample Prediction Can Change a Contribution
Take two hypothetical observations with \(A_1=1,Y_1=10\) and \(A_2=0,Y_2=4\). Suppose held-out predictions for both have \(\hat e=0.50\), \(\hat m_1=9\), and \(\hat m_0=5\). The resulting scores are
The estimate from these two scores is 6. Now imagine an in-sample learner that uses each observation’s own outcome and predicts its observed treatment-arm outcome exactly: for observation 1 it predicts \(m_1(X_1)=10\), and for observation 2 it predicts \(m_0(X_2)=4\). Suppose its other-arm predictions remain 5 and 9, respectively. Its scores become
The in-sample calculation gives 5 instead of 6 because the observed-arm residuals have been driven to zero. This tiny example does not establish which estimate is closer to a true ATE; it demonstrates that using an observation to train its own nuisance prediction can alter its score. Cross-fitting prevents that direct form of leakage.
Diagnostics, Alternatives, and Limits
A correct fold loop is necessary but not sufficient. Check that each observation is evaluated exactly once, and verify that every prediction used for fold \(k\) came from a model trained without fold \(k\). Store fold labels or training indices and audit them; a prediction table alone may not reveal accidental leakage. Also inspect propensity predictions for values near 0 or 1, outcome predictions for implausible values, and the distribution of score contributions. A small number of large weighted residuals can dominate the estimate. Compare fold-specific score averages as a stability check, while remembering that fold differences can reflect sampling variation.
Cross-fitting does not resolve unmeasured confounding, violations of positivity, or poor nuisance prediction. A learner can be flexible yet perform badly, particularly when the training data are limited or treatment groups are sparse. Nor does cross-fitting automatically justify confidence intervals: inference relies on assumptions about the sampling process and nuisance estimation beyond the act of splitting. The introductory material on AIPW and the earlier material on estimated propensity scores provide relevant context, but the key operational point here is that all nuisance predictions in an evaluation score must be held out.
| Approach | How nuisance predictions are made | Main trade-off |
|---|---|---|
| In-sample AIPW with flexible learners | Fit and evaluate on the same observations | Uses all data for fitting, but own-observation overfitting can distort score contributions. |
| Single sample split | Train on one part and evaluate on a separate part | Clear separation, but only the evaluation subset contributes to the estimate. |
| Cross-fit AIPW | Rotate held-out folds so each observation is evaluated once | Uses each observation for evaluation and for training other folds’ models, but requires careful fold bookkeeping. |
The main alternative, a single split, also separates fitting from evaluation but generally discards evaluation opportunities for the training subset. Cross-fitting reuses observations across roles in different folds without using an observation to train the particular nuisance fits used for its own score. Compared with parametric full-sample AIPW, it can accommodate more adaptive nuisance learners, but it does not guarantee that those learners are well chosen or that overlap is adequate.
Check Your Understanding
Answer each question, then compare your reasoning with the solutions.
- In a three-fold cross-fit, which observations may be used to train the nuisance models whose predictions are evaluated on fold 2?
- Fold A has four observations and mean score 1.5; fold B has six observations and mean score 2.5. What is the overall mean score?
- Does cross-fitting ensure positivity, or prevent bias from unmeasured confounding?
- Why can a flexible learner’s in-sample residual be misleading in an AIPW score?
- What should be checked if a few observations have very large AIPW contributions?
Solutions. (1) Train on folds 1 and 3, excluding every observation in fold 2. (2) Weight by fold size: \(\{4(1.5)+6(2.5)\}/10=(6+15)/10=2.100\). (3) No. Cross-fitting changes how nuisance predictions are evaluated; it neither guarantees overlap nor removes unmeasured confounding. (4) The learner may have adapted to that observation’s outcome, making the residual artificially small and changing its score contribution. (5) Audit the held-out propensity predictions for values near 0 or 1, verify that the correct fold-specific fits were used, and examine the corresponding residuals and predictions. Large contributions can indicate instability or model problems, not just a coding issue.