Robust X-Learner: Breaking the Curse of Imbalance and Heavy Tails via Robust Cross-Imputation

TL;DR

This work addresses the joint challenges of extreme imbalance and heavy-tailed outcomes in conditional average treatment effect estimation. It introduces the Robust X-Learner (RX-Learner), which replaces the standard MSE objective with a gamma-divergence (equivalently the Welsch loss under Gaussian core) to down-weight outliers, and optimizes stably via a Majorization-Minimization based proxy Hessian. RX-Learner preserves cross-imputation benefits while mitigating Outlier Smearing by applying robust base learners, robust cross-imputation, and robust aggregation, validated on semi-synthetic Criteo uplift data and synthetic simulations. The results show dramatic reductions in core population error (PEHE) and strong stability against contamination, delivering reliable CATE estimates for industrial decision-making and budget allocation. Overall, RX-Learner offers a practical, theoretically grounded approach to structural robustness in CATE estimation under real-world data messiness.

Abstract

Estimating Heterogeneous Treatment Effects (HTE) in industrial applications such as AdTech and healthcare presents a dual challenge: extreme class imbalance and heavy-tailed outcome distributions. While the X-Learner framework effectively addresses imbalance through cross-imputation, we demonstrate that it is fundamentally vulnerable to "Outlier Smearing" when reliant on Mean Squared Error (MSE) minimization. In this failure mode, the bias from a few extreme observations ("whales") in the minority group is propagated to the entire majority group during the imputation step, corrupting the estimated treatment effect structure. To resolve this, we propose the Robust X-Learner (RX-Learner). This framework integrates a redescending γ-divergence objective -- structurally equivalent to the Welsch loss under Gaussian assumptions -- into the gradient boosting machinery. We further stabilize the non-convex optimization using a Proxy Hessian strategy grounded in Majorization-Minimization (MM) principles. Empirical evaluation on a semi-synthetic Criteo Uplift dataset demonstrates that the RX-Learner reduces the Precision in Estimation of Heterogeneous Effect (PEHE) metric by 98.6% compared to the standard X-Learner, effectively decoupling the stable "Core" population from the volatile "Periphery".

Figures (4)

• Figure 1: Qualitative Comparison on 1D Data. The MSE-based learner (Red Dashed Line) is severely pulled upward by the outliers (Red Stars), failing to capture the true function structure. The RX-Learner (Blue Solid Line) effectively ignores the outliers, recovering the true underlying generator.
• Figure 2: Sensitivity Analysis. PEHE error vs. Contamination Rate. The Standard X-Learner's error explodes immediately. The RX-Learner remains stable, demonstrating a high breakdown point.
• Figure 3: Stability Analysis over 5 Trials. The box plot of Core-PEHE shows that the Baseline X-Learner suffers from extreme variance, making it operationally risky for campaign management. In contrast, the RX-Learner maintains consistently low error, demonstrating robust convergence across different noise realizations.
• Figure 4: Visualizing Outlier Smearing. Comparison of Predicted vs. True CATE on the Core population. (Left) The Baseline X-Learner is heavily biased upward by the outliers, illustrating the "Smearing" effect where extreme noise is mistaken for signal. (Right) The RX-Learner effectively suppresses the outlier influence, recovering the true underlying causal heterogeneity.