Disentangled Representation via Variational AutoEncoder for Continuous Treatment Effect Estimation

TL;DR

This work tackles continuous treatment effect estimation from observational data by introducing DRVAE, a variational autoencoder that disentangles covariates into four latent factors: instrumental Γ, confounding Δ, adjustment Υ, and external noise E. The model learns balanced representations to support counterfactual inference for continuous treatments, optimizing a loss that combines an ELBO with auxiliary and regularization terms, and infers counterfactual outcomes by sampling Δ and Υ. Empirical results on synthetic and semi-synthetic datasets show DRVAE outperforms nine state-of-the-art methods across multiple metrics and demonstrates robustness to non-causal noise while preserving the continuity of the ADRF curve. The approach advances causal representation learning by isolating true confounding signals from noise, enabling more accurate and stable estimation of dose–response relationships in complex, high-dimensional settings.

Abstract

Continuous treatment effect estimation holds significant practical importance across various decision-making and assessment domains, such as healthcare and the military. However, current methods for estimating dose-response curves hinge on balancing the entire representation by treating all covariates as confounding variables. Although various approaches disentangle covariates into different factors for treatment effect estimation, they are confined to binary treatment settings. Moreover, observational data are often tainted with non-causal noise information that is imperceptible to the human. Hence, in this paper, we propose a novel Dose-Response curve estimator via Variational AutoEncoder (DRVAE) disentangled covariates representation. Our model is dedicated to disentangling covariates into instrumental factors, confounding factors, adjustment factors, and external noise factors, thereby facilitating the estimation of treatment effects under continuous treatment settings by balancing the disentangled confounding factors. Extensive results on synthetic and semi-synthetic datasets demonstrate that our model outperforms the current state-of-the-art methods.

Figures (6)

• Figure 1: (a) The potential causal graph used by VCNet-liked methods (the varying coefficient structure proposed by VCNet called VCNet). All covariates are treated as confounder factors here. (b) The latent causal graph of this work. Where dashed arrows signify representative relationships, while solid arrows indicate causal relationships.
• Figure 2: The framework of DRVAE.
• Figure 3: A comparative analysis of the performance trends of the DRVAE against other benchmark methods on varying Simu($k$) datasets. The increase in external noise factors does not significantly impact the performance of our model.
• Figure 4: Estimated ADRF on the testing set from a typical run of DRVAE and other five baseline methods.
• Figure 5: Latent factors distribution.