2024/01/22 by Seungyeon Lee, Lee, Seungyeon, Ruoqi Liu +7 · 1 citation
Computer Science · Mathematics · #Advanced Causal Inference Techniques #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning in Healthcare #Methodology (stat.ME)
paper · pdf · doi:10.48550/arxiv.2401.12369
openalex publication_date 2024/01/22 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
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.