Skip to content

X-learner

The X-learner estimates conditional treatment effects using cross-group outcome imputation, two effect regressions and a weighted blend. Understand its imbalance advantage and causal limits.

The X-learner estimates a conditional average treatment effect by imputing missing outcomes, fitting effect regressions separately in the treated and control groups, and blending their predictions. In an execution experiment with many orders using the incumbent rule and few using a new rule, it offers a way to use the large control sample to help estimate how the new rule's benefit varies across order characteristics. It is a recipe for combining regression models, rather than one particular model family.

Cross the predictions, then learn the effect

Let W=1W=1 denote treatment, W=0W=0 control, XX pre-treatment features and YY the observed outcome. First fit two outcome regressions: μ^0(x)\widehat\mu_0(x) using control observations, and μ^1(x)\widehat\mu_1(x) using treated observations. Construct the pseudo-outcomes

D~i1=Yi−μ^0(Xi)for Wi=1,D~i0=μ^1(Xi)−Yifor Wi=0.\widetilde D_i^1=Y_i-\widehat\mu_0(X_i)\quad\text{for }W_i=1, \qquad \widetilde D_i^0=\widehat\mu_1(X_i)-Y_i\quad\text{for }W_i=0.
Each observed outcome is paired with the other group's outcome prediction. Both differences use the treatment-minus-control sign. Next regress the treated pseudo-outcomes on their features to obtain τ^1(x)\widehat\tau_1(x), and the control pseudo-outcomes on their features to obtain τ^0(x)\widehat\tau_0(x). The subscripts identify the group supplying the second-stage observations, not two different causal targets.

The final prediction is

τ^X(x)=g(x)τ^0(x)+{1−g(x)}τ^1(x),0≤g(x)≤1.\widehat\tau_X(x)=g(x)\widehat\tau_0(x)+\{1-g(x)\}\widehat\tau_1(x),\qquad 0\le g(x)\le1.
This follows Künzel et al., X-Learner, equations 5–9 and Remark 1. Here gg weights the control-group effect regression. An estimated propensity score, e^(x)≈P(W=1∣X=x)\widehat e(x)\approx\mathbb P(W=1\mid X=x), is a suggested choice for gg, not a mandatory or universally optimal weight.

Two imputed differences are not a fitted effect surface

Consider fictional execution benefits in basis points per order, relative to a fixed benchmark and after consistently measured costs; higher is better. Treatment means switching to the new execution rule. The numbers below are stipulated observations and predictions, with no uncertainty estimates supplied.
ObservationObserved outcomePrediction under the other ruleImputed effect
Treated order7 bpControl: 3 bp7−3=47-3=4 bp
Control order2 bpTreatment: 6 bp6−2=46-2=4 bp

These are noisy training labels, not observed individual causal effects. A prediction of what the first order would have earned under control does not reveal its missing potential outcome.

Separately, suppose second-stage models trained on a larger dataset predict τ^0(x)=2\widehat\tau_0(x)=2 bp and τ^1(x)=4\widehat\tau_1(x)=4 bp at a chosen feature value. These predictions are stipulated; they are not fits from the two rows above. With g(x)=0.25g(x)=0.25,
τ^X(x)=0.25(2)+0.75(4)=3.5 bp.\widehat\tau_X(x)=0.25(2)+0.75(4)=3.5\ \text{bp}.

The result lies between the two predictions and is closer to 4 bp because that prediction receives three-quarters of the weight. Reversing which model gets gg would answer a different calculation.

Fictional execution-benefit arithmetic in basis points: treated outcome 7 minus predicted control 3 gives 4; predicted treatment 6 minus control outcome 2 gives 4. Separately stipulated effect predictions 2 and 4 blend to 3.5 at weight g equals 0.25.
Fictional execution-benefit arithmetic in basis points: treated outcome 7 minus predicted control 3 gives 4; predicted treatment 6 minus control outcome 2 gives 4. Separately stipulated effect predictions 2 and 4 blend to 3.5 at weight g equals 0.25.
Original fictional illustration. Top: two imputation examples. Bottom: separately stipulated effect-model predictions at one feature value, not fits from those two observations. The blend uses weight g on the control-group effect regression.

Why imbalance can help—and where it cannot

A large control sample can support a good control-outcome fit. Subtracting that fit from treated outcomes can leave an effect function simpler than either outcome surface. When treatment is rare locally and g=e^g=\widehat e, the blend emphasises τ^1\widehat\tau_1: the treated-group effect regression whose labels use the control-outcome model. The reverse logic applies with abundant treated observations. These are advantages under suitable structure and estimation quality, not an automatic benefit from imbalance. Künzel et al., Intuition Behind the Metalearners give the mechanism.
Imbalance also differs from absent causal overlap. For causal interpretation, the baseline setup requires well-defined actions and consistency, no unaccounted interference, conditional exchangeability of potential outcomes and assignment given XX, and both assignments possible at the features being studied. Under those conditions,
τ(x)=E[Y(1)−Y(0)∣X=x]=E[Y∣W=1,X=x]−E[Y∣W=0,X=x].\tau(x)=\mathbb E[Y(1)-Y(0)\mid X=x] =\mathbb E[Y\mid W=1,X=x]-\mathbb E[Y\mid W=0,X=x].
Nie and Wager, §2 state the potential-outcome and exchangeability setup; Kennedy, §2 lists consistency, positivity and no unmeasured confounding. If an unrecorded urgency signal drives both execution-rule choice and outcomes, the regressions do not establish comparability. If a feature region contains only control orders by design, more control observations do not identify treatment response there. A fitted curve can still extrapolate into that region.

Separate it from S-, T-, R- and DR-learning

An S-learner uses one outcome model with treatment as a feature and contrasts its two treatment settings. A T-learner subtracts two separately fitted outcome predictions. The X-learner adds effect-label regression and blending. Nie and Wager, §6.1 describe these baselines.
An R-learner instead fits an effect function through a loss relating outcome residuals to treatment residuals. A DR-learner regresses an augmented inverse-probability-weighted pseudo-outcome on features. A propensity-weighted X-learner blend does not confer the latter's double robustness: gg combines effect predictions; it does not supply the DR pseudo-outcome's inverse-probability residual correction. See Nie and Wager, §§2–3 and 5 and Kennedy, §§4.1–4.2.
For flexible models, separate nuisance fitting, tuning and final evaluation. Cross-fitting can keep a held-out fold's outcomes out of the models used to construct its pseudo-outcomes; it does not change the X-learner into an orthogonal or doubly robust procedure. Second-stage predictions also need independent evaluation, with time and dependence respected for financial data. No individual causal label becomes available merely by holding an observation out.
Conditional Expectation (membership; public preview) explains the averaging foundation: regression targets a conditional mean of noisy labels. It does not make each label an individual counterfactual.

Sources

  • Künzel, Sekhon, Bickel and Yu (2019), Metalearners for estimating heterogeneous treatment effects using machine learning. PNAS article, X-Learner, equations 5–9; Remark 1; Intuition Behind the Metalearners.
  • Nie and Wager, Quasi-Oracle Estimation of Heterogeneous Treatment Effects. Primary manuscript, §§2–3, 5 and 6.1.
  • Kennedy, Towards Optimal Doubly Robust Estimation of Heterogeneous Causal Effects. Primary manuscript, §§2 and 4.1–4.2.
Last updated
  • causal inference
  • heterogeneous treatment effects
  • metalearners
  • imputation
← Back to the glossaryBack to top ↑