Structure Maintained Representation Learning Neural Network for Causal Inference

TL;DR

This paper introduces a novel structure keeper in representation learning for causal inference, improving individual treatment effect estimation by maintaining correlations between covariates and representations, and demonstrating superior performance on real and simulated data.

Contribution

The paper proposes a structure keeper mechanism that enhances representation learning for causal inference, addressing the tradeoff between distribution balance and information retention.

Findings

Outperforms state-of-the-art methods on simulated data

Achieves better treatment effect estimation on MIMIC-III data

Effectively maintains correlation between covariates and representations

Abstract

Recent developments in causal inference have greatly shifted the interest from estimating the average treatment effect to the individual treatment effect. In this article, we improve the predictive accuracy of representation learning and adversarial networks in estimating individual treatment effects by introducing a structure keeper which maintains the correlation between the baseline covariates and their corresponding representations in the high dimensional space. We train a discriminator at the end of representation layers to trade off representation balance and information loss. We show that the proposed discriminator minimizes an upper bound of the treatment estimation error. We can address the tradeoff between distribution balance and information loss by considering the correlations between the learned representation space and the original covariate feature space. We conduct extensive experiments with simulated and real-world observational data to show that our proposed Structure Maintained Representation Learning (SMRL) algorithm outperforms state-of-the-art methods. We also demonstrate the algorithms on real electronic health record data from the MIMIC-III database.