SubgroupTE: Advancing Treatment Effect Estimation with Subgroup Identification

Abstract

Precise estimation of treatment effects is crucial for accurately evaluating the intervention. 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 often overlook the diversity of treatment effects across potential subgroups that have varying treatment effects and characteristics, treating the entire population as a homogeneous group. This limitation restricts the ability to precisely estimate treatment effects and provide targeted 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 responses and more precisely estimates treatment effects by considering subgroup-specific treatment effects in the estimation process. In addition, we introduce an expectation-maximization (EM)-based training process that iteratively optimizes estimation and subgrouping networks to improve both estimation and subgroup identification. Comprehensive experiments on the synthetic and semi-synthetic datasets demonstrate the outstanding performance of SubgroupTE compared to the existing works for treatment effect estimation and subgrouping models. Additionally, a real-world study demonstrates the capabilities of SubgroupTE in enhancing targeted treatment recommendations for patients with opioid use disorder (OUD) by incorporating subgroup identification with treatment effect estimation.

Keywords: Deep learning; Opioid use disorder; Subgroup analysis; Treatment effect estimation.

Fig. 1.

An overview of SubgroupTE. The proposed framework consists of three modules: 1) feature representation network, 2) subgrouping network, and 3) subgroup-informed prediction network. First, the feature representation network transforms the raw data into a meaningful representation for estimating treatment effects. Second, the subgrouping network assigns a subgroup probability vector to each data sample, aiming to maximize the homogeneity of estimated treatment effects within subgroups. Finally, the subgroup-informed prediction network integrates both subgroup information and patient representation to estimate treatment effects by subgroup.

Fig. 2.

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.

Fig. 3.

Illustration of the trends in PEHE and variance within and across subgroups during the training phase.

Fig. 4.

Sensitivity analysis conducted for (a) Coefficient and (b) Number of subgroups on the semi-synthetic dataset.

Fig. 5.

Illustration of the cohort selection criteria. The index date indicates the first prescription date of the drug. The baseline and follow-up periods encompass all dates before and after the index date, respectively.

Fig. 6.

The boxplots of the treatment effect distribution for the identified subgroups on the opioid dataset.

Fig. 7.

The heatmap of the relative ratios of variables for demographics and diagnosis codes among the three subgroups. These relative ratios are calculated using the formula ${\pi }_{k,i}/{\sum }_{k=1}^K{\pi }_{k,i}$, where ${\pi }_{k,i}$ represents the ratio of the $i$-th variable within the $k$-th subgroup. SMD: Substance-related mental disorders; NSD: Other nervous system disorders; SDB: Spondylosis, intervertebral disc disorders, or other back problems; CTB: Other connective tissue disease.