Meta-Learners for Conditional Treatment Effects

Why Use Meta-Learners?

A conditional treatment effect describes how the average effect of a treatment varies with observed characteristics. Estimating it requires more than predicting outcomes well: the analysis must compare the treated and untreated outcome patterns at the same covariate values. Meta-learners provide different ways to organize that comparison using outcome and treatment-assignment models.

The four learners in this tutorial make distinct choices about how to estimate those models and combine them. The S-learner fits one outcome model, the T-learner fits separate outcome models by treatment group, the X-learner imputes effects and combines estimates using the propensity score, and the R-learner fits effects through a residualized loss. None is uniformly best. Their behavior depends on overlap, treatment-group sizes, and whether the fitted models capture outcome and effect patterns well.

Definition: Let $W\in\{0,1\}$ indicate treatment, $Y$ be the observed outcome, and $X$ be observed covariates. Let $Y(1)$ and $Y(0)$ be the potential outcomes under treatment and control. The conditional average treatment effect (CATE) is $$ \tau(x)=E[Y(1)-Y(0)\mid X=x]. $$ Under consistency, conditional exchangeability $(Y(1),Y(0))\perp W\mid X$, and positivity $0<P(W=1\mid X=x)<1$ for covariate values of interest, it can be identified as $E[Y\mid W=1,X=x]-E[Y\mid W=0,X=x]$.

Here consistency means the observed outcome is $Y=Y(W)$. Conditional exchangeability rules out unmeasured confounding after conditioning on $X$, and positivity requires both treatment conditions to be possible at the covariate values being studied. These assumptions are substantive; choosing a learner does not make them true. A CATE is an average within a covariate-defined group, not an observed or generally identifiable effect for each individual.

S- and T-Learners: One Outcome Model or Two?

Write $m_w(x)=E[Y\mid W=w,X=x]$ for the conditional outcome mean under treatment status $w$, and $e(x)=P(W=1\mid X=x)$ for the propensity score. The S- and T-learners estimate the two outcome means in different ways, then take their difference. Their target contrast is $\tau(x)=m_1(x)-m_0(x)$, under the causal assumptions above.

Definition: An S-learner fits one outcome model $\hat f(x,w)$ using both covariates and treatment status. Its estimated effect is $\hat\tau_S(x)=\hat f(x,1)-\hat f(x,0)$. A T-learner fits separate outcome models $\hat m_1(x)$ among treated units and $\hat m_0(x)$ among controls, then estimates $\hat\tau_T(x)=\hat m_1(x)-\hat m_0(x)$.

An S-learner pools observations to fit a single model. This can be efficient when the outcome patterns are similar across groups and the model allows treatment to interact with relevant covariates. But if the fitting procedure heavily shrinks the treatment contribution, or the model omits important treatment interactions, estimated effects can be too similar across $x$. A T-learner allows the outcome relationships to differ freely between groups, but each model is trained on only one group. It can therefore be unstable when a group is small or has sparse coverage of some covariate values.

Worked Example: Comparing S- and T-Learners

Suppose a hypothetical outcome is measured in points, and $x$ is a binary covariate. A fitted S-learner uses $\hat f(x,w)=10+2x+3w$. A fitted T-learner uses $\hat m_0(x)=10+2x$ and $\hat m_1(x)=13+7x$. Calculate both estimated effects at $x=0$ and $x=1$.

For the S-learner, the predicted outcomes at $x=0$ are $\hat f(0,1)=10+0+3=13$ and $\hat f(0,0)=10+0+0=10$, so $\hat\tau_S(0)=13-10=3$ points. At $x=1$, the predictions are $15$ and $12$, giving $\hat\tau_S(1)=15-12=3$ points. This fitted model imposes a constant effect because it contains no interaction between $x$ and $w$.

For the T-learner, at $x=0$ the estimates are $\hat m_1(0)=13$ and $\hat m_0(0)=10$, so $\hat\tau_T(0)=3$ points. At $x=1$, they are $\hat m_1(1)=20$ and $\hat m_0(1)=12$, so $\hat\tau_T(1)=8$ points. The difference from the S-learner is a modeling choice, not a calculation error: the separate models allow a varying effect, while the stated S-model cannot represent one. These fitted equations are stipulated for illustration, not claimed to have been estimated from an unreported data set.

X-Learner: Impute Effects, Then Combine

The X-learner first fits the two outcome models. It uses them to impute missing counterfactual outcomes for observed units: treated units have observed $Y(1)$ but not $Y(0)$, while controls have observed $Y(0)$ but not $Y(1)$. It then fits models to these imputed effect values separately and combines the resulting estimates. Thus the outcome models are especially important because their predictions enter directly into the imputed effects.

Definition: The X-learner forms $D^{(1)}=Y-\hat m_0(X)$ among treated units and $D^{(0)}=\hat m_1(X)-Y$ among controls. It fits effect models $\hat\tau_1(x)$ to $D^{(1)}$ and $\hat\tau_0(x)$ to $D^{(0)}$. A common propensity-weighted combination is $\hat\tau_X(x)=\hat e(x)\hat\tau_0(x)+[1-\hat e(x)]\hat\tau_1(x)$.

The superscripts identify the group used to fit each imputed-effect model: $\hat\tau_1$ is fitted using treated units, and $\hat\tau_0$ using controls. The combination assigns weight $\hat e(x)$ to the control-fitted effect and weight $1-\hat e(x)$ to the treated-fitted effect. When treated units are rare at $x$, this gives greater weight to the treated-fitted effect model, whose imputed outcomes use the outcome model estimated from the more numerous controls. This can be useful under imbalance, but it does not guarantee accuracy: errors in either outcome model can contaminate the imputed effects, and a near-zero or near-one propensity still signals weak overlap.

Worked Example: X-Learner Under Treatment Imbalance

Consider a hypothetical covariate group with estimated propensity $\hat e(x)=0.20$, so treatment is relatively uncommon. Suppose the fitted outcome models give $\hat m_0(x)=10$ and $\hat m_1(x)=16$. In this group, two treated units have outcomes 17 and 18, while two control units have outcomes 11 and 12. For simplicity, suppose the fitted effect models predict the respective averages of their group’s imputed effects at $x$.

For the treated units, the imputed effects are $17-10=7$ and $18-10=8$. Therefore, $\hat\tau_1(x)=(7+8)/2=7.5$ outcome units. For the controls, they are $16-11=5$ and $16-12=4$, so $\hat\tau_0(x)=(5+4)/2=4.5$ outcome units.

Combining them gives $\hat\tau_X(x)=0.20(4.5)+0.80(7.5)=0.90+6.00=6.90$ outcome units. This calculation shows how the propensity weighting works; it does not establish that either imputed-effect model is accurate. In particular, the estimate depends on the correctness of the outcome predictions used to fill in the missing potential outcomes.

R-Learner: Compare Residualized Outcomes

The R-learner uses the outcome regression $m(x)=E[Y\mid X=x]$ and propensity score $e(x)$ to remove predictable components of outcome and treatment. It then chooses an effect function to make the remaining outcome resemble the remaining treatment multiplied by that effect. This is a way to estimate heterogeneity without fitting separate conditional outcome models as the final step.

Definition: Given estimates $\hat m(x)$ and $\hat e(x)$, an R-learner estimates $\tau(x)$ by minimizing a regularized version of $$ \frac{1}{n}\sum_{i=1}^n\left\{Y_i-\hat m(X_i)-[W_i-\hat e(X_i)]\tau(X_i)\right\}^2. $$ Regularization or restrictions on the function class may be needed to avoid an overly variable fit.

The population target of this loss can be justified directly. Under the causal assumptions, write $m(x)=E[Y\mid X=x]$. Conditional on $X=x$, the treatment residual has mean zero and variance $e(x)[1-e(x)]$. Also, $E[(W-e(x))(Y-m(x))\mid X=x]=e(x)[1-e(x)]\tau(x)$. Therefore, the conditional expected squared loss, as a function of a candidate value $t$, has derivative $$ -2E[(W-e(x))(Y-m(x))\mid X=x]+2tE[(W-e(x))^2\mid X=x]. $$ Setting this derivative to zero gives $t=\tau(x)$, provided $0<e(x)<1$. This identifies the population target of the residualized loss when the nuisance functions are the true ones. In practice they are estimated, and limited overlap can make the fit unstable because little residual treatment variation is available at some $x$.

Worked Example: A Constant-Effect R-Learner Fit

Suppose four hypothetical observations have residualized treatments $W_i-\hat e(X_i)$ equal to $0.5,-0.5,0.5,-0.5$, and residualized outcomes $Y_i-\hat m(X_i)$ equal to $4,-2,6,-4$, respectively. Fit a constant effect $\tau(X_i)=c$ by minimizing the unregularized R-loss.

For a constant $c$, the loss is $L(c)=\sum_{i=1}^4(r_{Y,i}-r_{W,i}c)^2$, where $r_{Y,i}=Y_i-\hat m(X_i)$ and $r_{W,i}=W_i-\hat e(X_i)$. Differentiating and setting the derivative to zero yields $$ \hat c=\frac{\sum_i r_{W,i}r_{Y,i}}{\sum_i r_{W,i}^2}. $$

The numerator is $0.5(4)+(-0.5)(-2)+0.5(6)+(-0.5)(-4)=2+1+3+2=8$. The denominator is $0.5^2+(-0.5)^2+0.5^2+(-0.5)^2=1$. Thus $\hat c=8/1=8$ outcome units per treatment indicator. This is an illustrative residualized fit, not evidence of a true effect; its interpretation relies on the identifying assumptions and adequate nuisance estimates.

Choosing Among Learners

Treatment-group imbalance and outcome-model quality influence these methods in different ways. The table summarizes the main trade-offs; no row replaces checking overlap, model fit, and the scientific plausibility of the identifying assumptions.

Learner	How it constructs the effect	Potential advantage	Important vulnerability
S-learner	Difference between predictions from one pooled outcome model	Shares information across groups and can be parsimonious	A weak treatment contribution or omitted interactions can hide heterogeneity
T-learner	Difference between separate treated and control outcome models	Allows outcome relationships to differ flexibly by group	The smaller group may yield a noisy model or poor covariate coverage
X-learner	Models imputed effects separately, then propensity-weights them	Can use the larger group’s outcome model to help impute effects when groups are imbalanced	Errors in outcome predictions flow into the imputed effects; weak overlap remains a problem
R-learner	Fits a function to residualized outcome and treatment variation	Targets treatment-effect variation after removing predictable components	Needs useful residual treatment variation; nuisance error and regularization choices matter

An outcome model can predict $Y$ accurately overall while estimating $\tau(x)$ poorly. For example, if treatment effects are small relative to the variation in baseline outcomes, a model that predicts the baseline well but smooths away treatment-by-covariate differences may have low prediction error and still miss the heterogeneity of interest. Conversely, a highly flexible effect estimate can vary substantially in regions with few treated or few control observations. The target population and covariate regions where both treatments are observed should therefore guide model assessment.

Common Pitfall

A common mistake is to select the learner with the best overall outcome prediction and assume it must give the best CATE estimates. That reasoning is incomplete: outcome prediction error is not the same as error in the difference $m_1(x)-m_0(x)$. The correction is to examine whether the chosen models can represent the relevant treatment-effect patterns, whether each treatment group provides adequate information across covariate values, and whether predictions used to construct effects are credible in the target region. A flexible learner cannot compensate for absent overlap or unmeasured confounding.

Key takeaway: S-, T-, X-, and R-learners differ in how they organize outcome modeling and effect estimation. Their estimates have a causal interpretation only under appropriate identification assumptions, and their relative performance depends on treatment-group balance, model quality, and support in the covariate regions of interest.

Check Your Understanding

Use the definitions and assumptions above to answer the following questions.

What conditional outcome models does a T-learner fit, and how does it construct its estimated effect?
In the X-learner combination shown here, what happens to the weight on $\hat\tau_1(x)$ as $\hat e(x)$ becomes small?
Why can an S-learner with no treatment-by-covariate interaction fail to represent a varying treatment effect?
What role does residual treatment variation play in the R-learner, and why is it limited when the propensity is near zero or one?
State the assumptions needed for the conditional difference in observed outcome means to identify the CATE.