Learning Subgroups with Maximum Treatment Effects without Causal Heuristics

2025-11-25

AI Summary

This paper addresses the core problem of identifying subgroups with the largest average treatment effect (ATE) in precision medicine, public policy, and education. We propose a novel subgroup discovery method that does not rely on causal heuristics or point estimators. Grounded in structural causal models, we rigorously formulate optimal subgroup identification as a supervised learning task: under a partitioned structural model assumption, we prove its equivalence to recovering the underlying data-generating mechanism. Our approach adopts a partition-based supervised learning framework—exemplified by CART—to jointly learn both subgroup structure and treatment effect prediction. Extensive experiments on synthetic and semi-synthetic datasets demonstrate that our method significantly outperforms state-of-the-art baselines in both accuracy and robustness of effect estimation. To our knowledge, this is the first end-to-end subgroup learning method that achieves high performance without invoking causal assumptions or heuristic approximations.

Technology Category

• Machine Learning: Causal Learning
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Causality

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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.

Problem

• Discovering subgroups with maximum treatment effects for targeted decision making
• Addressing limitations of heuristic-based causal subgroup discovery methods
• Developing partition-based models to identify optimal subgroups without causal heuristics

Innovation

• Uses structural causal model framework for subgroup discovery
• Reduces subgroup discovery to standard supervised learning
• Applies CART method to identify maximum treatment effect subgroups