Learning Subgroups with Maximum Treatment Effects without Causal Heuristics
Lincen Yang, Zhong Li, Matthijs van Leeuwen, Saber Salehkaleybar
TL;DR
The paper addresses discovering the subgroup that yields the maximum average treatment effect under a structural causal model. It proposes a partition-based theory showing that the maximizer must have homogeneous pointwise treatment effects, which allows reducing the task to standard supervised learning; the method is instantiated with CART and evaluated using honest inference on synthetic and semi-synthetic data. Results indicate superior performance of this simple, causally grounded approach over several heuristic-based baselines, suggesting causal heuristics may be unnecessary for maximum-effect subgroup discovery under the partition-based model. The work provides a principled, general recipe: learn a supervised partition to recover the data-generating structure, then estimate subgroup effects and select the best performing subgroup, with practical implications for precision medicine, policy, and education.
Abstract
Discovering subgroups with the maximum average treatment effect is crucial for targeted decision making in domains such as precision medicine, public policy, and education. While most prior work is formulated in the potential outcome framework, the corresponding structural causal model (SCM) for this task has been largely overlooked. In practice, two approaches dominate. The first estimates pointwise conditional treatment effects and then fits a tree on those estimates, effectively turning subgroup estimation into the harder problem of accurate pointwise estimation. The second constructs decision trees or rule sets with ad-hoc 'causal' heuristics, typically without rigorous justification for why a given heuristic may be used or whether such heuristics are necessary at all. We address these issues by studying the problem directly under the SCM framework. Under the assumption of a partition-based model, we show that optimal subgroup discovery reduces to recovering the data-generating models and hence a standard supervised learning problem (regression or classification). This allows us to adopt any partition-based methods to learn the subgroup from data. We instantiate the approach with CART, arguably one of the most widely used tree-based methods, to learn the subgroup with maximum treatment effect. Finally, on a large collection of synthetic and semi-synthetic datasets, we compare our method against a wide range of baselines and find that our approach, which avoids such causal heuristics, more accurately identifies subgroups with maximum treatment effect. Our source code is available at https://github.com/ylincen/causal-subgroup.