Learning Subgroups with Maximum Treatment Effects without Causal Heuristics

AI summary

Overview

Research area: Causal machine learning and subgroup discovery — specifically, finding the subgroup of a population that benefits most from a treatment, studied through structural causal models (SCMs) rather than the potential-outcome framework.

Technical level: Advanced. The paper's central contribution is a set of formal theorems proved with do-calculus and SCM notation, though the practical takeaway (fit a standard decision tree) is simple enough for a general machine learning audience.

Scope: The paper proves that under a partition-based SCM, discovering the maximum-effect subgroup reduces to recovering the data-generating partition — a standard supervised learning problem — and then validates that claim empirically against seven causal baselines on synthetic and semi-synthetic data.

What This Paper Is About

Practitioners who want to find "which subgroup benefits most from this treatment" usually rely on either (a) estimating treatment effects for every individual and then fitting a tree to those estimates, or (b) growing trees or rule sets with hand-crafted "causal" splitting criteria. Both routes embed assumptions that have never been rigorously justified for the maximum-effect goal.

This paper asks whether those special-purpose causal heuristics are necessary at all. Working directly in the structural causal model framework, the authors show that the maximum-effect subgroup must be one of the regions of the underlying data-generating partition, which means the task collapses into ordinary classification or regression — no bespoke causal machinery required.

Key Contributions

1. A theoretical grounding of maximum-effect subgroup discovery in SCMs. The authors state the problem formally under a structural causal model (with the DAG X → T → Y and no hidden confounders), define the subgroup treatment effect A(Q), and prove it is identifiable from observational data (Proposition 1).

2. A homogeneity theorem for optimal subgroups (Theorem 1). For any subset Q′ ⊂ Q with A(Q′) ≤ A(Q), the complement Q ∖ Q′ must satisfy A(Q ∖ Q′) ≥ A(Q), with equality only when A(Q′) = A(Q). Consequently, a maximum-effect subgroup must have a homogeneous pointwise conditional treatment effect — otherwise it could be split and one part would have a higher average effect.

3. The main reduction theorem (Theorem 2). Under a partition-based model where the feature space is divided into disjoint subsets {K_i} and P(Y | X, do(T = t)) is constant within each K_i, the subgroup maximizing A(Q) is necessarily one of the K_i. Therefore, learning the maximum-effect subgroup reduces to recovering the data-generating partition — i.e., to standard supervised learning — followed by evaluating each learned subset and picking the best.

4. Empirical validation of the reduction with plain CART. The authors instantiate the approach with CART (Gini index for discrete targets, MSE for continuous targets) plus cost-complexity pruning and cross-validation, use "honest" inference borrowed from Causal Tree (train on one split, estimate subgroup effects on a held-out split), and show it beats seven causal baselines on synthetic and semi-synthetic benchmarks. Code is released at https://github.com/ylincen/causal-subgroup.

Main Findings

• Causal heuristics are unnecessary for the maximum-effect objective. Theorem 2 shows that any partition-based learner used for classification or regression — tree-based or rule-based — can serve as the subgroup learner, because the target of the search is the data-generating partition itself, not a specially crafted causal criterion.

• Effect-maximizing heuristics are statistically fragile. Naively maximizing Â(Q) during tree growth requires accurate estimation of both P(Y = 1 | T = 1, X ∈ Q) and P(Y = 1 | T = 0, X ∈ Q). Under treatment/control imbalance, one term pushes toward shrinking the subgroup and the other toward expanding it — conflicting objectives that produce high variance or overfitting. Optimizing the Gini index avoids this, deferring effect estimation to an honest stage where both treatment arms are adequately supported.

• Strong performance on semi-synthetic data. Across 77 semi-synthetic datasets generated by the ACIC-2016 simulator from real Infant Health and Development Program (IHDP) features, the proposed CART-based method achieved the highest mean ground-truth subgroup treatment effect at 10.540, and ranked first on 51.9% of all datasets.

• Baselines lag on semi-synthetic data. Mean ground-truth subgroup treatment effects were 7.843 (CausalTree), 7.451 (DistillTree), 7.410 (CURLS), 6.280 (InteractionTree), 5.135 (QUINT), and 4.622 (SIDES). Proportions of datasets ranked first were 0.180 (CURLS), 0.143 (CausalTree), 0.130 (DistillTree), 0.039 (InteractionTree), 0.013 (SIDES), and 0.000 (QUINT).

• Synthetic simulations confirm the picture. In Simulation 1 (features and treatment independent) and Simulation 2 (treatment dependent on X₁), the method was competitive even though several heuristic-driven causal methods already began to fail on these simple two-subgroup cases. In Simulation 3 (a five-feature rule-list simulator with three nested rules), the method clearly outperformed the five heuristics-driven baselines on both the absolute effect error |Â(Q̂) − Â(Q_gt)| and the Jaccard similarity between Q̂ and the ground-truth subgroup Q_gt. DistillTree was a close competitor on effect error in Simulation 3 but performed substantially worse on Jaccard similarity. Simulations varied sample sizes from 1000 to 5000 and were repeated 50 times, with error bars representing standard error.

• The two-step family inherits a harder problem. Methods such as Virtual Twins and DistillTree first estimate pointwise conditional treatment effects and then fit a model to those estimates. The authors argue pointwise estimation is substantially harder than subgroup-level estimation, making the learned subgroups highly sensitive to estimation error — consistent with DistillTree's weaker subgroup recovery.

• Model-based recursive partitioning targets a different object. Because such methods allow within-subgroup variation in effect, they do not target the maximum-effect subgroup as defined here, which Theorem 1 shows must be effect-homogeneous.

Methodology in Plain English

The authors begin by formalizing what they are looking for: a region Q of the feature space whose average treatment effect — the difference between the outcome probability under do(T = 1) and do(T = 0) within that region — is as large as possible. Under the standard no-hidden-confounder assumption, this quantity can be written using only observational conditional probabilities, so it is identifiable.

They then prove a structural fact: if an optimal region could be split into two parts, one of those parts would have to carry an effect at least as large as the whole. That means the optimum cannot contain internal effect variation — it must be effect-homogeneous.

Next they assume the world is generated by a partition-based model: the feature space is carved into disjoint pieces, and within each piece the outcome distribution is the same for every point. Under that assumption they prove the optimal subgroup is exactly one of those pieces. Finding it therefore means recovering the partition, which is what an ordinary decision tree or rule learner does when minimizing classification or regression loss.

The instantiation is deliberately plain: fit a CART classification tree with the Gini index (discrete outcomes) or a CART regression tree with MSE (continuous outcomes), prune by cost-complexity with cross-validation, then split the data into training and test halves. The tree is grown on the training half; each leaf defines a candidate subgroup via its root-to-leaf path (dropping any node that conditions on T); and the subgroup treatment effect is estimated on the held-out half using empirical conditional probabilities. The leaf with the highest honest estimate is returned as Q̂.

Why This Matters

Impact on research. The paper reframes maximum-effect subgroup discovery as a supervised learning problem rather than a causal estimation problem, and its theoretical result challenges the necessity of the large family of bespoke causal splitting criteria and regularizers proposed over the past decade. It also points out that the two-step strategy — estimate pointwise effects, then cluster or tree them — solves a harder problem than necessary. The result suggests that the growing literature of custom causal trees and rule sets should be compared against a plain CART baseline; notably, a heuristic-free method ranked first on 51.9% of the 77 semi-synthetic datasets and had a mean ground-truth effect of 10.540 versus at most 7.843 among baselines.

Real-world applications:

• Precision medicine: identifying which patient subgroup benefits most from a given treatment plan, so therapy can be targeted rather than applied uniformly.
• Public policy: finding the population segment that responds most strongly to an intervention such as an educational program or a social benefit.
• Education: locating the student subgroup with the largest gain from a specific instructional approach.
• Sequential covering: iteratively removing the discovered subgroup and re-running the method on remaining instances, using the maximum-effect objective as a building block for describing broad treatment-effect heterogeneity.

Industry relevance. Any organization running A/B tests or randomized trials — technology platforms, insurers, pharmaceutical companies, marketing teams — faces the question of which user segment responds most to an intervention. The paper's practical message is that an off-the-shelf, well-tuned decision tree with honest effect estimation can match or beat specialized causal tooling, lowering the methodological and engineering barrier to entry. The released source code makes adoption straightforward.

Future Directions

• Generalize beyond the partition-based assumption. Theorem 2 and its empirical validation rest on the data being generated by a finite partition with homogeneous conditional outcome distributions. The paper notes an extension of Theorem 2 to the hidden-confounder case in the supplementary materials, but broader model classes are left open.
• Test the reduction with more powerful partition learners. The authors explicitly state that they chose CART to validate the general approach rather than to propose a new algorithm; they point to advanced tree and rule-based methods and leave their evaluation to future work.
• Move beyond a single optimal subgroup. The paper frames the maximum-effect objective as a foundation for sequential covering or divide-and-conquer, so systematic procedures for iteratively peeling off subgroups are a natural extension.
• Revisit the design of causal heuristics. Given that the theory questions their necessity, an open question is whether any of the existing specialized criteria — QUINT's weighted sum, Interaction Tree's standardized child difference, Causal Tree's unbiased MSE proxy — offer advantages under conditions not covered by the simulations reported here. The experiments on semi-synthetic data, and the effect-estimation portion of the synthetic study, are described in the paper content only up to Table 1.

Target Audience

This paper is most valuable to causal inference and machine learning researchers working on heterogeneous treatment effects and subgroup discovery, particularly those building or benchmarking causal trees and rule sets. It also suits applied statisticians and data scientists in medicine, public policy, and education who need interpretable, high-effect subgroups from trial data and want to know whether specialized causal tooling is warranted. Readers need comfort with structural causal models and do-calculus to follow the proofs, but the practical takeaway — fit a standard tree, estimate effects honestly on held-out data — is accessible to anyone familiar with decision trees and randomized experiments.

Authors’ 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.