Meta-Learners for CATE Estimation

"While methods like Double Machine Learning excel at estimating the Average Treatment Effect (ATE), many applications require a more granular understanding: how does the effect of a treatment $T$ on an outcome $Y$ vary across individuals with different characteristics $X$? This varying effect is the Conditional Average Treatment Effect, or CATE:"

$$\tau(x) = E[Y(1) - Y(0) | X=x]$$

Estimating CATE in high dimensions presents unique challenges. We need methods that can flexibly model the relationship between covariates $X$ and the treatment effect itself, without imposing strong parametric assumptions. Meta-learners provide a powerful framework for achieving this by leveraging existing supervised machine learning algorithms (the "base learners") to construct CATE estimators. They are termed "meta"-learners because they operate on top of standard prediction models.

We will examine three prominent meta-learners: the S-Learner, T-Learner, and X-Learner. Each offers a different strategy for repurposing supervised learning algorithms for CATE estimation.

The S-Learner: Simplicity through Feature Inclusion

The S-Learner (Single-Learner) takes the most straightforward approach. It includes the treatment indicator $T$ as a regular feature alongside the covariates $X$ in a single model trained to predict the outcome $Y$.

Let $\mu(x, t) = E[Y | X=x, T=t]$ be the expected outcome conditional on covariates and treatment assignment. The S-Learner trains a single supervised learning model $\hat{\mu}$ (e.g., gradient boosting, random forest, neural network) using the full dataset $(X_i, T_i, Y_i)$ to approximate $\mu(x, t)$.

The CATE is then estimated by taking the difference between the model's predictions with $T$ set to 1 and $T$ set to 0 for a given $x$:

$$\hat{\tau}_S(x) = \hat{\mu}(x, T=1) - \hat{\mu}(x, T=0)$$

Strengths:

• Simplicity: Easy to implement using any standard supervised learning library.
• Direct Outcome Modeling: Directly models the response surface $E[Y|X, T]$.

Weaknesses:

• Treatment Effect Regularization: If the treatment effect $\tau(x)$ is small relative to the main effect of $X$ on $Y$, regularization within the base learner $\hat{\mu}$ might attenuate or even zero out the estimated effect. The model prioritizes overall predictive accuracy for $Y$, not specifically the accuracy of the difference related to $T$.
• Model Specificity: Assumes the base learner is suitable for capturing both the main outcome relationship and the treatment effect heterogeneity using the same functional form and regularization. This might not hold if, for instance, the outcome $Y$ is smoothly related to $X$, but the treatment effect $\tau(x)$ has sharp discontinuities.

The S-Learner serves as a useful baseline but often struggles when treatment effects are subtle or have a structure significantly different from the baseline outcome function.

The T-Learner: Separate Models for Treatment and Control

The T-Learner (Two-Learner) adopts a more direct approach to modeling the potential outcomes. It builds two separate supervised learning models: one for the outcome under treatment and one for the outcome under control.

1. Model for Treated Group: Train a model $\hat{\mu}_1$ to predict $Y$ using only the data points where $T=1$. This model estimates $\mu_1(x) = E[Y | X=x, T=1]$.
2. Model for Control Group: Train a second model $\hat{\mu}_0$ to predict $Y$ using only the data points where $T=0$. This model estimates $\mu_0(x) = E[Y | X=x, T=0]$.

The CATE is then estimated as the difference between the predictions of these two models:

$$\hat{\tau}_T(x) = \hat{\mu}_1(x) - \hat{\mu}_0(x)$$

Flow of the T-Learner approach, splitting data to train separate outcome models for treated and control groups.

Strengths:

• Direct Potential Outcome Modeling: Explicitly models the response surfaces for the treated and control groups separately.
• Flexibility: Allows using different base learners or hyperparameter settings for $\hat{\mu}_1$ and $\hat{\mu}_0$, potentially accommodating different complexities in each group.

Weaknesses:

• Data Sparsity: If one treatment group is much smaller than the other (imbalanced treatment assignment), the model for the smaller group might perform poorly due to insufficient data.
• Error Propagation: The final CATE estimate combines errors from two independent models.
• Ignores Shared Information: Fails to leverage potential similarities in the relationship between $X$ and $Y$ across the two treatment groups. Information learned by $\hat{\mu}_1$ about the effect of $X$ does not directly inform $\hat{\mu}_0$, and vice versa.

The T-Learner often performs better than the S-Learner when the treatment effect is substantial, but its performance can degrade with imbalanced treatments.

The X-Learner: Leveraging Imputed Treatment Effects

The X-Learner, proposed by Künzel et al. (2019), is designed to address the shortcomings of both S and T-Learners, particularly in scenarios with imbalanced treatment groups or complex CATE functions. It employs a multi-stage estimation strategy.

Stage 1: Estimate Outcome Models (like T-Learner) First, estimate the separate outcome models $\hat{\mu}_1(x)$ and $\hat{\mu}_0(x)$ exactly as in the T-Learner, using the treated ($T=1$) and control ($T=0$) data, respectively.

Stage 2: Impute Individualized Treatment Effects Next, use the models from Stage 1 to impute the counterfactual outcomes for each individual and calculate imputed treatment effects:

• For individuals in the treated group ($T_i=1$), impute the effect as: $$\tilde{D}_i^1 = Y_i^{\text{obs}} - \hat{\mu}_0(X_i)$$ This represents the observed outcome minus the predicted outcome had they not received the treatment.
• For individuals in the control group ($T_i=0$), impute the effect as: $$\tilde{D}_i^0 = \hat{\mu}_1(X_i) - Y_i^{\text{obs}}$$ This represents the predicted outcome had they received the treatment minus their observed outcome.

Stage 3: Estimate CATE using Imputed Effects Now, treat the imputed effects $\tilde{D}^1$ and $\tilde{D}^0$ as target variables. Train two new supervised learning models to predict these imputed effects based on the covariates $X$:

• Train model $\hat{\tau}_1(x)$ using the dataset $(X_i, \tilde{D}_i^1)$ for all units where $T_i=1$.
• Train model $\hat{\tau}_0(x)$ using the dataset $(X_i, \tilde{D}_i^0)$ for all units where $T_i=0$. These models directly learn the relationship between covariates $X$ and the estimated treatment effects within each group.

Stage 4: Combine Estimates with Weighting Finally, combine the two CATE estimates $\hat{\tau}_1(x)$ and $\hat{\tau}_0(x)$ using a weighting function $g(x)$. A common choice for $g(x)$ is an estimate of the propensity score, $\hat{e}(x) = P(T=1 | X=x)$:

$$\hat{\tau}_X(x) = g(x) \hat{\tau}_0(x) + (1 - g(x)) \hat{\tau}_1(x)$$

This weighting scheme gives more influence to the estimate derived from the larger group for a given $x$. For example, if propensity scores are low ($g(x) \approx 0$), meaning individuals with covariates $x$ are rarely treated, the final estimate relies more on $\hat{\tau}_1(x)$, which was learned using the imputed effects for the treated group (where we have more information about $Y(1)$). Conversely, if $g(x) \approx 1$, the estimate relies more on $\hat{\tau}_0(x)$.

Strengths:

• Effective with Imbalanced Data: By imputing effects and modeling them separately, it uses data from the larger group to improve estimates relevant to the smaller group.
• Direct CATE Modeling: Stage 3 specifically targets the CATE function $\tau(x)$.
• Asymptotic Properties: Possesses desirable theoretical properties under certain conditions.

Weaknesses:

• Complexity: Involves multiple estimation steps, increasing implementation complexity and potential points of failure.
• Error Propagation: Errors from Stage 1 models ($\hat{\mu}_0, \hat{\mu}_1$) propagate into the imputed effects ($\tilde{D}^0, \tilde{D}^1$) and subsequently into the Stage 3 models ($\hat{\tau}_0, \hat{\tau}_1$).
• Propensity Score Estimation: Requires an estimate of the propensity score $g(x)$ for weighting, introducing another potential source of error if the propensity model is misspecified.

The X-Learner is often the most sophisticated and best-performing meta-learner, especially when dealing with significant heterogeneity or data imbalances, but its complexity demands careful implementation and validation.

Choosing and Implementing Meta-Learners

The choice between S, T, and X-Learners depends on the specific characteristics of the problem:

• S-Learner: A simple baseline, potentially suitable if treatment effects are expected to be large and smoothly related to covariates in a way similar to the main outcome function.
• T-Learner: A good choice when treatment groups are relatively balanced and separate modeling seems appropriate.
• X-Learner: Generally preferred when treatment groups are imbalanced, CATE is complex, or maximizing the use of all available data is critical.

Implementation Considerations:

• Base Learners: The performance of any meta-learner heavily depends on the choice and tuning of the underlying supervised ML models ($\hat{\mu}$, $\hat{\mu}_0$, $\hat{\mu}_1$, $\hat{\tau}_0$, $\hat{\tau}_1$). Models capable of handling high dimensions and complex interactions (e.g., Gradient Boosting, Random Forests, Neural Networks) are typically used. Be mindful of how regularization in these base learners might impact CATE estimation (especially for the S-Learner).
• Propensity Scores (for X-Learner): The weighting function $g(x)$ in the X-Learner requires a propensity score model. The quality of this model can influence the final CATE estimates.
• Software: Libraries like causalml and dowhy in Python offer implementations of these meta-learners, simplifying their application. For example, causalml.inference.meta provides classes like BaseSLearner, BaseTLearner, and BaseXLearner that allow plugging in scikit-learn compatible estimators as base learners.

Meta-learners provide a flexible and powerful bridge between standard supervised learning and causal effect estimation. By cleverly structuring the prediction tasks, they allow us to estimate complex, heterogeneous treatment effects using familiar machine learning tools, a significant step towards understanding individualized impacts in high-dimensional settings. Remember that, like all methods relying on covariate adjustment, meta-learners fundamentally depend on the unconfoundedness assumption: $(Y(1), Y(0)) \perp T | X$. They estimate CATE based on the observed covariates $X$, assuming these are sufficient to control for confounding bias. Evaluating the robustness of these estimates, as discussed later, remains an important part of the workflow.