Tutorials › Econometrics (Graduate) › Cross-Fitting Nuisance Models for Heterogeneity

Treatment Effect Heterogeneity · Tutorial 837 of 952

Cross-Fitting Nuisance Models for Heterogeneity

Learn to construct out-of-fold treatment-effect scores, understand why cross-fitting can reduce overfitting bias, and distinguish it from honest forest estimation.

Advanced 10 min read

What You'll Learn

  • Define fold-specific nuisance predictions and out-of-fold treatment-effect scores
  • Apply the augmented inverse-probability-weighted score for a binary treatment
  • Explain how orthogonality makes nuisance-estimation errors enter jointly
  • Work through a two-fold calculation with hypothetical observations
  • Distinguish cross-fitting nuisance models from forest honesty
  • Identify the assumptions and limitations of cross-fitting

Why Cross-Fit Nuisance Models?

A conditional average treatment effect (CATE) model is often built from a score that combines an outcome model and a treatment-assignment model. These are nuisance models: they help construct a score for the effect of interest, but are not themselves the target. Flexible fitting can make their predictions adapt closely to the same observations whose scores they help construct. The resulting score may then reflect overfitting as well as treatment-effect information.

Cross-fitting separates nuisance-model training from score construction for each observation. Each observation receives predictions from models fitted without that observation. This separation can limit overfitting bias while allowing flexible nuisance estimators. It does not, by itself, establish causal identification, guarantee accurate nuisance predictions, or eliminate all finite-sample bias.

Definition: Cross-fitting divides the sample into folds, fits nuisance models using observations outside a fold, and uses those fitted models to predict nuisance quantities for observations inside that fold. A prediction used to construct an observation's score is therefore out of fold: that observation was not used to fit the nuisance model producing the prediction.

The Score and Its Assumptions

Let \(W\) be a binary treatment, \(Y\) the observed outcome, and \(X\) a vector of observed covariates. Write \(Y(1)\) and \(Y(0)\) for potential outcomes under treatment and control. The target conditional effect is \(\tau(x)=E[Y(1)-Y(0)\mid X=x]\). Define the outcome regressions \(\mu_w(x)=E[Y\mid W=w,X=x]\), for \(w\in\{0,1\}\), and the propensity score \(e(x)=P(W=1\mid X=x)\).

For a causal interpretation using these observed-data regressions, assume consistency, conditional exchangeability \((Y(1),Y(0))\perp W\mid X\), and overlap \(0<e(X)<1\) on the covariate region of interest. Consistency means the observed outcome equals the potential outcome under the treatment actually received. Exchangeability rules out unmeasured confounding conditional on \(X\); overlap ensures both treatment states are possible at each relevant \(X\). Under consistency and exchangeability, \(\mu_1(x)-\mu_0(x)=\tau(x)\). Cross-fitting does not make these assumptions true.

A commonly used orthogonal score for the CATE is the augmented inverse-probability-weighted score. With nuisance functions \(\mu_1,\mu_0,e\), define

$$ \psi(Y,W,X;\eta) = \mu_1(X)-\mu_0(X) +\frac{W}{e(X)}\{Y-\mu_1(X)\} -\frac{1-W}{1-e(X)}\{Y-\mu_0(X)\}, $$

where \(\eta=(\mu_1,\mu_0,e)\) denotes the collection of nuisance functions. If these functions are the true conditional means and propensity score, then \(E[\psi\mid X=x]=\tau(x)\). The residual terms adjust the difference in fitted outcome means using observed outcomes, with weights determined by treatment probabilities.

In practice, use estimates \(\hat\mu_1,\hat\mu_0,\hat e\) in this formula. For an observation in fold \(k\), the estimates are fitted on the other folds and evaluated at that observation's \(X\). Denote these predictions by \(\hat\mu_{1,-k}(X_i)\), \(\hat\mu_{0,-k}(X_i)\), and \(\hat e_{-k}(X_i)\). The resulting \(\hat\psi_i\) is then an outcome for a subsequent CATE learner, such as a regression of scores on \(X\).

A Two-Fold Score Calculation

Worked Example: Constructing Out-of-Fold Scores

Consider four hypothetical observations assigned to two folds. The nuisance predictions shown for each fold were produced using only the other fold. For simplicity, suppose the models return constant predictions within each fold. The scores use the formula above.

ObservationFold\(W\)\(Y\)\(\hat\mu_1\)\(\hat\mu_0\)\(\hat e\)
A11121070.40
B1051070.40
C219860.50
D204860.50

For A, the fitted mean difference is \(10-7=3\), and the treated residual contribution is \((12-10)/0.40=5\). Thus

$$ \hat\psi_A=3+\frac{12-10}{0.40}=3+5=8. $$

For B, \(W=0\), so the treated residual contribution is zero. The control contribution is subtracted:

$$ \hat\psi_B=3-\frac{5-7}{1-0.40} =3-\frac{-2}{0.60} =6.333. $$

For C, the fitted mean difference is \(8-6=2\). Its treated residual contribution is \((9-8)/0.50=2\), giving \(\hat\psi_C=4\). For D, the control residual is \(4-6=-2\), so \(\hat\psi_D=2-(-2)/0.50=6\).

The predictions for A and B come from models trained on C and D; those for C and D come from models trained on A and B. The fold-specific training step, not the particular numerical score, is what makes these predictions cross-fitted. Four observations are far too few to establish reliable effects or nuisance models; this calculation illustrates the workflow, not a recommended sample size.

Why Orthogonality Helps

The score is called orthogonal because small perturbations in the nuisance functions do not generally create a first-order change in its conditional expectation at the truth. This property is one reason cross-fitting is useful with flexible nuisance estimators. The following calculation shows the relevant product structure.

Fix \(X=x\), and suppose the nuisance predictions are held fixed independently of the evaluation observation. Let \(\hat e=\hat e(x)\), \(\hat\mu_w=\hat\mu_w(x)\), and define errors \(\delta_e=\hat e-e(x)\) and \(\delta_w=\hat\mu_w-\mu_w(x)\). Conditional on \(X=x\), the expected score with these predictions is

$$ \begin{aligned} E[\hat\psi\mid X=x] &=\hat\mu_1-\hat\mu_0 +\frac{e(x)}{\hat e}\{\mu_1(x)-\hat\mu_1\} -\frac{1-e(x)}{1-\hat e}\{\mu_0(x)-\hat\mu_0\}\\ &=\tau(x) +\delta_e\left\{\frac{\delta_1}{\hat e} +\frac{\delta_0}{1-\hat e}\right\}. \end{aligned} $$

To see the last equality, the treated-mean error contributes \(\delta_1\{1-e(x)/\hat e\}=\delta_1\delta_e/\hat e\). The control-mean error contributes \(\delta_0\{(1-e(x))/(1-\hat e)-1\}=\delta_0\delta_e/(1-\hat e)\). Adding these to \(\mu_1(x)-\mu_0(x)\) gives the displayed expression. The conditional bias is therefore a product involving propensity error and outcome-regression error, rather than a first-order term in either error alone.

This calculation is the orthogonality argument, not a guarantee of good performance under arbitrary estimation. If propensity predictions approach zero or one, the denominators can amplify errors. Useful results also require appropriate overlap, controlled nuisance-estimation errors, and adequate moments of the score. Cross-fitting helps control the dependence between fitted nuisance functions and evaluation observations; it does not ensure these other conditions.

Worked Example: Reading the Product Error

Suppose at a particular covariate value the true propensity is \(e=0.50\), while the fitted propensity is \(\hat e=0.54\). Let the fitted treated mean be one outcome unit too high, so \(\delta_1=1\), and the fitted control mean be half an outcome unit too low, so \(\delta_0=-0.50\). The conditional score bias from the formula is

$$ (\hat e-e)\left\{\frac{\delta_1}{\hat e} +\frac{\delta_0}{1-\hat e}\right\} = 0.04\left\{\frac{1}{0.54}-\frac{0.50}{0.46}\right\} \approx 0.0306. $$

The bias is approximately \(0.0306\) outcome units at this covariate value. It is not zero, but it is a product of errors: if the propensity prediction were exact, this conditional bias expression would be zero even with these outcome-model errors. Likewise, if both outcome predictions were exact, the expression would be zero despite propensity error, provided the fitted propensity remains away from zero and one. This is a population conditional calculation; a realized score for an individual will also vary because of the observed outcome.

The Cross-Fitting Workflow

1
Partition the observations.
Assign the \(n\) observations to \(K\) folds. In ordinary \(K\)-fold cross-fitting, each observation belongs to exactly one evaluation fold.
2
Fit nuisance models outside each fold.
For fold \(k\), fit the outcome regressions and propensity model using observations not in fold \(k\). Use the same exclusion rule for every nuisance prediction used in that fold's score.
3
Predict and construct scores inside the fold.
Predict \(\hat\mu_{1,-k}(X_i)\), \(\hat\mu_{0,-k}(X_i)\), and \(\hat e_{-k}(X_i)\) for each \(i\) in fold \(k\), then substitute them into the score formula.
4
Combine out-of-fold scores.
After repeating the process for every fold, each observation has one score based on nuisance predictions fitted without it. Use these scores for the intended CATE-learning step, while preserving the distinction between estimated scores and the underlying conditional effect.

Fold assignments should be made before fitting the fold-specific nuisance models. In data with clusters or repeated observations, the split should respect the independent sampling units: observations from the same cluster should not be divided across training and evaluation folds if doing so would leak cluster-specific information. The exact split design depends on the sampling structure and target population.

Cross-Fitting Is Not Forest Honesty

Cross-fitting and forest honesty both involve sample separation, but they separate data for different purposes. Cross-fitting asks whether an observation helped fit the nuisance predictions used to construct its score. Forest honesty asks whether the same observations are used both to choose a tree's partition and to estimate effects within its leaves. These are distinct design choices, and using one does not automatically implement the other.

ProcedureWhat is separated?Primary purpose
Cross-fittingNuisance-model training observations and score-evaluation observationsLimit bias from evaluating nuisance-based scores on observations used to fit the nuisance models
Forest honestyObservations used to determine a forest partition and observations used to estimate effects in that partitionReduce adaptive reuse of outcomes when learning and estimating forest structure

A workflow may cross-fit nuisance models and then use their scores in a learner; whether that learner is also honest is a separate question. Conversely, a forest can use an honest partitioning scheme without using cross-fitted nuisance models. Do not infer one property from the label or design of the other.

Worked Example: Diagnosing In-Sample Leakage

Suppose a hypothetical flexible treated-outcome model is fitted on an observation A itself and memorizes its treated outcome \(Y=12\), predicting \(\hat\mu_1=12\). Let the model predict \(\hat\mu_0=7\), and let \(\hat e=0.50\). If this in-sample prediction is used for A, the score is

$$ \hat\psi_A=(12-7)+\frac{12-12}{0.50}=5. $$

The treated residual correction is zero because the model has fitted that observation's outcome exactly. If an out-of-fold model instead predicted \(\hat\mu_1=10\), with the other predictions unchanged, the score would be

$$ \hat\psi_A=(10-7)+\frac{12-10}{0.50}=3+4=7. $$

The example does not establish that the second score is always better or that the difference must have this direction. It shows how using an observation to fit a flexible nuisance model can change its own residual and hence its score. Cross-fitting prevents this particular form of own-observation fitting; it does not prevent overfitting relationships that recur across other observations or errors caused by poor model specification.

Common Pitfalls and Practical Checks

  • Treating cross-fitting as causal identification. Sample splitting cannot remove unmeasured confounding. The causal interpretation still depends on consistency, conditional exchangeability, and overlap.
  • Assuming out-of-fold means unbiased. Cross-fitting limits a source of overfitting bias, but finite-sample bias and nuisance-estimation error can remain. The product-error calculation also relies on denominators that are not too small.
  • Using the wrong predictions for a fold. If an observation's nuisance model was fitted using that observation, its score is not out of fold. Keep fold membership and training exclusions explicit for every nuisance model.
  • Confusing score with individual effect. A score is a constructed, noisy outcome whose conditional expectation targets the CATE under the stated assumptions. It is not an observed individual treatment effect.
  • Confusing cross-fitting with honesty. Cross-fitting separates nuisance fitting from score evaluation; forest honesty separates partition choice from within-partition effect estimation.

The practical value of cross-fitting is clearest when nuisance models are flexible enough that in-sample predictions could be overly adapted to their training observations. A disciplined fold workflow gives each score a more defensible out-of-sample nuisance input. Whether this improves a particular analysis depends on sample size, fold design, overlap, nuisance accuracy, and the downstream CATE learner.

Key takeaway: Cross-fitting assigns each observation a treatment-effect score built from nuisance predictions fitted without that observation. This reduces own-observation overfitting concerns and works naturally with an orthogonal score, but does not replace causal assumptions, overlap checks, accurate nuisance estimation, or separate decisions about forest honesty.

Check Your Understanding

Answer each question using the definitions and workflow above.

  1. In two-fold cross-fitting, which observations are used to fit nuisance models that predict for observations in fold 1?
  2. State the three causal assumptions used to interpret the observed-data CATE as a treatment effect, and explain why cross-fitting does not establish them.
  3. In the conditional bias expression, what happens if the propensity prediction is exact but the outcome-regression predictions are not?
  4. Why is an out-of-fold score not an observed individual treatment effect?
  5. What is the distinction between cross-fitting nuisance models and forest honesty?
Clinical Biostats