Disentangled representation for sequential treatment effect estimation

Abstract

Background and objective: Treatment effect estimation, as a fundamental problem in causal inference, focuses on estimating the outcome difference between different treatments. However, in clinical observational data, some patient covariates (such as gender, age) not only affect the outcomes but also affect the treatment assignment. Such covariates, named as confounders, produce distribution discrepancies between different treatment groups, thereby introducing the selection bias for the estimation of treatment effects. The situation is even more complicated in longitudinal data, because the confounders are time-varying that are subject to patient history and meanwhile affect the future outcomes and treatment assignments. Existing methods mainly work on cross-sectional data obtained at a specific time point, but cannot process the time-varying confounders hidden in the longitudinal data.

Methods: In this study, we address this problem for the first time by disentangled representation learning, which considers the observational data as consisting of three components, including outcome-specific factors, treatment-specific factors, and time-varying confounders. Based on this, the proposed approach adopts a recurrent neural network-based framework to process sequential information and learn the disentangled representations of the components from longitudinal observational sequences, captures the posterior distributions of latent factors by multi-task learning strategy. Moreover, mutual information-based regularization is adopted to eliminate the time-varying confounders. In this way, the association between patient history and treatment assignment is removed and the estimation can be effectively conducted.

Results: We evaluate our model in a realistic set-up using a model of tumor growth. The proposed model achieves the best performance over benchmark models for both one-step ahead prediction (0.70% vs 0.74% for the-state-of-the-art model, when γ = 3. Measured by normalized root mean square error, the lower the better) and five-step ahead prediction (1.47% vs 1.83%) in most cases. By increasing the effect of confounders, our proposed model always shows superiority against the state-of-the-art model. In addition, we adopted T-SNE to visualize the disentangled representations and present the effectiveness of disentanglement explicitly and intuitively.

Conclusions: The experimental results indicate the powerful capacity of our model in learning disentangled representations from longitudinal observational data and dealing with the time-varying confounders, and demonstrate the surpassing performance achieved by our proposed model on dynamic treatment effect estimation.

Keywords: Causal inference; Deep learning; Representation learning; Sequential data; Treatment effect estimation.