Heterogeneous effects and causal ML: meta-learners, double ML and causal forests

Before this lesson: potential outcomes and the CATE (lesson 02); conditional exchangeability, positivity, the propensity score and AIPW (lesson 05); regularization (ridge and lasso), cross-validation, decision trees, random forests and boosting (standard supervised learning); projections (linear algebra 03) help for section 3.3. Time: about 110 minutes, plus the notebook. This is the heaviest lesson of the track. The two folded deeper blocks can be skipped on a first reading.

---

Why this matters #

Many decisions are about whom to treat, not whether to treat: which users get a discount, which patients get a drug with side effects, which job seekers get a training program. A model that predicts "who will buy" or "who will churn" answers the wrong question. The right question is "for whom does the treatment change the outcome?". This is uplift modeling in marketing and personalised medicine in health. Tech companies look for effects that differ between user groups in every large A/B test.

The second use is estimating one effect with many covariates. With 100 confounders, a linear regression can be badly misspecified, and a plain ML model gives biased effects (section 3.3). Double/debiased machine learning (Chernozhukov et al. 2018) is now a standard tool in economics, epidemiology and industry: it gives a $\sqrt n$-consistent estimate and a valid confidence interval while gradient boosting or random forests model the confounding. The ideas of this lesson (orthogonal scores, cross-fitting, honest estimation) also appear in off-policy evaluation (lesson 11) and in modern semiparametric statistics.

1. Predict first #

2. The picture #

2.1 The effect is a gap between two curves #

Take one covariate $x$. The outcome regressions are $\mu_0(x) = \mathbb{E}[Y \mid T = 0, X = x]$ (the average outcome of controls with features $x$) and $\mu_1(x) = \mathbb{E}[Y \mid T = 1, X = x]$. Under the assumptions of lesson 05, the CATE is the vertical gap between the two curves:

Why is this hard, even with good ML models? Three reasons.

• The gap is small and simple; the curves are big and complex. In the figure below, the baseline $\mu_0(x)$ goes up and down by about 4 units, while the effect changes slowly between $0.5s$ and $1.5s$. A model that predicts $Y$ well spends its effort on the big wiggles, not on the small gap.
• Errors do not cancel. $\hat\tau = \hat\mu_1 - \hat\mu_0$ carries the error of both models. If $\hat\mu_1$ is fitted on few points, its error is large, and all of it lands in $\hat\tau$.
• Regularization works against the gap. Regularization pulls a model toward simple functions. If $T$ is one input among many, "simple" can mean "ignore $T$", which sets $\hat\tau$ to 0.

A meta-learner is a recipe that turns any regression method (a forest, boosting, a neural network, a kernel smoother) into a CATE estimator. Four common recipes, in words:

• S-learner ("single"): fit one model $\hat\mu(x, t)$ with $T$ as an extra input; then $\hat\tau(x) = \hat\mu(x, 1) - \hat\mu(x, 0)$.
• T-learner ("two"): fit $\hat\mu_1$ on the treated and $\hat\mu_0$ on the controls; subtract.
• X-learner: first a T-learner. Then impute an effect for every unit with the other group's model, fit a model to these effects, and mix the two effect models.
• DR-learner: build a "doubly robust" pseudo-outcome for every unit (the AIPW term of lesson 05) and regress it on $x$.

Interactive figure: loading…

The figure shows the main trade-offs. No learner wins everywhere; section 3.2 gives a table.

2.2 Partialling out: residual on residual #

Now take a different question: one number $\theta$, the effect of $T$ on $Y$, when many covariates $X$ confound it. Linear regression of $Y$ on $T$ and $X$ answers it only if the confounding is linear. The Frisch–Waugh–Lovell (FWL) theorem shows another way to compute the coefficient of $T$:

1. Predict $T$ from $X$ and keep the residual $\tilde T = T - \hat m(X)$: the part of the treatment that $X$ does not explain.
2. Predict $Y$ from $X$ and keep the residual $\tilde Y = Y - \hat \ell(X)$.
3. Regress $\tilde Y$ on $\tilde T$. The slope is the effect.

With linear predictions in steps 1 and 2, the slope is exactly the multiple-regression coefficient. Double machine learning replaces the linear predictions by ML models. That is the whole idea.

Interactive figure: loading…

3. Definitions and main results #

3.1 The CATE and its identification #

Definition. The conditional average treatment effect (CATE) is $\tau(x) = \mathbb{E}[Y(1) - Y(0) \mid X = x]$. The ATE is its average: $\tau_{\mathrm{ATE}} = \mathbb{E}[\tau(X)]$.

The CATE is not the individual effect $\tau_i = Y_i(1) - Y_i(0)$. Two people with the same $x$ can react differently; $\tau(x)$ is their average. With richer $x$, $\tau(x)$ gets closer to individual effects, but it never becomes them.

Assumptions (lesson 05, now used for every $x$):

1. Consistency: $Y = Y(T)$.
2. Conditional exchangeability: $(Y(0), Y(1)) \perp\!\!\!\perp T \mid X$ (no unmeasured confounding).
3. Positivity (overlap): $0 < e(x) < 1$ for every $x$ we want to evaluate.

Theorem (identification of the CATE). Under 1–3, $\tau(x) = \mu_1(x) - \mu_0(x)$.

Proof (for $\mu_1$; $\mu_0$ is the same).

1. Positivity gives $P(T = 1 \mid X = x) > 0$, so $\mu_1(x) = \mathbb{E}[Y \mid T = 1, X = x]$ is defined.
2. Consistency: $\mathbb{E}[Y \mid T = 1, X = x] = \mathbb{E}[Y(1) \mid T = 1, X = x]$.
3. Exchangeability: $\mathbb{E}[Y(1) \mid T = 1, X = x] = \mathbb{E}[Y(1) \mid X = x]$.
4. Subtract the same chain for $t = 0$. $\blacksquare$

Positivity can be checked (look at $\hat e(x)$). Exchangeability cannot: machine learning does not remove a confounder that is not in $X$. Everything below is about estimation: the estimand is fixed by the theorem, and the methods differ in how well they estimate it.

3.2 Meta-learners #

Let $\mathcal{A}$ be any regression method ("fit $Z$ on $W$"). The four learners of section 2, written out:

Learner	Recipe	Good when	Weak when
S	fit $\hat\mu(x, t)$ on all data; $\hat\tau(x) = \hat\mu(x, 1) - \hat\mu(x, 0)$	effect zero or small; little data	real effect, strong regularization: $\hat\tau$ shrinks to 0
T	fit $\hat\mu_1$ on treated, $\hat\mu_0$ on controls; subtract	much data in both groups; $\mu_1, \mu_0$ very different	one group small; $\tau$ simpler than $\mu_t$
X	T-learner, then impute effects, fit $\hat\tau_1, \hat\tau_0$, mix with $\hat e$	unbalanced groups; smooth $\tau$	both groups small; poor $\hat e$
DR	regress the AIPW pseudo-outcome on $x$ (cross-fitted)	good overlap; want robustness	$\hat e(x)$ near 0 or 1: huge pseudo-outcomes

X-learner (Künzel, Sekhon, Bickel and Yu 2019), step by step:

1. Fit $\hat\mu_0$ on the controls and $\hat\mu_1$ on the treated.
2. Impute individual effects. For a treated unit, $D_i^1 = Y_i - \hat\mu_0(X_i)$ (observed outcome minus the predicted outcome without treatment). For a control unit, $D_i^0 = \hat\mu_1(X_i) - Y_i$.
3. Fit $\hat\tau_1$ on the treated with target $D^1$, and $\hat\tau_0$ on the controls with target $D^0$.
4. Mix: $\hat\tau(x) = g(x)\,\hat\tau_0(x) + (1 - g(x))\,\hat\tau_1(x)$ with $g(x) = \hat e(x)$.

Why this weight? With few treated units ($\hat e$ small), $\hat\tau_1$ is the better model: it uses the control model $\hat\mu_0$, fitted on many points, and its target $D^1$ is the smooth effect plus noise. So the weight $1 - \hat e(x)$ on $\hat\tau_1$ is large. With few controls, the roles swap.

DR-learner (Kennedy 2023). Split the data. On one part, fit $\hat\mu_0, \hat\mu_1, \hat e$. On the other part, compute for each unit the pseudo-outcome

and regress $\hat\varphi$ on $X$. Swap the parts and average. Its average over all units is the AIPW estimate of the ATE from lesson 05.

Theorem (the pseudo-outcome is doubly robust). Let $\bar\mu_t, \bar e$ be fixed functions (the fitted models, from the other part of the data). Under assumptions 1–3,

So the target of the regression is exactly $\tau(x)$ if $\bar e = e$, or if both $\bar\mu_t = \mu_t$.

Proof

1. Given $X = x$: $\mathbb{E}[T Y \mid X = x] = e(x)\,\mu_1(x)$ and $\mathbb{E}[T \mid X = x] = e(x)$. So $\mathbb{E}\big[T(Y - \bar\mu_1(X))/\bar e(X) \mid X = x\big] = \frac{e(x)}{\bar e(x)}\big(\mu_1(x) - \bar\mu_1(x)\big)$.
2. In the same way, $\mathbb{E}\big[(1 - T)(Y - \bar\mu_0(X))/(1 - \bar e(X)) \mid X = x\big] = \frac{1 - e(x)}{1 - \bar e(x)}\big(\mu_0(x) - \bar\mu_0(x)\big)$.
3. Add $\bar\mu_1 - \bar\mu_0$ and write $\bar\mu_1 + \frac{e}{\bar e}(\mu_1 - \bar\mu_1) = \mu_1 + (\mu_1 - \bar\mu_1)\frac{e - \bar e}{\bar e}$. For the control part, $\bar\mu_0 + \frac{1 - e}{1 - \bar e}(\mu_0 - \bar\mu_0) = \mu_0 - (\mu_0 - \bar\mu_0)\frac{e - \bar e}{1 - \bar e}$.
4. Subtract: $\mathbb{E}[\varphi \mid X = x] = \mu_1 - \mu_0 + (e - \bar e)\big[\frac{\mu_1 - \bar\mu_1}{\bar e} + \frac{\mu_0 - \bar\mu_0}{1 - \bar e}\big]$, and $\mu_1 - \mu_0 = \tau(x)$ by section 3.1. $\blacksquare$

The error is a product of a propensity error and an outcome-model error. Kennedy (2023) shows that, because of this, the DR-learner can reach the error rate of an "oracle" that regresses the true $\varphi$ on $X$, even when the nuisance models converge more slowly. The same product structure is the key to double ML in the next section.

3.3 Double machine learning for one effect #

The model. The partially linear model (Robinson 1988) is

$g$ and $m$ can be any functions (they carry the confounding). The effect of $T$ is the same number $\theta$ for everybody. $T$ can be binary or continuous (a dose, a price). Under conditional exchangeability, $\theta$ is the causal effect.

Step 1: Frisch–Waugh–Lovell. Take expectations given $X$: $\ell(X) = \mathbb{E}[Y \mid X] = \theta\,m(X) + g(X)$. Subtract this from $Y$:

This is Robinson's transformation: the residual of $Y$ is $\theta$ times the residual of $T$, plus noise. So $\theta$ is the slope of a regression through the origin: $\theta = \mathbb{E}[\tilde T \tilde Y] / \mathbb{E}[\tilde T^2]$, with $\tilde T = T - m(X)$ and $\tilde Y = Y - \ell(X)$. When $m$ and $\ell$ are linear projections, this is the FWL theorem.

FWL theorem for linear regression

Let $W$ be the $n \times p$ matrix of covariates (with a column of ones), $P = W(W^\top W)^{-1}W^\top$ the projection onto its column space (linear algebra 03), and $M = I - P$. OLS of $Y$ on $(T, W)$ gives $Y = \hat\beta T + W\hat\gamma + r$, with a residual $r$ orthogonal to $T$ and to every column of $W$.

1. Multiply by $M$: $MW = 0$, and $Mr = r$ (because $Pr = 0$). So $MY = \hat\beta\,MT + r$.
2. Multiply by $(MT)^\top$: $(MT)^\top MY = \hat\beta\,(MT)^\top MT + T^\top M r$.
3. $T^\top M r = T^\top r - T^\top P r = 0 - 0$. So $\hat\beta = \dfrac{(MT)^\top (MY)}{(MT)^\top (MT)}$.
4. $MT = \tilde T$ and $MY = \tilde Y$ are the residuals of $T$ and $Y$ on $W$. So $\hat\beta$ is the slope of $\tilde Y$ on $\tilde T$. $\blacksquare$

Step 2: replace the projections by ML models. Fit $\hat m(X) \approx \mathbb{E}[T \mid X]$ and $\hat\ell(X) \approx \mathbb{E}[Y \mid X]$ with any ML method. Both are ordinary prediction problems, so they can be tuned by cross-validation.

Step 3: the DML estimator with cross-fitting (Chernozhukov et al. 2018):

1. Split the $n$ units at random into $K$ folds (often $K = 5$).
2. For each fold $k$: fit $\hat\ell^{(-k)}$ and $\hat m^{(-k)}$ on the other folds, and compute the residuals of the units in fold $k$: $\tilde Y_i = Y_i - \hat\ell^{(-k)}(X_i)$, $\tilde T_i = T_i - \hat m^{(-k)}(X_i)$.
3. Pool all units:

4. The 95% confidence interval is $\hat\theta \pm 1.96\,\widehat{\mathrm{SE}}$. In practice, repeat with several random splits and report the median.

Why not just plug in an ML model? A natural "plug-in" estimator fits $\hat g$ with ML and then regresses $Y - \hat g(X)$ on $T$. Its error contains the term $\frac{1}{n}\sum_i T_i\,(g - \hat g)(X_i) \approx \mathbb{E}[m(X)\,(g - \hat g)(X)]$: the error of $\hat g$ enters directly (first order), multiplied by the confounding $m(X)$. ML models are regularized, so $\hat g - g$ shrinks more slowly than $1/\sqrt n$, and this regularization bias dominates the statistical noise. The confidence interval then misses $\theta$.

Why DML is different: errors multiply. Write the nuisance errors as $\Delta_m = \hat m - m$ and $\Delta_\ell = \hat\ell - \ell$. Then $\tilde T = v - \Delta_m$ and, from Robinson's transformation, $\tilde Y - \theta\tilde T = \varepsilon - \Delta_\ell + \theta\,\Delta_m$. So

• The first term gives the normal distribution of $\hat\theta$, with standard error of order $1/\sqrt n$.
• The middle terms have mean zero if the nuisance models were fitted on other data, because $\mathbb{E}[v \mid X] = 0$ and $\mathbb{E}[\varepsilon \mid X] = 0$. This is what cross-fitting buys.
• The last term is the only bias, and it is a product of two errors. If each error is about $10\%$, the bias is about $1\%$.

Interactive figure: loading…

Neyman orthogonality. The DML estimator solves the moment equation $\frac1n \sum_i \psi(W_i; \hat\theta, \hat\eta) = 0$ with the score $\psi(W; \theta, \eta) = \big(Y - \ell(X) - \theta(T - m(X))\big)\big(T - m(X)\big)$ and nuisance $\eta = (\ell, m)$. A score is Neyman orthogonal if, at the true values,

In words: a small change in the nuisance functions does not change the moment condition to first order.

Proof that the DML score is orthogonal. Take a direction $(\delta_\ell, \delta_m)$, functions of $X$.

1. Differentiate inside the expectation at $r = 0$: $\mathbb{E}\big[(-\delta_\ell + \theta_0\,\delta_m)(X)\,(T - m(X))\big] + \mathbb{E}\big[(Y - \ell(X) - \theta_0(T - m(X)))\,(-\delta_m(X))\big]$.
2. First term: condition on $X$; $\mathbb{E}[T - m(X) \mid X] = \mathbb{E}[v \mid X] = 0$. So it is 0.
3. Second term: $Y - \ell - \theta_0(T - m) = \varepsilon$ by Robinson's transformation, and $\mathbb{E}[\varepsilon\,\delta_m(X)] = \mathbb{E}\big[\delta_m(X)\,\mathbb{E}[\varepsilon \mid X]\big] = 0$. $\blacksquare$

The plug-in score $\psi = (Y - \theta T - g(X))\,T$ is not orthogonal: its derivative in the direction $\delta_g$ is $-\mathbb{E}[\delta_g(X)\,T] = -\mathbb{E}[\delta_g(X)\,m(X)]$, which is not 0 when $T$ is confounded. This is the regularization bias above.

Cross-fitting. If $\hat\ell$ was fitted on unit $i$ itself, it has partly learned unit $i$'s noise $\theta v_i + \varepsilon_i$. Then $\Delta_\ell(X_i)$ is correlated with $v_i$, the middle term $-v_i\,\Delta_\ell(X_i)$ no longer has mean zero, and $\hat\theta$ is biased (toward 0 in the figure). Fitting the nuisances on other folds removes this dependence. Cross-fitting does not waste data: every unit is used for the final estimate, and every unit helps fit the nuisance models of the other folds. (This is the same idea as cross-validation: never evaluate a model on the points it was trained on.)

Theorem (DML, informal; Chernozhukov et al. 2018, Theorem 4.1). In the partially linear model, assume $\mathbb{E}[v^2] > 0$, bounded moments, cross-fitting, and nuisance errors (in mean square) that satisfy

Then $\sqrt n\,(\hat\theta - \theta) \to \mathcal N(0, \sigma^2)$ with $\sigma^2 = \mathbb{E}[v^2\varepsilon^2] / (\mathbb{E}[v^2])^2$, and the interval of step 4 has coverage $95\%$ in large samples.

The rate condition holds, for example, if each nuisance error is $o(n^{-1/4})$: much slower than the $n^{-1/2}$ of the final estimate. Random forests, boosting and lasso reach such rates under suitable conditions. This is the precise meaning of "two slow errors give one fast estimate".

3.4 AIPW is double ML for the ATE #

For a binary $T$ and heterogeneous effects, the target $\tau_{\mathrm{ATE}} = \mathbb{E}[\mu_1(X) - \mu_0(X)]$ has the orthogonal score $\psi = \varphi(W; \mu_0, \mu_1, e) - \tau$, where $\varphi$ is the AIPW pseudo-outcome of section 3.2. The DML recipe gives

with $\hat\mu_t$ and $\hat e$ cross-fitted. This is the AIPW estimator of lesson 05, now with ML models and cross-fitting. The product form of its bias (section 3.2) gives the rate condition $\|\hat e - e\|_2 \cdot \|\hat\mu_t - \mu_t\|_2 = o_P(n^{-1/2})$ for $t = 0, 1$, plus $\hat e$ bounded away from 0 and 1. Under it, $\hat\tau_{\mathrm{AIPW}}$ is asymptotically normal and reaches the smallest possible asymptotic variance (the semiparametric efficiency bound, Hahn 1998). In practice, clip $\hat e$ to an interval such as $[0.01, 0.99]$ and report how many units were affected.

3.5 Causal trees and causal forests #

A regression tree splits the covariate space into boxes (leaves) and predicts the mean of $Y$ in each leaf. A causal tree (Athey and Imbens 2016) predicts the effect in each leaf: in a randomized experiment, the difference between the mean outcome of treated and control units in the leaf. Two changes make this work.

1. Split to find different effects. A regression tree chooses the split that most reduces the squared error of $Y$. We cannot compute the squared error of $\tau$. Athey and Imbens show that, up to terms that do not depend on the split and a variance penalty, minimizing the expected squared error of $\hat\tau$ is the same as maximizing $\sum_{\text{leaves}} n_{\text{leaf}}\,\hat\tau_{\text{leaf}}^2$. So a good split makes the effects of its two children as different as possible.

2. Honesty. A tree is honest if one subsample chooses the splits and a separate subsample estimates the leaf effects. Why: the split search looks at many candidate splits and picks the one where the two sides look most different. Part of that difference is noise. If the same units then estimate the leaf effects, the noise stays in the estimates: the effects look more extreme than they are, and the usual confidence interval is too narrow. With honesty, the estimation units are fresh. Given the tree, the leaf difference in means is an unbiased estimate of the average effect in that leaf, and the usual interval $\hat\tau_{\text{leaf}} \pm 1.96\sqrt{s_1^2/n_1 + s_0^2/n_0}$ is valid.

Interactive figure: loading…

Causal forests (Wager and Athey 2018) average many honest causal trees, each grown on a random subsample (drawn without replacement) and with random subsets of features at each split: $\hat\tau(x) = \frac1B\sum_{b=1}^B \hat\tau_b(x)$. Under conditions on the subsample size, the leaf sizes and the smoothness of $\tau$, they prove that $\hat\tau(x)$ is consistent and asymptotically normal at each $x$, with a variance that can be estimated (by the infinitesimal jackknife). So a causal forest gives a confidence interval for $\tau(x)$, which most ML methods cannot.

Generalized random forests and the R-learner

Forest weights. A forest defines how similar a training point $X_i$ is to a target point $x$: $\alpha_i(x) = \frac1B \sum_b \frac{\mathbf 1\{X_i \in L_b(x)\}}{|L_b(x)|}$, where $L_b(x)$ is the leaf of tree $b$ that contains $x$. The weights are non-negative and add up to 1. A random forest prediction is $\sum_i \alpha_i(x) Y_i$: an adaptive nearest-neighbor average.

Generalized random forests (Athey, Tibshirani and Wager 2019) use these weights to solve a local moment equation $\sum_i \alpha_i(x)\,\psi_{\theta(x)}(O_i) = 0$ for any estimating equation $\psi$. For the CATE, the grf causal forest first computes out-of-bag nuisance estimates $\hat\ell(X_i)$ and $\hat e(X_i)$ ("local centering"), then solves the residual-on-residual equation of section 3.3 locally:

$$\hat\tau(x) = \frac{\sum_i \alpha_i(x)\,\tilde T_i\,\tilde Y_i}{\sum_i \alpha_i(x)\,\tilde T_i^2}, \qquad \tilde T_i = T_i - \hat e(X_i),\; \tilde Y_i = Y_i - \hat\ell(X_i).$$

So a causal forest is "double ML, solved separately near each $x$, with neighbors chosen by a forest". The same machinery gives quantile forests, instrumental-variable forests and others. Exercise 4 of the notebook builds a small version with scikit-learn.

R-learner (Nie and Wager 2021). Robinson's transformation also holds with a varying effect: if $Y = \mu_0(X) + \tau(X)\,T + \varepsilon$, then $Y - \ell(X) = \tau(X)\,(T - e(X)) + \varepsilon$. So the true $\tau$ minimizes the R-loss

$$\hat\tau = \arg\min_\tau\; \frac1n\sum_{i=1}^n \Big[\big(Y_i - \hat\ell(X_i)\big) - \tau(X_i)\,\big(T_i - \hat e(X_i)\big)\Big]^2 + \Lambda(\tau),$$

with cross-fitted $\hat\ell, \hat e$ and any regularizer $\Lambda$. Equivalently: regress the pseudo-outcome $(Y_i - \hat\ell)/(T_i - \hat e)$ on $X_i$ with sample weights $(T_i - \hat e)^2$, with any ML method that accepts weights. Nie and Wager prove a quasi-oracle property: with nuisance errors of order $o(n^{-1/4})$, $\hat\tau$ has the same error rate as if $\ell$ and $e$ were known. The R-loss is also a way to compare CATE models on validation data.

3.6 Evaluating a CATE model without ground truth #

For a prediction model we compute the test error. For a CATE model we cannot: $\tau_i = Y_i(1) - Y_i(0)$ is never observed (the fundamental problem of lesson 02). Four tools, all best used on a randomized held-out set (the propensity $e$ is known and there is no confounding):

1. The transformed outcome. With known $e(x)$, define

Theorem. Under assumptions 1–3, $\mathbb{E}[Y^* \mid X = x] = \tau(x)$.

Proof. $\mathbb{E}[TY/e(X) \mid X = x] = e(x)\,\mu_1(x)/e(x) = \mu_1(x)$ and $\mathbb{E}[(1 - T)Y/(1 - e(X)) \mid X = x] = \mu_0(x)$. Subtract and use section 3.1. $\blacksquare$

So $Y^*$ is a noisy but unbiased "label" for the effect. For a model $\hat\tau$ fitted on other data, $\mathbb{E}[(Y^* - \hat\tau(X))^2] = \mathbb{E}[(\tau(X) - \hat\tau(X))^2] + \mathbb{E}[(Y^* - \tau(X))^2]$ (the cross term vanishes because $\mathbb{E}[Y^* - \tau(X) \mid X] = 0$). The second term does not depend on the model. So the mean squared error against $Y^*$ ranks models correctly, although its value is much larger than the true error. (The AIPW pseudo-outcome $\hat\varphi$ is a lower-variance label of the same kind.)

2. Uplift and Qini curves. Sort the test units by $\hat\tau$, from high to low. For the top $k$ units, compare the treated and control outcomes among them:

where $\bar Y_T(k)$ and $\bar Y_C(k)$ are the mean outcomes of the treated and control units among the top $k$, and $n_T(k)$ is the number of treated among them. $U(k)$ estimates the number of extra successes if we treat these $k$ units (the uplift curve); $Q(k)$ is the Qini curve of Radcliffe (2007), the same quantity counted in the treated group. A good model rises fast; random targeting is the straight line from 0 to $U(n)$; the area between the curve and that line (the Qini coefficient) summarizes the model. Unlike a churn or response model, the uplift curve measures the effect of targeting.

Interactive figure: loading…

3. Calibration by groups. Sort the test units into, say, 5 groups by $\hat\tau$. In each group, estimate the average effect (difference in means in an RCT, or the mean of $\hat\varphi$) and compare it with the mean of $\hat\tau$ in the group. A useful model has group effects that increase from group 1 to group 5 and lie near the predictions. These group average treatment effects (GATES) come with confidence intervals, so the comparison is honest.

4. The best linear predictor test (Chernozhukov, Demirer, Duflo and Fernández-Val). Treat the model's score $S(X) = \hat\tau(X)$ as a feature. The best linear predictor of $\tau(X)$ from $S$ is $\beta_1 + \beta_2\,(S - \mathbb{E}S)$, with $\beta_1 = \tau_{\mathrm{ATE}}$ and $\beta_2 = \mathrm{Cov}(\tau, S)/\mathrm{Var}(S)$. Both can be estimated in an RCT by a linear regression of $Y$ on a constant, $T - e$ and $(T - e)(S - \bar S)$ (plus optional baseline controls), with weights $1/(e(1 - e))$. $\beta_2 = 0$ means the model found no heterogeneity that is real; $\beta_2 = 1$ means its scores have the right scale. A test of $\beta_2 = 0$ is a test for heterogeneity.

From effects to decisions. If treating costs $c$ (in units of $Y$), the best rule given $x$ is "treat if $\tau(x) > c$". The plug-in policy treats if $\hat\tau(x) > c$. Its value can be estimated on a randomized test set (lesson 11 does this for general policies). Athey and Wager (2021) learn a policy directly: choose a rule from a simple class (for example, a depth-2 tree) that maximizes the average AIPW score of the actions it takes. This gives rules that are easy to explain and come with regret guarantees.

3.7 A practical guide #

1. Write the assumptions first. Draw the graph (lessons 03–04). Put only pre-treatment variables in $X$. Causal ML adjusts only for what is in $X$; it cannot fix a missing confounder (use the sensitivity analysis of lesson 05).
2. Check overlap. Plot $\hat e(x)$ for treated and controls. With values near 0 or 1, AIPW and the DR-learner become unstable. Clip $\hat e$, restrict the population (and say so: the estimand changes), or use overlap weights.
3. Nuisance models. Use strong prediction methods (gradient boosting, random forests, lasso with interactions) and tune them by cross-validation for prediction. For a binary $T$, use a calibrated classifier for $\hat e$.
4. Cross-fit with $K = 5$ folds; repeat over a few random splits and report the median estimate.
5. Report in this order: the ATE (AIPW or DML) with a confidence interval; a test for heterogeneity (BLP, GATES); then the CATE model, its uplift curve on held-out randomized data, and the policy it implies.
6. Treat discovered subgroups as hypotheses. A forest searches many splits. Confirm a surprising subgroup on new data, ideally a new experiment.

Answers to section 1. (1) No. The likely buyers may buy without a coupon, so the coupon adds little for them, and for some it can even lower purchases. Targeting should rank by the effect $\tau(x)$, estimated with an uplift model; the response model curve in the uplift figure shows the loss. (2) $\hat\tau$ is shrunk toward 0: this is the regularization bias of the S-learner (section 2.1). (3) Much smaller. In DML the nuisance errors enter only as a product (section 3.3), so two errors of about $10\%$ give a bias of the order of $10\% \times 10\% = 1\%$. A plug-in estimator would carry the $10\%$ error directly.

4. Worked example #

A shop sends a coupon to some customers and measures spending $Y$ in the next month. The only covariate is $X$: $X = 0$ for new customers, $X = 1$ for returning customers. Returning customers were more likely to get the coupon. We assume that $X$ is the only confounder. The data, as cell sizes and mean spending:

Cell	$n$	mean $Y$
$X = 0$, coupon ($T = 1$)	20	12
$X = 0$, no coupon ($T = 0$)	80	10
$X = 1$, coupon	80	26
$X = 1$, no coupon	20	20

Step 1: the naive comparison. Mean spending with a coupon: $(20 \cdot 12 + 80 \cdot 26)/100 = 23.2$. Without: $(80 \cdot 10 + 20 \cdot 20)/100 = 12$. Naive difference: $11.2$. It is confounded: the coupon group is mostly returning customers, who spend more anyway.

Step 2: the CATE (T-learner with cell means). $\hat\tau(0) = 12 - 10 = 2$ and $\hat\tau(1) = 26 - 20 = 6$. The coupon helps returning customers three times as much. With a discrete $X$ and one mean per cell, the S- (with an interaction), T-, X- and DR-learners all give these same two numbers; they differ only when models must smooth or extrapolate.

Step 3: ATE and ATT. Both groups have 100 customers, so $\hat\tau_{\mathrm{ATE}} = 0.5 \cdot 2 + 0.5 \cdot 6 = 4$. The treated are 20 new and 80 returning customers, so $\hat\tau_{\mathrm{ATT}} = (20 \cdot 2 + 80 \cdot 6)/100 = 5.2$.

Step 4: double ML by hand. The nuisance functions are cell averages.

• $\hat e(0) = 20/100 = 0.2$ and $\hat e(1) = 80/100 = 0.8$.
• $\hat\ell(0) = 0.2 \cdot 12 + 0.8 \cdot 10 = 10.4$ and $\hat\ell(1) = 0.8 \cdot 26 + 0.2 \cdot 20 = 24.8$.

Residuals $\tilde T = T - \hat e(X)$ and $\tilde Y = Y - \hat\ell(X)$. Inside a cell, $\tilde T$ is the same for every unit, so $\sum \tilde T\,\tilde Y = n\,\tilde T\,(\bar Y_{\text{cell}} - \hat\ell)$ (the individual noise around the cell mean sums to 0):

Cell	$n$	$\tilde T$	$\bar Y_{\text{cell}} - \hat\ell$	$n\,\tilde T\,(\bar Y_{\text{cell}} - \hat\ell)$	$n\,\tilde T^2$
$X = 0$, $T = 1$	20	$0.8$	$12 - 10.4 = 1.6$	$25.6$	$12.8$
$X = 0$, $T = 0$	80	$-0.2$	$10 - 10.4 = -0.4$	$6.4$	$3.2$
$X = 1$, $T = 1$	80	$0.2$	$26 - 24.8 = 1.2$	$19.2$	$3.2$
$X = 1$, $T = 0$	20	$-0.8$	$20 - 24.8 = -4.8$	$76.8$	$12.8$
sum				$128$	$32$

So $\hat\theta = 128/32 = 4$. Check with the formula of the warning box in section 3.3: the weights are $n_x\,\hat e(x)(1 - \hat e(x)) = 100 \cdot 0.16 = 16$ for both strata, so $\hat\theta = (16 \cdot 2 + 16 \cdot 6)/32 = 4$. Here it equals the ATE because the two strata have the same weight. Check-yourself question 2 changes one number and shows that it does not always.

Step 5: one AIPW pseudo-outcome. A returning customer without a coupon spent $Y = 17$. With $\hat\mu_1(1) = 26$, $\hat\mu_0(1) = 20$, $\hat e(1) = 0.8$:

One pseudo-outcome is far from the cell effect 6: the weight $1/(1 - 0.8) = 5$ multiplies the residual by 5. But the residuals of the cell average to 0, so the mean of $\hat\varphi$ over the cell is exactly 6. This is why DR-learners need many units, or good overlap, to be stable.

Step 6: an uplift curve with 8 units. A randomized test, units already sorted by $\hat\tau$ from high to low; $T$ = coupon, $Y$ = bought (1) or not (0):

rank	1	2	3	4	5	6	7	8
$T$	1	0	1	0	1	0	1	0
$Y$	1	0	1	0	0	0	0	1

• Top 2: treated $\bar Y_T = 1$, control $\bar Y_C = 0$; $U(2) = (1 - 0) \cdot 2 = 2$.
• Top 4: $\bar Y_T = 2/2 = 1$, $\bar Y_C = 0/2 = 0$; $U(4) = 4$.
• Top 6: $\bar Y_T = 2/3$, $\bar Y_C = 0/3 = 0$; $U(6) = (2/3) \cdot 6 = 4$.
• All 8: $\bar Y_T = 2/4$, $\bar Y_C = 1/4$; $U(8) = (1/4) \cdot 8 = 2$.

The curve is $0, 2, 4, 4, 2$: treating the top half gives about 4 extra purchases, treating more adds nothing and then loses. If a coupon costs $0.3$ purchases, the net gain is $4 - 4 \cdot 0.3 = 2.8$ for the top half and $4 - 6 \cdot 0.3 = 2.2$ for the top 75%. So we treat the top half. (With 8 units these numbers are very noisy; real uplift curves need thousands of units, as in the figure.)

5. In code #

NumPy and scikit-learn check the worked example and run DML with random forests on simulated data.

import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import KFold

# (a) the two-stratum example: cell sizes and cell means of Y
n  = {(0, 1): 20, (0, 0): 80, (1, 1): 80, (1, 0): 20}      # (X, T): count
ym = {(0, 1): 12., (0, 0): 10., (1, 1): 26., (1, 0): 20.}  # (X, T): mean Y
x = np.concatenate([[k[0]] * v for k, v in n.items()])
t = np.concatenate([[k[1]] * v for k, v in n.items()]).astype(float)
dev = np.concatenate([np.resize([-1., 1.], v) for v in n.values()])   # noise that averages to 0 in each cell
y = np.array([ym[(a, int(b))] for a, b in zip(x, t)]) + dev

tau_hat = {s: float(y[(x == s) & (t == 1)].mean() - y[(x == s) & (t == 0)].mean()) for s in (0, 1)}
e_hat = {s: float(t[x == s].mean()) for s in (0, 1)}
print(tau_hat, e_hat)                                  # {0: 2.0, 1: 6.0} {0: 0.2, 1: 0.8}
print(round(y[t == 1].mean() - y[t == 0].mean(), 3))   # 11.2  (naive)
print(np.mean([tau_hat[s] for s in x]))                # 4.0   (ATE)
print(np.mean([tau_hat[s] for s in x[t == 1]]))        # 5.2   (ATT)

# residual on residual, with the cell averages as nuisance functions
t_res = t - np.array([e_hat[s] for s in x])
l_hat = {s: y[x == s].mean() for s in (0, 1)}
y_res = y - np.array([l_hat[s] for s in x])
print(round(t_res @ y_res, 3), round(t_res @ t_res, 3), round((t_res @ y_res) / (t_res @ t_res), 3))
# 128.0 32.0 4.0

# (b) double ML with random forests and 2-fold cross-fitting on simulated data
rng = np.random.default_rng(0)
N = 2000
X = rng.normal(size=(N, 5))
m = np.tanh(2 * X[:, 0]) + 0.5 * X[:, 1]                 # E[T | X]
g = 2 * np.tanh(2 * X[:, 0]) + np.sin(2 * X[:, 1])       # effect of X on Y
T = m + rng.normal(size=N)
Y = 0.5 * T + g + rng.normal(size=N)                      # true theta = 0.5

print(round(np.polyfit(T, Y, 1)[0], 3))                    # 1.245  (naive slope, confounded)
l_hat_, m_hat_ = np.zeros(N), np.zeros(N)
for train, test in KFold(2, shuffle=True, random_state=0).split(X):
    rf = dict(n_estimators=200, min_samples_leaf=5, random_state=0)
    l_hat_[test] = RandomForestRegressor(**rf).fit(X[train], Y[train]).predict(X[test])
    m_hat_[test] = RandomForestRegressor(**rf).fit(X[train], T[train]).predict(X[test])
Tt, Yt = T - m_hat_, Y - l_hat_
theta = (Tt @ Yt) / (Tt @ Tt)
se = np.sqrt(np.sum(Tt**2 * (Yt - theta * Tt) ** 2)) / (Tt @ Tt)
print(round(theta, 3), round(se, 3))                       # 0.514 0.022

# (c) one AIPW pseudo-outcome: a control unit in stratum X = 1 with Y = 17
mu1, mu0, e = 26., 20., 0.8
print(round(mu1 - mu0 - (17 - mu0) / (1 - e), 3))         # 21.0

# (d) uplift curve of 8 randomized units, already sorted by predicted effect
tt = np.array([1, 0, 1, 0, 1, 0, 1, 0]); yy = np.array([1, 0, 1, 0, 0, 0, 0, 1])
for k in (2, 4, 6, 8):
    a, b = tt[:k] == 1, tt[:k] == 0
    print(k, float((yy[:k][a].mean() - yy[:k][b].mean()) * k))   # 2 2.0 / 4 4.0 / 6 4.0 / 8 2.0

The DML estimate $0.514$ has a 95% interval of about $[0.47, 0.56]$, which contains the true $0.5$; the naive slope $1.245$ is far off. The notebook adds the plug-in estimator, repeats the experiment over many seeds, and builds the meta-learners, a small causal forest and an uplift curve.

6. Where it shows up #

• Uplift modeling and targeting. Coupons, retention offers, notifications and ads are sent to the customers with the largest estimated effect, not the largest purchase or churn probability. Models are trained on randomized campaigns and evaluated with Qini curves.
• Heterogeneous effects in A/B tests. Large online experiments report the ATE and then look for segments (new vs returning users, device, country) where the effect differs, with causal forests or the BLP/GATES tests, and with care about multiple testing.
• Medicine. Causal forests and DR-learners on trial or registry data estimate which patients benefit most from a treatment, and which ones carry risk without benefit.
• Economics and policy. DML is standard for effects with many controls: the original paper estimates the effect of 401(k) eligibility on savings; later work uses it for price elasticities and labour-market programs. Causal forests target job-training and welfare programs.
• Machine learning for science. In agriculture and ecology, the effect of an intervention (fertilizer, irrigation, a conservation policy) depends on soil, weather and location; causal forests map where it helps. In climate and Earth science, DML-style partialling out is used to separate the effect of one driver from many correlated background variables.

7. Common mistakes #

• Targeting with a response or churn model. It ranks by the outcome, not by the effect. The people most likely to buy or to stay are often the ones the treatment changes least. Use an uplift model and check it with an uplift curve on randomized data.
• Reading the effect from an S-learner with strong regularization. The model can ignore $T$, and $\hat\tau$ shrinks toward 0. Compare with a T-, X- or DR-learner, and check that $T$ is actually used.
• No cross-fitting. Residuals from nuisance models that saw the same units are too small and correlated with the noise: biased $\hat\theta$ and too narrow intervals. Always predict each unit with a model trained on other folds.
• "ML handles confounding." DML and causal forests adjust only for the measured $X$. An unmeasured confounder biases them as much as it biases regression. Use domain knowledge and a sensitivity analysis (lesson 05).
• Many features, including bad controls. Putting post-treatment variables (mediators, colliders) into $X$ "because the model can handle it" breaks identification (lessons 03–04). Only pre-treatment covariates belong in $X$.
• Reporting the partially linear $\theta$ as the ATE when effects vary. It is an overlap-weighted average (section 3.3). For the ATE use AIPW.
• Ignoring overlap. With $\hat e$ near 0 or 1, pseudo-outcomes explode. Check the histogram of $\hat e$, and clip or restrict, saying how the estimand changes.
• Evaluating a CATE model in-sample, or with $Y$ as the label. Use a held-out randomized set, the transformed outcome, uplift curves and group effects.

8. Check yourself #

1. In an X-learner, at some $x$: $\hat e(x) = 0.1$, $\hat\tau_0(x) = 3$, $\hat\tau_1(x) = 1$. Compute $\hat\tau(x)$, and say why the X-learner trusts $\hat\tau_1$ more here.
   Answer

   $\hat\tau(x) = 0.1 \cdot 3 + 0.9 \cdot 1 = 1.2$. Only 10% of units like this are treated. $\hat\tau_1$ uses the imputed effects $Y_i - \hat\mu_0(X_i)$ of the treated, where $\hat\mu_0$ is fitted on the large control group and is accurate. $\hat\tau_0$ uses $\hat\mu_1$, which is fitted on few points and is less accurate.
2. Change the worked example: among returning customers, 50 get the coupon (mean spending 26) and 50 do not (mean 20). Nothing else changes. What do the ATE and the partially linear DML estimate become?
   Answer

   The CATEs stay $2$ and $6$, so the ATE stays $4$. Now $\hat e(1) = 0.5$, so the weights are $100 \cdot 0.16 = 16$ for new and $100 \cdot 0.25 = 25$ for returning customers: $\hat\theta = (16 \cdot 2 + 25 \cdot 6)/41 = 182/41 \approx 4.44$. The partially linear estimate gives more weight to the stratum where the treatment varies more. It is not the ATE when effects differ.
3. In an RCT with $e = 0.5$, compute the transformed outcome $Y^*$ for a treated unit with $Y = 1$, a treated unit with $Y = 0$ and a control unit with $Y = 1$. Why can single values of $Y^*$ be far from any plausible effect?
   Answer

   $Y^* = Y(T - e)/(e(1 - e)) = Y(T - 0.5)/0.25$: treated with $Y = 1$ gives $2$; treated with $Y = 0$ gives $0$; control with $Y = 1$ gives $-2$. Each $Y^*$ uses one outcome only, so it is very noisy. Only its conditional mean equals $\tau(x)$. That is why it ranks models correctly on average but its MSE level is large.
4. Does cross-fitting with $K = 5$ folds mean that the final estimate uses only one fifth of the data?
   Answer

   No. Every unit gets residuals from models trained on the other four folds, and $\hat\theta$ pools all $n$ units. Every unit is also used to train the nuisance models for the other folds. The only cost is that each nuisance model is trained on $80\%$ of the data.
5. At depth 2, the honest tree in the figure has a larger MSE than the adaptive tree; at depth 6, a smaller one. Explain both.
   Answer

   Honest trees use half the units to choose splits and half to estimate effects, so each part is noisier: at depth 2 this loss dominates. Deep adaptive trees choose many splits on noise and then estimate the leaves on the same noise, so their leaf effects are too extreme; this overfitting dominates at depth 6. Coverage is the clearer difference: about 95% for honest leaves, about 75–80% for adaptive leaves.
6. A DML analysis with random forests gives $\hat\theta = 0.51$ (SE $0.02$); OLS of $Y$ on $T$ and $X$ gives $0.90$. Which assumptions must hold before you call $0.51$ the causal effect?
   Answer

   No unmeasured confounding given $X$ (conditional exchangeability), only pre-treatment covariates in $X$, consistency, and enough variation of $T$ given $X$ (overlap, $\mathbb{E}[\mathrm{Var}(T \mid X)] > 0$). The partially linear form must hold, or $0.51$ must be read as the weighted average $\theta^*$. The nuisance models must be good enough for the rate condition. The difference from OLS suggests nonlinear confounding that OLS missed.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Define the CATE and say when it equals $\mu_1(x) - \mu_0(x)$.

   Show answer

2. Why does a regularized S-learner tend to underestimate effects?

   Show answer

3. How does the X-learner combine its two effect models?

   Show answer

4. Write the DML estimator for the partially linear model.

   Show answer

5. What is Neyman orthogonality, in words?

   Show answer

6. What does cross-fitting prevent?

   Show answer

7. State the rate condition for DML in the partially linear model.

   Show answer

8. What makes a causal tree "honest", and what does it buy?

   Show answer

9. In an RCT with known $e(x)$, what is the transformed outcome and its key property?

   Show answer

10. What does an uplift curve show?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How would you decide which users get a discount?" — Not by purchase or churn probability. Run a randomized campaign (or use logged data with all confounders measured), estimate the CATE with an uplift model (X-/DR-learner, causal forest), evaluate it with an uplift/Qini curve on a randomized holdout, and treat users with $\hat\tau(x)$ above the cost of the discount.
• "What is double ML, and why cross-fit?" — Partial out $X$ from $T$ and $Y$ with ML models and regress residual on residual. The score is Neyman orthogonal, so nuisance errors enter as a product and the estimate is $\sqrt n$-consistent with a valid CI. Cross-fitting stops the nuisance models from fitting the noise of the units they are evaluated on.
• "S-, T- or X-learner?" — S if the effect is small and data scarce (but it shrinks effects); T with plenty of data in both arms; X when the arms are unbalanced; DR or R when overlap is good and you want robustness and fast rates.
• "How do you evaluate a CATE model if you never see individual effects?" — On a randomized holdout: uplift/Qini curves, calibration by groups (GATES), the BLP test, MSE against the transformed outcome for ranking models.
• "Why are causal forests honest?" — To avoid choosing splits and estimating effects on the same noise; it makes leaf estimates unbiased given the tree and gives valid confidence intervals.

10. Practice #

The notebook for this lesson has five exercises, each with a check cell and a folded solution:

1. Simulate data with a known $\tau(x)$ and build the S-, T- and X-learners with scikit-learn; compare their error against the truth with 10% and 50% treated.
2. Write dml_plr with $K$-fold cross-fitting and a standard error; compare the plug-in estimator, DML without cross-fitting and DML with cross-fitting over many seeds (bias and coverage).
3. Build the DR-learner from cross-fitted AIPW pseudo-outcomes and check its ATE against the truth.
4. Build a small honest causal forest: a random forest on one half gives the neighborhoods, and a weighted residual-on-residual regression on the other half gives $\hat\tau(x)$.
5. Compute an uplift curve on a simulated randomized test set and choose the best budget for a given cost per treatment.

11. Watch and read #

6.1 - CATE Preliminaries and OutlineBrady Neal · Introduction to Causal Inference, lecture 6 · 4 min

The CATE and the plan of the lecture: conditional outcome models, X-learner, double ML, causal trees.

6.2 - Conditional Outcome ModelingBrady Neal · Introduction to Causal Inference, lecture 6 · 10 min

The S-learner and the T-learner, and why the S-learner can ignore the treatment.

6.3 - TARNet and X-LearnerBrady Neal · Introduction to Causal Inference, lecture 6 · 7 min

The X-learner step by step and its weighting by the propensity score.

6.5 - Doubly Robust Methods, Matching, Double Machine Learning, and Causal TreesBrady Neal · Introduction to Causal Inference, lecture 6 · 8 min

A short overview of doubly robust estimation, double ML and causal trees.

Double Machine Learning for Causal and Treatment EffectsBecker Friedman Institute University of Chicago · Victor Chernozhukov · 40 min

The author of DML explains regularization bias, Neyman orthogonality and cross-fitting.

Susan Athey Guest Talk - Estimating Heterogeneous Treatment EffectsBrady Neal - Causal Inference · 57 min

Honest causal trees, causal forests and policy learning, from one of their authors.

Conditional Average Treatment Effects: ForestsStanford Graduate School of Business · 37 min

Causal forests and generalized random forests in more depth.

Dr. Juan Orduz: Introduction to Uplift ModelingPyData · 44 min

Meta-learners for uplift and how to read uplift and cumulative gain curves in practice.

• Chernozhukov, Hansen, Kallus, Spindler and Syrgkanis, Applied Causal Inference Powered by ML and AI (free): chapters 4 and 10 (orthogonality and double ML), 14–15 (heterogeneous effects, meta-learners, model evaluation).
• Stefan Wager, Causal Inference: A Statistical Learning Approach (free PDF): chapter 3 (doubly robust methods, debiased ML), chapter 4 (heterogeneous effects, R-learner), chapter 5 (policy learning).
• Chernozhukov et al. (2018), Double/debiased machine learning for treatment and structural parameters.
• Künzel, Sekhon, Bickel and Yu (2019), Meta-learners for estimating heterogeneous treatment effects using machine learning.
• Kennedy (2023), Towards optimal doubly robust estimation of heterogeneous causal effects.
• Nie and Wager (2021), Quasi-oracle estimation of heterogeneous treatment effects (the R-learner).
• Athey and Imbens (2016), Recursive partitioning for heterogeneous causal effects (honest causal trees).
• Wager and Athey (2018), Estimation and inference of heterogeneous treatment effects using random forests; Athey, Tibshirani and Wager (2019), Generalized random forests.
• Chernozhukov, Demirer, Duflo and Fernández-Val, Generic machine learning inference on heterogeneous treatment effects in randomized experiments (BLP and GATES).
• Matheus Facure, Causal Inference for the Brave and True, part II: meta-learners, debiased ML and evaluation of CATE models, with Python code.