# SubgroupTE: Advancing Treatment Effect Estimation with Subgroup Identification

arXiv: https://arxiv.org/abs/2401.12369 · HTML: https://mlsanity.com/paper/2401.12369/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 47 references that are in the map.
- First posted on arXiv 22 Jan 2024.

## Builds on (most cited first, 16 of 16)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Learning Counterfactual Representations for Estimating Individual Dose-Response Curves](https://mlsanity.com/paper/1902.00981/) (2019)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [VCNet and Functional Targeted Regularization For Learning Causal Effects of Continuous Treatments](https://mlsanity.com/paper/2103.07861/) (2021)
- [Topic Discovery via Latent Space Clustering of Pretrained Language Model Representations](https://mlsanity.com/paper/2202.04582/) (2022)
- [Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2202.01336/) (2022)
- [Causal Rule Sets for Identifying Subgroups with Enhanced Treatment Effect](https://mlsanity.com/paper/1710.05426/) (2017)
- [Interpretable Subgroup Discovery in Treatment Effect Estimation with Application to Opioid Prescribing Guidelines](https://mlsanity.com/paper/1905.03297/) (2019)
- [Robust Recursive Partitioning for Heterogeneous Treatment Effects with Uncertainty Quantification](https://mlsanity.com/paper/2006.07917/) (2020)
- [Unsupervised Machine Learning for the Discovery of Latent Disease Clusters and Patient Subgroups Using Electronic Health Records](https://mlsanity.com/paper/1905.10309/) (2019)
- [Identifying Heterogeneous Treatment Effects in Multiple Outcomes using Joint Confidence Intervals](https://mlsanity.com/paper/2212.01437/) (2022)
- [CAPITAL: Optimal Subgroup Identification via Constrained Policy Tree Search](https://mlsanity.com/paper/2110.05636/) (2021)
- [Tree-based Subgroup Discovery In Electronic Health Records: Heterogeneity of Treatment Effects for DTG-containing Therapies](https://mlsanity.com/paper/2208.14329/) (2022)

## Built on by (most cited first, 1 of 1)

- [Causal Clustering for Conditional Average Treatment Effects Estimation and Subgroup Discovery](https://mlsanity.com/paper/2509.05775/) (2025)

## Abstract (arXiv)

Precise estimation of treatment effects is crucial for evaluating intervention effectiveness. While deep learning models have exhibited promising performance in learning counterfactual representations for treatment effect estimation (TEE), a major limitation in most of these models is that they treat the entire population as a homogeneous group, overlooking the diversity of treatment effects across potential subgroups that have varying treatment effects. This limitation restricts the ability to precisely estimate treatment effects and provide subgroup-specific treatment recommendations. In this paper, we propose a novel treatment effect estimation model, named SubgroupTE, which incorporates subgroup identification in TEE. SubgroupTE identifies heterogeneous subgroups with different treatment responses and more precisely estimates treatment effects by considering subgroup-specific causal effects. In addition, SubgroupTE iteratively optimizes subgrouping and treatment effect estimation networks to enhance both estimation and subgroup identification. Comprehensive experiments on the synthetic and semi-synthetic datasets exhibit the outstanding performance of SubgroupTE compared with the state-of-the-art models on treatment effect estimation. Additionally, a real-world study demonstrates the capabilities of SubgroupTE in enhancing personalized treatment recommendations for patients with opioid use disorder (OUD) by advancing treatment effect estimation with subgroup identification.