Learning Subgroups with Maximum Treatment Effects without Causal Heuristics

TL;DR

This paper introduces a novel method for discovering subgroups with maximum treatment effects by directly modeling the problem within the structural causal model framework, avoiding heuristic-based approaches.

Contribution

It reformulates subgroup discovery as a supervised learning problem under the SCM framework and demonstrates how to effectively learn optimal subgroups using standard tree-based methods like CART.

Findings

Outperforms baseline methods in accurately identifying subgroups with maximum treatment effects.

Avoids ad-hoc causal heuristics, leading to more reliable subgroup discovery.

Validated on synthetic and semi-synthetic datasets.

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.