On the Effects of Irrelevant Variables in Treatment Effect Estimation with Deep Disentanglement

TL;DR

The paper tackles treatment effect estimation from observational data by addressing the role of irrelevant pre-treatment variables. It introduces DRI-ITE, a four-factor disentanglement framework that explicitly models an irrelevant latent space $\\Omega$ via an autoencoder-based reconstruction, coupled with orthogonal regularization to prevent leakage into other latent factors $\\\Gamma,\\\Delta,\\\Upsilon$. The approach integrates four encoders, regression/classification heads, and a reconstruction objective to robustly identify latent factors and improve ITE predictions, demonstrated on synthetic data and real benchmarks IHDP and Jobs, with consistent gains in PEHE and policy risk as $\\\Omega$ grows. The results suggest that explicit handling of irrelevant variables yields more reliable counterfactual estimates in high-dimensional settings, offering practical benefits for observational causal inference frameworks.

Abstract

Estimating treatment effects from observational data is paramount in healthcare, education, and economics, but current deep disentanglement-based methods to address selection bias are insufficiently handling irrelevant variables. We demonstrate in experiments that this leads to prediction errors. We disentangle pre-treatment variables with a deep embedding method and explicitly identify and represent irrelevant variables, additionally to instrumental, confounding and adjustment latent factors. To this end, we introduce a reconstruction objective and create an embedding space for irrelevant variables using an attached autoencoder. Instead of relying on serendipitous suppression of irrelevant variables as in previous deep disentanglement approaches, we explicitly force irrelevant variables into this embedding space and employ orthogonalization to prevent irrelevant information from leaking into the latent space representations of the other factors. Our experiments with synthetic and real-world benchmark datasets show that we can better identify irrelevant variables and more precisely predict treatment effects than previous methods, while prediction quality degrades less when additional irrelevant variables are introduced.

Figures (7)

• Figure 1: Average PEHE error on IHDP dataset against number of irrelevant variable dimensions (smaller the better). PEHE generally degrades with more irrelevant factors but our method is less affected.
• Figure 2: (Top): Illustrates the rise in PEHE as the number of irrelevant variables grows based on a baseline approach (TVAE). (Left): visualizes the individual contributions of variables towards learning the encoder for miscellaneous factors. Notably, the contribution of irrelevant variables mirrors that of relevant ones, underscoring the limitations of TVAE in disentangling irrelevant variables. (Right): shows average contribution of each variable in PEHE increase by using permutation of variables. Irrelevant variables are significantly participating in PEHE increase.
• Figure 3: Underlying factors of $\mathcal{X}$. Observe that $\Omega$ has no associated downstream task with any observed variable.
• Figure 4: High level architecture of DRI-ITE.
• Figure 5: The visualization of feature contributions on each latent factor representational network is conducted for the dataset with dimensions 8, 8, 8, 15 ($\Gamma$,$\Delta$,$\Upsilon$,$\Omega$) utilizing the $\bar{W}$ criterion based on DRI-ITE (ours). The top row visualizes all individual features, where high values are expected for the features between dotted lines, the bottom row represents the average over all features that are supposed to be represented by that particular network compared to the average weight of wrongly represented features.