Stable Heterogeneous Treatment Effect Estimation across Out-of-Distribution Populations

TL;DR

The paper tackles stable heterogeneous treatment effect estimation under distribution shift by introducing Stable Balanced Representation Learning with Hierarchical-Attention Paradigm (SBRL-HAP). The approach couples a Balancing Regularizer and an Independence Regularizer within a Hierarchical-Attention Paradigm to jointly reduce selection bias and decorrelate unstable features, enabling robust HTE estimation across out-of-distribution populations. Empirical results on synthetic and real-world data show that SBRL-HAP yields notable reductions in PEHE and ATE bias (average improvements of about 10% and up to 14%, respectively) compared with state-of-the-art methods, with the largest gains under pronounced distribution shifts. The framework is modular and backboned on existing balanced representation methods, highlighting practical impact for real-world decision-making under evolving population distributions, such as in medicine and policy evaluation.

Abstract

Heterogeneous treatment effect (HTE) estimation is vital for understanding the change of treatment effect across individuals or subgroups. Most existing HTE estimation methods focus on addressing selection bias induced by imbalanced distributions of confounders between treated and control units, but ignore distribution shifts across populations. Thereby, their applicability has been limited to the in-distribution (ID) population, which shares a similar distribution with the training dataset. In real-world applications, where population distributions are subject to continuous changes, there is an urgent need for stable HTE estimation across out-of-distribution (OOD) populations, which, however, remains an open problem. As pioneers in resolving this problem, we propose a novel Stable Balanced Representation Learning with Hierarchical-Attention Paradigm (SBRL-HAP) framework, which consists of 1) Balancing Regularizer for eliminating selection bias, 2) Independence Regularizer for addressing the distribution shift issue, 3) Hierarchical-Attention Paradigm for coordination between balance and independence. In this way, SBRL-HAP regresses counterfactual outcomes using ID data, while ensuring the resulting HTE estimation can be successfully generalized to out-of-distribution scenarios, thereby enhancing the model's applicability in real-world settings. Extensive experiments conducted on synthetic and real-world datasets demonstrate the effectiveness of our SBRL-HAP in achieving stable HTE estimation across OOD populations, with an average 10% reduction in the error metric PEHE and 11% decrease in the ATE bias, compared to the SOTA methods.

Figures (6)

• Figure 1: Two main challenges in stable HTE estimation across OOD populations: (C1) selection bias from imbalanced treatment assignment, and (C2) distribution shift across various populations. The former is manifested as imbalanced distributions of covariates (e.g., age) between treated (i.e., T=1) and control (i.e., T=0) units in a specific population. The latter occurs frequently in real-world applications, resulting in out-of-distribution populations that have distinct covariate distributions from the training dataset. This work is among the first to synergistically resolve both selection bias and distribution shift.
• Figure 2: The framework of Stable Balanced Representation Learning with Hierarchical-Attention Paradigm (SBRL-HAP). SBRL-HAP consists of three modules: i. Balancing Regularizer restricts IPM for balanced representation, ii. Independence Regularizer eliminates feature dependence measured by HSIC-RFF for generalization, and iii. Hierarchical-Attention Paradigm decorrelates features comprehensively with a hierarchy for dispelling the interaction between balance and independence. With high flexibility, SBRL-HAP can be plugged into most balanced representation methods by replacing the neural network backbone.
• Figure 3: Results of PEHE on synthetic data $\text{Syn}\_16\_16\_16\_2$ with different bias rate $\rho$ for the testing set. All models are trained with $\rho=2.5$.
• Figure 4: Results of $F_1$ scores on synthetic data $\text{Syn}\_16\_16\_16\_2$ with different bias rate $\rho$ for the testing set. All models are trained with $\rho=2.5$.
• Figure 5: Nonlinear correlation among features in the balanced representation. As shown, the feature correlation is reduced by our SBRL, and further decreased by incorporating HAP.

Theorems & Definitions (2)

• Definition 3.1: Individual Treatment Effect
• Definition 3.2: Average Treatment Effect