Uncovering Treatment-Responsive Subgroups Under Direct Bias Control

ICLR 2026 Conference Submission21143 Authors

Keywords: causal inference, heterogeneous treatment effects, covariate balancing, weighting estimators

Abstract: Identifying subpopulations that benefit most or least from a treatment is central to scientific research and policy analysis. We propose an optimization-based framework for learning such subgroups from observational data. The proposed methods discover subgroups exhibiting maximal treatment-effect heterogeneity while enforcing covariate balance, thus directly controlling confounding without explicitly modeling treatment or outcome mechanisms. The framework accommodates flexible subgroup definitions, allowing additional constraints such as fairness criteria to be incorporated. We show that our approach admits flexible nonparametric estimators and enjoys finite-sample error guarantees. We also introduce a principled rule for subgroup assignment based on observed covariates. Simulated and real-world experiments demonstrate substantial improvements over existing approaches.

Primary Area: causal reasoning

Submission Number: 21143