[Paper Review] Estimating Counterfactual Treatment Outcomes over Time Through Adversarially Balanced Representations

Ioana Bica, Ahmed M. Alaa+2 more|arXiv (Cornell University)|Feb 10, 2020

Machine Learning in Healthcare60 references31 citations

TL;DR

The Counterfactual Recurrent Network (CRN) uses adversarially balanced representations to estimate time-varying counterfactual treatment outcomes from observational data, enabling optimal treatment decisions over time.

ABSTRACT

Identifying when to give treatments to patients and how to select among multiple treatments over time are important medical problems with a few existing solutions. In this paper, we introduce the Counterfactual Recurrent Network (CRN), a novel sequence-to-sequence model that leverages the increasingly available patient observational data to estimate treatment effects over time and answer such medical questions. To handle the bias from time-varying confounders, covariates affecting the treatment assignment policy in the observational data, CRN uses domain adversarial training to build balancing representations of the patient history. At each timestep, CRN constructs a treatment invariant representation which removes the association between patient history and treatment assignments and thus can be reliably used for making counterfactual predictions. On a simulated model of tumour growth, with varying degree of time-dependent confounding, we show how our model achieves lower error in estimating counterfactuals and in choosing the correct treatment and timing of treatment than current state-of-the-art methods.

Motivation & Objective

• Motivate and address the need to estimate treatment effects over time from observational data.
• Propose a sequence-to-sequence model that accounts for time-varying confounding via treatment-invariant representations.
• Enable counterfactual prediction of future outcomes under planned treatment sequences to inform treatment timing and selection.
• Leverage domain adversarial training to balance representations and improve causal inference in longitudinal data.

Proposed method

• Introduce the Counterfactual Recurrent Network (CRN), a seq2seq architecture that combines an encoder with a decoder to predict counterfactual outcomes under planned treatment sequences.
• Use an RNN encoder to build a treatment-invariant representation Phi(H_t) that balances the history with respect to current treatment, via a gradient reversal layer.
• Incorporate a treatment classifier G_a and an outcome predictor G_y, training to maximize the treatment-discrimination loss and minimize the outcome loss.
• Apply domain adversarial training across multiple treatment domains at each timestep to reduce time-dependent confounding bias.
• Decode counterfactual futures by initializing the decoder from the balanced representation and auto-regressively predicting outcomes under planned treatments.
• Provide uncertainty estimation via variational dropout and allow irregular sampling with PhasedLSTM variants if needed.
• The training objective combines the outcome prediction loss with a weighted treatment-discrimination loss, optimized via a minimax game capturing domain invariance.

Experimental results

Research questions

• RQ1 Can a balancing representation learned via adversarial training remove time-dependent confounding in longitudinal treatment settings?
• RQ2 Do counterfactual outcome estimates and treatment recommendations from CRN outperform state-of-the-art methods under varying degrees of confounding?
• RQ3 Can CRN provide reliable uncertainty estimates for counterfactual predictions and support clinical decision-making about start/stop and choice of treatments over time?

Key findings

• CRN achieves lower counterfactual prediction error than state-of-the-art methods in a tumor growth simulation across varying levels of time-dependent confounding.
• CRN with adversarial balancing (λ>0) significantly outperforms CRN without adversarial training, achieving up to 48.1% RMSE improvement in one-step ahead counterfactuals.
• Compared to Recurrent Marginal Structural Networks (RMSN), CRN provides about 17% lower error when confounding is strong (γ>6).
• CRN yields higher accuracy in selecting the correct treatment and timing across forecasting horizons, with robust performance across different policy settings.
• T-SNE visualizations indicate the balancing representation Phi(H_t) is not predictive of the treatment at time t, confirming treatment invariance.
• The approach demonstrates applicability to real data (MIMIC-III) beyond simulated tumor growth, supporting potential clinical decision support use.