Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
48,923 characters · 51 sections · 27 citation commands
Robust X-Learner: Breaking the Curse of Imbalance and Heavy Tails via Robust Cross-Imputation
In the era of large-scale observational data, the focus of causal inference has decisively shifted from estimating the Average Treatment Effect (ATE) to understanding the Conditional Average Treatment Effect (CATE). In domains ranging from precision medicine to digital marketing, decision-makers require granular insights: identifying which patient subgroups respond to a specific drug, or which user segments yield incremental lift from an advertisement ate_to_cate.
However, the transition from academic theory to industrial practice is often obstructed by the inherent messiness of real-world data generating processes (DGP). While pioneering approaches like recursive partitioning ate_to_cate have successfully addressed the challenge of high-dimensional covariates, industrial data often exhibits specific pathologies—such as heavy-tailed noise and extreme outliers—that violate the sub-Gaussian assumptions typically required for these estimators to function optimally.
We identify two distinct but often co-occurring barriers to effective CATE estimation in modern applications:
While the X-Learner was explicitly designed to handle the first challenge (Imbalance) by employing a cross-imputation strategy xlearner, it remains defenseless against the second (Heavy Tails).
A naive approach might suggest using Robust Double Machine Learning (DML) dml, R-Learners rlearner, or tree-based methods like Causal Forests ate_to_cate. However, standard splitting criteria in causal trees are typically based on MSE reductions, which remain sensitive to extreme values. Furthermore, R-Learners rely on inverse probability weighting (IPW), which becomes numerically unstable when the overlap assumption is threatened by extreme imbalance (i.e., propensity scores $\pi(x)$ approach 0 or 1).
Conversely, the standard X-Learner handles imbalance well but fails catastrophically under heavy tails due to a phenomenon we term Outlier Smearing.
Consider the first stage of the X-Learner, where response functions $\mu_1(x)$ and $\mu_0(x)$ are estimated. If the treatment group is small and contains even a single outlier, an MSE-based $\hat{\mu}_1(x)$ will be significantly biased upward. In the subsequent imputation step, the pseudo-outcome for the control group is calculated as:
Because $\hat{\mu}_1(x)$ is contaminated globally or locally by a few whales, this bias is effectively “smeared” across every imputed value in the control group. Consequently, the final CATE estimator learns a distorted structure, mistaking the noise of the outlier for a genuine treatment effect, even in regions where the control data is dense and clean.
To overcome this dilemma, we propose the Robust X-Learner (RX-Learner), a unified framework that preserves the structural advantages of the X-Learner while enforcing robustness against heavy tails.
Our approach rests on the concept of Core-Periphery Causal Inference. We posit that the data consists of a stable “Core” mechanism and a volatile “Periphery” (tail). Our goal is to estimate the causal effect on the Core without being misled by the Periphery. The RX-Learner achieves this through three key technical innovations:
We validate the RX-Learner on a semi-synthetic derivation of the Criteo Uplift v2.1 dataset. Results demonstrate that our method reduces the error on the Core population by 98.6% compared to the baseline, confirming that robustification is essential for industrial-scale CATE estimation.
The remainder of this paper is organized as follows: Section 2 details the mathematical preliminaries. Section 3 analyzes the failure modes of DML and standard X-Learners. Section 4 presents the RX-Learner algorithm. Section 5 provides simulation results demonstrating the method's superiority, and Section 6 discusses the application to AdTech data.
In this Section, we establish the theoretical framework for estimating Conditional Average Treatment Effects (CATE). We first review the standard potential outcomes framework and the specific mechanics of the original X-Learner. Subsequently, we formally define the “Dual Challenge” of data asymmetry and heavy-tailed contamination, utilizing a mixture model to characterize the “Core-Periphery” structure of industrial data. Finally, we discuss the theoretical properties of robust loss functions relevant to our proposed solution.
We adopt the Neyman-Rubin potential outcomes framework rubin1974. We consider a set of $N$ independent and identically distributed (i.i.d.) units. For each unit $i$, we observe a feature vector $X_i \in \mathcal{X} \subseteq \mathbb{R}^d$, a binary treatment assignment $W_i \in \{0, 1\}$, and an observed outcome $Y_i \in \mathbb{R}$.
We posit the existence of two potential outcomes: $Y_i(1)$ corresponding to the treatment condition, and $Y_i(0)$ corresponding to the control condition. The observed outcome is realized as:
Our primary estimand is the Conditional Average Treatment Effect (CATE), defined as:
where $\mu_w(x) = \mathbb{E}[Y(w) \mid X = x]$ represents the response function for the treatment arm $w \in \{0, 1\}$.
To identify $\tau(x)$ from observational data, we invoke the standard assumptions of Causal Inference:
The X-Learner, introduced by Künzel et al. (2019) xlearner, is a meta-learner specifically designed to maximize information utilization when sample sizes are unbalanced. It proceeds in three steps:
Standard implementations assume that the base learners in Steps 1 and 3 minimize the Mean Squared Error (MSE), implicitly assuming Gaussian residuals with constant variance.
Real-world data, particularly in AdTech and healthcare, violates the implicit assumptions of the standard X-Learner. We formally define the two specific pathologies addressed in this paper.
We consider the regime where the sample sizes of the treatment and control groups are highly asymmetric. Without loss of generality, assume the treatment group is the minority:
In this setting, the response function $\hat{\mu}_1(x)$ is estimated with high variance due to data sparsity, while $\hat{\mu}_0(x)$ is estimated with high precision.
We relax the assumption of sub-Gaussian noise. Instead, we introduce a Core-Periphery Contamination Model. We model the observed outcome $Y$ as being generated from a mixture distribution:
where the error term $\varepsilon_i$ is drawn from a contaminated density:
Under this model, the standard MSE objective function becomes dominated by the tail component $h_{\text{tail}}$. A single observation from the tail can induce an arbitrarily large gradient update, shifting the estimated mean $\hat{\mu}$ away from the Core expectation.
To address Pathology 2, we require a loss function that is robust to the Periphery component. In this work, we focus on the family of density-power divergences.
Specifically, the $\gamma$-divergence fujisawa is known to provide robust parameter estimation by minimizing the discrepancy between a power of the model density and the data density. While generally applicable, its implementation in gradient boosting requires specification of the base distribution.
Connection to Welsch Loss: If we assume the underlying Core distribution is Gaussian, $\phi_{\text{core}} \sim \mathcal{N}(\mu, \sigma^2)$, the minimization of the $\gamma$-divergence mathematically collapses to minimizing the following loss function with respect to the residual $r_i = y_i - \hat{y}_i$:
This formulation is functionally equivalent to the Welsch (or Leclerc) Loss used in robust statistics holland1977.
Crucially, this loss function is non-convex and redescending. Its influence function $\psi(r) = \partial \mathcal{L} / \partial r$ satisfies:
This property effectively performs an “Oracle Refinement”: residuals that are statistically impossible under the Core distribution (i.e., Whales) are assigned near-zero weight, allowing the model to fit the Core structure without distortion. In Section 4, we will demonstrate how to optimize this non-convex objective stably within a boosting framework using MM algorithms hunter2004.
In this Section, we theoretically analyze why existing state-of-the-art methods are insufficient for the dual challenge of imbalance and heavy tails. We first show that Double Machine Learning (DML) dml, even when equipped with robust loss functions, suffers from variance explosion due to propensity score instability. We then dissect the standard X-Learner to demonstrate the mechanism of “Outlier Smearing,” proving that blind information borrowing becomes pollution spreading in the presence of whales.
Double Machine Learning (DML), specifically the R-Learner rlearner, frames CATE estimation as a weighted regression problem based on the Robinson transformation robinson1988. The objective function is typically defined as finding $\tau(x)$ that minimizes:
where $\hat{\mu}(x) = \mathbb{E}[Y|X=x]$ and $\hat{\pi}(x) = P(W=1|X=x)$.
While DML is efficient under standard conditions, its reliance on the variation of treatment assignment creates a fundamental vulnerability in highly imbalanced settings (e.g., $N_1 \ll N_0$).
The effective variance of the R-Learner estimator scales inversely with the variance of the treatment assignment. Theoretically, the asymptotic variance bound involves a term proportional to:
In our setting of extreme imbalance, $\pi(x) \to 0$ for the vast majority of the covariate space $\mathcal{X}$. This leads to a denominator approaching zero, causing the variance of the estimator to explode.
A common counter-argument is to replace the squared loss with a robust loss function $\rho(\cdot)$ (e.g., Huber huber1964 or Quantile loss) to handle outliers. However, this Robust DML approach fails to address the structural issue. Even with a robust loss, the identification of $\tau(x)$ relies on the variation of $W$ after conditioning on $X$. When $\pi(x) \approx 0$, there is almost no variation in $W$ to exploit. The estimator effectively attempts to divide a robustified residual by a near-zero weight, resulting in numerical instability and wide confidence intervals that render the estimate useless for business decision-making.
The X-Learner is structurally superior for imbalance because it does not divide by the propensity score. Instead, it utilizes Cross-Imputation to borrow strength from the control group to estimate the treated group's effect, and vice versa.
However, we prove here that without robust base learners, this strength-borrowing mechanism facilitates the propagation of errors—a phenomenon we term Outlier Smearing.
Consider the standard X-Learner Step 1 (Response Estimation). Let the treatment group ($W=1$) be the minority with sample size $N_1$, containing a single extreme outlier (whale) at index $k$ with value:
The base learner $\hat{\mu}_1$, trained via MSE minimization on the small sample $N_1$, will be heavily biased towards this outlier. For a local region $\mathcal{N}(X_k)$ around the outlier, the estimated response is shifted:
where $\delta > 0$ is the bias induced by the whale. Note that because $N_1$ is small, the leverage of the single outlier is high, making $\delta$ significant.
Now, consider Step 2 (Imputation) for the Control Group (the majority, $N_0$). The pseudo-outcome $\tilde{D}^0$ is calculated as:
Substituting the biased estimator $\hat{\mu}_1$:
The Smearing Effect: Here lies the critical failure. The bias $\delta$, originating from a single unit in the minority group, is added as a constant shift to every single imputed data point in the majority group within that region.
In Step 3, the X-Learner aggregates these pseudo-outcomes. The model $\hat{\tau}_0(x)$ is trained to predict $\tilde{D}^0$. Since the target variable $\tilde{D}^0$ systematically contains the term $+\delta$, the resulting CATE estimate converges to:
Even though the control group $N_0$ is large and clean, the imputed labels provided to it are corrupted. This is “Smearing”: the pollution from the minority tail is smeared across the clean majority structure.
While the final weighting $g(x)\hat{\tau}_0 + (1 - g(x))\hat{\tau}_1$ usually down-weights $\hat{\tau}_0$ when $\pi(x)$ is small, in practice:
Table 1 summarizes the theoretical vulnerabilities of each method.
This analysis confirms that no existing method solves both problems simultaneously. We require a method that borrows information (like X-Learner) but rejects pollution (like Robust Statistics). This motivates the design of the RX-Learner in the next Section.
In this Section, we present the Robust X-Learner (RX-Learner), a novel framework designed to estimate CATE consistently in the presence of extreme imbalance and heavy-tailed contamination.
The RX-Learner preserves the structural advantage of the standard X-Learner (Cross-Imputation) but fundamentally alters the learning objective. By replacing Mean Squared Error (MSE) minimization with a robust divergence minimization and employing an MM-based optimization strategy, we achieve an “End-to-End” robustness that prevents the Smearing effect identified in Section 3.
The first critical step is to estimate the response functions $\mu_1(x)$ and $\mu_0(x)$ without bias from outliers.
As established in Section 2.4, we adopt the $\gamma$-divergence under a Gaussian assumption. For a regression setting with residuals $r_i = y_i - f(x_i)$, we minimize the empirical loss:
where $\gamma > 0$ is a tuning parameter controlling robustness (typically $\gamma \in [0.1, 1.0]$), and $\hat{\sigma}$ is a scale parameter.
The gradient of this loss function with respect to the prediction $F = f(x_i)$ reveals the mechanism of robustness:
where the adaptive weight $w_i(r_i)$ is given by:
The Redescending Property: This weight function $w_i$ is bell-shaped. As the residual magnitude $|r_i| \to \infty$ (i.e., for a “whale”), the weight $w_i \to 0$ exponentially. This creates an Oracle Refinement effect: the model automatically identifies and “soft-trims” outliers during the training process, fitting the function solely to the “Core” population. Unlike MSE, where outliers pull the fit, here outliers are effectively invisible to the gradient.
Robust estimators can suffer from “implosion” where $\hat{\sigma} \to 0$ to fit a single point perfectly. To prevent this, we employ a robust fixed-scale approach. We initialize $\hat{\sigma}$ using the Median Absolute Deviation (MAD) of the residuals from a robust preliminary fit (e.g., Least Absolute Deviation) hampel1986:
This scale serves as a fixed “anchor,” ensuring the definition of the “Core” remains stable throughout boosting.
Standard Gradient Boosting Decision Tree (GBDT) implementations (e.g., XGBoost) require convex objectives to utilize Newton-Raphson updates. Since our loss $\mathcal{L}_\gamma$ is non-convex, the Hessian can be negative, leading to instability.
Instead of a heuristic approximation, we ground our optimization in the Majorization-Minimization (MM) framework idier2001.
The Welsch loss function $\rho(r) = 1 - e^{-r^2}$ is a half-quadratic function geman1992. It admits a quadratic majorizer $Q(r; r^{(t)})$ at any current residual $r^{(t)}$:
Minimizing this quadratic upper bound $Q$ corresponds to solving a weighted least squares problem with fixed weights $w_i = w_i(r^{(t)})$.
Based on this principle, we implement the boosting step as follows. For iteration $t$:
Theorem (Monotone Descent): Since the update minimizes the majorizing function $Q$, and $Q(\theta, \theta) = \mathcal{L}_\gamma(\theta)$, the procedure guarantees that the objective function decreases monotonically: $\mathcal{L}_\gamma(\theta^{(t+1)}) \le \mathcal{L}_\gamma(\theta^{(t)})$. This ensures stability even without a convex loss surface.
With robust base learners $\hat{\mu}_{1,\gamma}$ and $\hat{\mu}_{0,\gamma}$ obtained, we proceed to imputation.
Prevention of Smearing: Since $\hat{\mu}_{1,\gamma}$ effectively ignored the whales in the training phase (due to $w_i \to 0$), the bias term $\delta$ (identified in Section 3) is now 0. Thus, the pseudo-outcomes for the control group $\tilde{D}^0_i$ are calculated using a “clean” response function, preventing the propagation of pollution.
The pseudo-outcomes $\tilde{D}$ may still contain outliers (specifically, if $Y_i$ itself is a whale, $\tilde{D}_i$ will be a whale). Therefore, standard regression in Step 3 is insufficient.
We apply the same $\gamma$-Boosting procedure from Step 1 to estimate the CATE models:
Finally, we aggregate utilizing an Inverse-Variance Weighting scheme. We estimate the local variance $\hat{\sigma}^2_w(x)$ of each estimator (via a separate variance prediction model or quantile regression) and combine them:
This formula correctly assigns higher weight to the estimator with higher precision (lower variance), naturally prioritizing the model trained on the larger/cleaner group.
In this Section, we rigorously evaluate the proposed RX-Learner against a comprehensive set of baselines using fully controlled synthetic data. Unlike the Criteo benchmark, this environment allows us to systematically vary contamination levels and verify the underlying robust mechanisms.
To gain intuition, we first generated a simple 1D dataset where the true treatment effect is fixed, but outcomes are corrupted by “Whales” (outliers). Figure (ref) visualizes the learned functions.
We evaluated performance in an “Extreme Pathology Scenario” designed to mimic AdTech pathologies: $N=2000$, Extreme Imbalance (2% Treated), and Asymmetric Pareto Noise ($\alpha=1.5$). Table (ref) summarizes the results across 5 independent trials.
Result: The RX-Learner outperforms the strongest baseline (Huber) by over 50% in PEHE scores. Standard strategies like Winsorization (1%) fail because the heavy tails result in outliers that are not confined to the top 1%, or because valid signals are discarded.
We varied the contamination rate from 0% to 20%. As shown in Figure (ref), the standard estimator degrades immediately. The RX-Learner maintains stability even under significant contamination.
To confirm the theoretical hypothesis that MSE "smears" outliers across the control group, we injected a single outlier of varying magnitude $\xi$ into the treatment group and measured the average prediction shift in the control group.
Analysis:
Finally, we tested a small-sample regime ($N=100$) with student-$t$ noise ($df=3$) to see if RX-Learner loses efficiency compared to Huber. Surprisingly, RX-Learner (PEHE 0.97) still outperformed Huber (PEHE 1.37). This suggests that even in small samples with moderate tails, the $\gamma$-divergence provides a sharper identification of the core structure than the Pseudo-Huber loss.
In this Section, we bridge the gap between theoretical guarantees and industrial reality. While the pure simulations in Section 5 confirmed the validity of the RX-Learner under idealized conditions, real-world AdTech data involves complex feature interactions, high dimensionality, and non-trivial correlations.
To demonstrate the external validity of our method, we conduct a comprehensive evaluation using a semi-synthetic benchmark derived from the Criteo Uplift Prediction Dataset (v2.1) criteo2018. This approach allows us to utilize realistic covariates while maintaining access to the ground truth CATE for precise error quantification.
Evaluating CATE estimators on real-world data is notoriously difficult due to the “Fundamental Problem of Causal Inference”---we never observe the counterfactual outcome for a given user. Standard offline metrics (e.g., Qini curves) are often noisy surrogates. By employing a Semi-Synthetic approach—using real covariates $X$ from Criteo but generating outcomes $Y$ via a known causal function—we combine the best of both worlds:
We subsampled $N=61,200$ users and applied a stress-test DGP designed to mimic the “Whale” phenomenon observed in mobile gaming and e-commerce.
The Displacement Challenge: Crucially, this setup creates a Displacement Ratio of $\approx 80:1$. In a standard MSE-based loss landscape, a single whale exerts an influence equivalent to 80 regular users. This allows us to strictly test whether the estimator tracks the population mass (the 80 users) or the singular outlier (the 1 whale).
We compared the proposed RX-Learner ($\gamma=0.2$) against the industry-standard X-Learner (XGBoost with MSE). Performance was measured via Core-PEHE (Precision in Estimation of Heterogeneous Effect on the non-outlier population) over 5 independent trials.
Table (ref) presents the quantitative results.
The Baseline X-Learner did not merely perform poorly; it exhibited catastrophic instability. The standard deviation ($24.54$) exceeded the mean ($18.15$). This indicates that the MSE-based model is not statistically consistent in this regime. In trials where the Pareto tail generated particularly extreme values (e.g., a user spending \$10,000), the model's predictions collapsed globally. For practitioners, this implies an unacceptable operational risk: deploying a standard X-Learner is akin to gambling, where a single outlier record can destroy the model's targeting logic overnight.
In stark contrast, the RX-Learner achieved a 98.6% reduction in error with negligible variance ($\sigma=0.022$). This confirms that the $\gamma$-divergence objective successfully functions as an automated gatekeeper. By down-weighting residuals that exceed the “Core” distribution, the learner effectively ignores the 80:1 leverage of the whales, recovering the subtle causal signal ($\approx 0.66$) buried in the noise.
While this experiment measures estimation error (PEHE), the implications for business KPIs (ROAS) are direct.
This comprehensive evaluation on the Criteo benchmark validates that “Outlier Smearing” is not just a theoretical artifact but a pervasive issue in realistic feature spaces. The RX-Learner provides a mathematically grounded solution that is robust enough for the heavy-tailed reality of modern AdTech.
In this work, we addressed a critical gap in the causal inference literature: the simultaneous presence of extreme class imbalance and heavy-tailed outcome distributions. While the standard X-Learner xlearner provides a robust structural framework for handling imbalance, we demonstrated that it is fundamentally vulnerable to Outlier Smearing when the base learners rely on Mean Squared Error (MSE).
Our theoretical analysis revealed that in asymmetric settings, a single outlier (whale) in the minority group does not merely corrupt local predictions; via the cross-imputation mechanism, it systematically shifts the pseudo-outcomes for the entire majority group. This transforms “information borrowing” into “pollution spreading.”
To resolve this, we introduced the Robust X-Learner (RX-Learner). Our contribution is three-fold:
While the RX-Learner represents a significant advance, several avenues for future research remain:
Data in the real world is rarely Gaussian and rarely balanced. As causal inference moves from the clean environment of clinical trials to the messy reality of industrial logs, our methods must evolve. The RX-Learner represents a step towards Structural Robustness—an approach that respects the geometry of the data (Core vs. Periphery) rather than forcing it to fit convenient statistical assumptions.