SubgroupTE: Advancing Treatment Effect Estimation with Subgroup Identification

TL;DR

SubgroupTE addresses heterogeneity in treatment effects by jointly identifying subgroups and estimating subgroup-specific causal effects within a Transformer-based representation. It introduces an EM-based training framework that alternates between updating subgroup centroids and optimizing three networks—feature representation, subgrouping, and subgroup-informed prediction—using KDE-based centroid shifts to maintain stable clustering. On synthetic and semi-synthetic data, SubgroupTE achieves state-of-the-art performance with substantial improvements in PEHE and \\epsilon_{ATE}, while also yielding interpretable subgroups with clearly separated TE distributions. A real-world opioid use disorder study demonstrates the method's practical value by revealing subgroup-specific treatment effects and guiding personalized treatment recommendations.

Abstract

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.

Figures (8)

• Figure 1: An overview of SubgroupTE.
• Figure 2: Training process for SubgroupTE
• Figure 3: Sensitivity analysis conducted for (a) Coefficient and (b) Number of subgroups on the semi-synthetic dataset. For (a), the performance of each coefficient is evaluated while fixing the remaining two coefficients at 1.
• Figure 4: The boxplots of the treatment effect distribution for the identified subgroups on the (a) synthetic and (b) semi-synthetic datasets. The box spans from the first quartile to the third quartile of the data, with a line indicating the median. The whiskers extend from the box to encompass the 5th to 95th percentiles.
• Figure 5: Illustration of the trends in PEHE and variance within and across subgroups during the training phase on the validation set of the synthetic dataset.