Technical Summary: Robust X-Learner (RX-Learner)

1. Problem Statement

The paper addresses a critical dual challenge in estimating Heterogeneous Treatment Effects (HTE) within industrial applications like AdTech and healthcare: extreme class imbalance and heavy-tailed outcome distributions.

While the standard X-Learner framework (Künzel et al., 2019) effectively mitigates imbalance through cross-imputation (borrowing information from the majority group to estimate effects for the minority), it fails catastrophically when data contains heavy-tailed outliers ("whales"). The authors identify a specific failure mode termed "Outlier Smearing."

In the standard X-Learner, which relies on Mean Squared Error (MSE) minimization:

1. Bias Propagation: If the minority group (e.g., treatment group) contains even a single extreme outlier, the MSE-based response function estimator ($\hat{\mu}_1$) becomes significantly biased.
2. Smearing Mechanism: During the cross-imputation step, this biased $\hat{\mu}_1$ is used to generate pseudo-outcomes for the entire majority (control) group. Consequently, the bias from a single "whale" is mathematically "smeared" across every imputed value in the clean majority group.
3. Result: 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.

Existing alternatives like Robust Double Machine Learning (R-Learner) fail in this context because they rely on inverse probability weighting (IPW), which suffers from variance explosion when propensity scores approach zero or one (extreme imbalance).

2. Methodology: The Robust X-Learner (RX-Learner)

The RX-Learner preserves the structural advantages of the X-Learner (cross-imputation) but fundamentally alters the learning objective to enforce robustness against heavy tails. The framework is built on the concept of Core-Periphery Causal Inference, aiming to estimate effects on the stable "Core" population while ignoring the volatile "Periphery."

2.1 Robust Base Learners via $\gamma$-Divergence

Instead of minimizing MSE, the RX-Learner minimizes a density-power divergence objective. Under the assumption that the "Core" distribution is Gaussian, this objective is mathematically equivalent to the Welsch (or Leclerc) loss:
$$L_{Robust}(r_i) \propto 1 - \exp\left(-\frac{\gamma r_i^2}{2\hat{\sigma}^2}\right)$$

• Redescending Influence: The loss function possesses a redescending influence function. As the residual magnitude $|r_i|$ increases (indicating a "whale"), the weight assigned to that observation decays exponentially to zero.
• Oracle Refinement: This mechanism allows the model to automatically "soft-trim" outliers during training, fitting the function solely to the Core population.
• Scale Estimation: To prevent singularity (implosion), the scale parameter $\hat{\sigma}$ is initialized and fixed using the Median Absolute Deviation (MAD) of residuals from a preliminary robust fit.

2.2 MM-Based Optimization (Proxy Hessian)

Since the Welsch loss is non-convex, standard gradient boosting implementations (which rely on convex Hessians) may become unstable. The authors introduce a Majorization-Minimization (MM) strategy:

• Quadratic Majorization: The Welsch loss is treated as a half-quadratic function. At each iteration, a quadratic upper bound (majorizer) is constructed.
• Proxy Hessian: The optimization is performed by minimizing this quadratic bound, which is equivalent to solving a weighted least squares problem where the weights act as a "Proxy Hessian."
• Guarantee: This approach ensures a monotonic decrease in the loss function, guaranteeing stability even without a convex loss surface.

2.3 Robust Cross-Imputation and Aggregation

• Step 2 (Imputation): Because the base learners ($\hat{\mu}_{1,\gamma}, \hat{\mu}_{0,\gamma}$) are robust, the bias term $\delta$ is effectively zero. The pseudo-outcomes generated for the majority group are therefore "clean," preventing the smearing effect at its source.
• Step 3 (Aggregation): The final CATE models are also trained using the same $\gamma$-boosting procedure on the pseudo-outcomes. The final estimate is a weighted average using an inverse-variance weighting scheme, prioritizing the estimator with higher precision.

3. Key Contributions

1. Theoretical Integration: The paper integrates $\gamma$-divergence minimization into the meta-learning framework, establishing a theoretical bridge between modern density-power divergence methods and classical robust statistics (Welsch loss). It identifies "Outlier Smearing" as a fundamental vulnerability of standard X-Learners in imbalanced, heavy-tailed settings.
2. Algorithmic Stability: It solves the non-convex optimization challenge inherent in robust estimation by grounding the "Proxy Hessian" strategy in the MM framework, enabling efficient implementation on standard Gradient Boosting machinery.
3. Empirical Validation: The method is validated on a semi-synthetic derivation of the Criteo Uplift dataset, demonstrating that it effectively decouples the stable Core population from the volatile Periphery.

4. Experimental Results

The authors evaluated the RX-Learner against standard baselines (MSE-based X-Learner, Huber loss, Winsorization, and Robust DML) using both pure synthetic data and a semi-synthetic Criteo benchmark.

• Synthetic Data (Extreme Pathology): In a scenario with 2% treatment imbalance and asymmetric Pareto noise, the RX-Learner reduced the Precision in Estimation of Heterogeneous Effect (PEHE) by over 50% compared to the next best baseline (Huber loss). Crucially, while MSE-based methods showed a linear increase in prediction shift proportional to outlier magnitude, the RX-Learner showed zero shift, confirming the rejection of outliers.
• Semi-Synthetic Criteo Benchmark:
  • Setup: $N=61,200$ users, 2% treatment rate, with "whales" (5% of treated users) having values 80x larger than the signal.
  • Baseline Failure: The standard X-Learner (MSE) exhibited catastrophic instability, with a standard deviation (24.54) exceeding the mean error (18.15).
  • RX-Learner Success: The RX-Learner achieved a 98.6% reduction in Core-PEHE (0.263 vs. 18.15) with negligible variance (0.022).
  • Mechanism Verification: The results confirmed that the baseline model suffered from "Outlier Smearing," where the model's predictions collapsed globally due to a single outlier, whereas the RX-Learner successfully recovered the subtle underlying causal signal.

5. Significance and Claims

The paper claims that the RX-Learner provides a mathematically grounded solution for industrial-scale CATE estimation where data is rarely Gaussian or balanced.

• Structural Robustness: The method moves beyond "convenient statistical assumptions" to respect the geometry of real-world data (Core vs. Periphery).
• Business Impact: By preventing "Whale Chasing" (bidding on users who are already converting organically), the RX-Learner offers a conservative but trustworthy estimator for budget allocation, ensuring ad spend targets incremental conversions rather than subsidizing outliers.
• Limitations: The authors note that the current implementation assumes a Gaussian Core distribution for computational stability. Future work is suggested to explore non-Gaussian extensions (e.g., Student-t or Cauchy) via Natural Gradient Boosting and dynamic tuning of the robustness parameter $\gamma$.

In summary, the RX-Learner resolves the tension between the X-Learner's ability to handle imbalance and the need for robustness against heavy tails, transforming "information borrowing" from a vector of "pollution spreading" into a stable estimation process.