Source-linked AI summary
Estimating Counterfactual Treatment Outcomes over Time Through Adversarially Balanced Representations
Ioana Bica, Ahmed M. Alaa, James Jordon, Mihaela van der Schaar
TL;DR
The paper addresses how to estimate and use treatment effects over time from observational data despite time-varying confounding. It introduces CRN, which learns treatment-invariant patient-history representations with domain adversarial training and predicts outcomes for future treatment sequences. In tumour-growth experiments and a real electronic-health-record dataset, CRN improved over existing methods in the reported medical decision-making scenarios.
Problem
Time-varying confounders bias observational estimates of treatment effects, while treatment timing and selection over time remain important medical questions.
Method
CRN is a sequence-to-sequence model that uses domain adversarial training to build treatment-invariant representations and estimate counterfactual outcomes for future treatment sequences.
Results
CRN improved over existing state-of-the-art methods in counterfactual prediction and in choosing treatment and treatment timing in tumour-growth experiments.
Takeaways & Limitations
CRN's counterfactual predictions can support clinical decision systems for selecting treatments, identifying treatment timing, and deciding when treatment is no longer needed.
Takeaways & Limitations
The approach relies on consistency, positivity, and no hidden confounders, and further work is needed for treatment combinations and individualized dosage effects.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Estimating treatment effects from observational data is difficult because time-dependent confounders link past treatments, evolving patient covariates, future treatment assignments, and outcomes. CRN addresses this challenge with adversarially balanced representations and sequence-to-sequence counterfactual prediction, improving treatment and timing decisions in tumour-growth experiments.
- Clinical trials provide strong causal evidence but are expensive, small, and restricted by narrow inclusion criteria.
- Time-dependent confounders are patient covariates affected by past treatments that influence subsequent treatments and outcomes.
- Static causal-inference methods cannot model evolving covariates or estimate effects for sequential treatment plans.
- CRN is a sequence-to-sequence architecture that uses domain adversarial training to construct treatment-invariant representations of patient history.
- CRN estimates counterfactual trajectories for future treatment plans, supporting treatment selection, treatment timing, and treatment stopping decisions.
- In a tumour-growth model, CRN outperformed current state-of-the-art methods in counterfactual prediction and in choosing treatment and treatment timing.
2 RELATED WORK
Prior work includes epidemiological methods and machine-learning approaches for causal inference, but longitudinal treatment effects remain difficult to estimate. CRN differs from recurrent weighting methods by using balancing representations rather than inverse-probability weights.
- Longitudinal causal-inference methods include g-computation, Structural Nested Models, and Marginal Structural Models.
- Earlier epidemiological approaches often use logistic or linear regression predictors that are unsuitable for complex time dependencies.
- Several methods model patient outcomes or treatment responses with Gaussian processes, dynamical systems, or propensity-weight estimation.
- RMSNs use recurrent networks to estimate inverse-probability treatment weights and forecast treatment responses, but retain the high-variance weighting problem.
- Static balancing-representation methods do not directly generalize to longitudinal settings with sequential treatment assignment and changing covariates.
3 PROBLEM FORMULATION
The paper formulates longitudinal treatment-effect estimation using patient histories, treatment sequences, and potential outcomes. Identification relies on consistency, positivity, and sequential strong ignorability.
- The potential-outcomes framework defines factual or counterfactual outcomes for every possible course of treatment.
- Patient history combines covariate history, previous treatment assignments, and baseline features.
- The target is the outcome under a possible future treatment sequence beginning at the current timestep.
- Identification assumes consistency, positivity, and no hidden confounders through sequential strong ignorability.
4 COUNTERFACTUAL RECURRENT NETWORK
CRN combines adversarial representation learning with a recurrent encoder-decoder to estimate outcomes under intended future treatment sequences. The encoder balances treatment information against outcome prediction, while the decoder generates counterfactual trajectories autoregressively.
- A supervised model trained on observational data is not reliable for causal prediction without correcting time-dependent confounding.
- CRN maps patient history to a treatment-invariant representation whose distribution is equalized across treatment groups.
- Encoder: The encoder processes treatment history, covariates, and baseline features while suppressing treatment predictability and preserving one-step outcome information.
- Decoder: The decoder initializes from the balanced representation and predicts outcomes for an intended sequence of future treatments.
- Decoder: During testing, predicted outcomes are fed back autoregressively, allowing comparisons of treatment starts, stops, timings, and choices.
- Variational dropout enables uncertainty intervals, and PhasedLSTM can adapt the model to irregularly sampled observations.
5 ADVERSARIALLY BALANCED REPRESENTATION OVER TIME
CRN builds treatment-invariant representations of patient history by adversarially balancing treatment domains while preserving outcome-prediction accuracy. Its minimax objective reaches equilibrium when the induced representations are invariant across treatments, targeting removal of time-dependent confounding bias.
- Treatment-invariant representations: At each timestep, the possible treatments define domains whose patient-history representations are required to have identical distributions.The target condition is P(Φ(H̄_t) | A_1) = ··· = P(Φ(H̄_t) | A_K).
- Adversarial training: Domain adversarial training makes the representation invariant to treatment while retaining low error for predicting the next outcome.The treatment classifier and outcome predictor impose competing objectives on the learned representation.
- Adversarial training: The treatment classifier is trained to discriminate domains, while the representation is trained to maximize treatment-classification loss and minimize outcome loss.This trade-off produces a representation that is difficult to use for treatment discrimination but useful for outcome prediction.
- Optimization: The overall optimization seeks a saddle point that balances domain discrimination with outcome estimation.The representation and outcome parameters minimize the objective, while the treatment-classifier parameters maximize it.
- Theoretical guarantee: Theorem 1 states that the minimax game reaches a global minimum exactly when learned representations are invariant across all treatments.The theorem connects the optimization objective to equality of the treatment-conditional representation distributions.
- Theoretical guarantee: A good representation supports low counterfactual-estimation error across treatments while minimizing the H-divergence between induced domain distributions.The method combines distributional alignment with empirical training performance.
6 EXPERIMENTS
The experiments evaluate CRN on a simulated tumour-growth model with time-varying confounding, comparing it with established benchmarks for counterfactual prediction and treatment recommendation. CRN achieves strong performance across counterfactual estimation, treatment selection, and treatment timing, with additional applicability shown on real ICU data.
- Experimental setup: The evaluation uses a pharmacokinetic-pharmacodynamic tumour-growth model with simulated chemotherapy and radiotherapy responses and controllable time-varying confounding.The confounding strength is controlled through γc and γr, which govern how strongly treatment assignment depends on patient history.
- Experimental setup: The benchmark suite includes MSMs, RMSNs, standard supervised RNN and linear models, and CRN without adversarial training.The ablated model, CRN (λ = 0), has the same architecture without adversarial training.
- Counterfactual prediction: 48.1% improvement: at γ = 10, CRN improves over the same architecture without domain adversarial training for one-step counterfactual estimation.The metric is normalized RMSE, with tumour volume normalized by Vmax = 1150cm3.
- Counterfactual prediction: CRN achieves the lowest RMSE across all confounding levels and improves by ∼17% over RMSNs when γ > 6.For five-step prediction, the results highlight limitations of IPTW-based handling because τ treatment weights can produce high variance.
- Treatment recommendation: CRN achieves the lowest RMSE and highest treatment and timing accuracy across the evaluated treatment policies in the recommendation task.The evaluation uses forecasting horizons and test sets designed to assess selecting both the treatment and when to administer it.
- Real-data applicability: Additional experiments on MIMIC III provide results concerning CRN’s applicability to real ICU electronic-health-record data.These experiments are reported in Appendix M and address more complex medical scenarios involving real data.
7 CONCLUSION
The paper addresses causal inference for time-dependent treatments with CRN, which builds treatment-invariant representations to reduce confounding bias. Experiments show improvements over existing methods, while the broader problem remains less studied and requires further methodological development.
- Time-dependent treatment effects remain less studied than static causal inference, motivating further methods and theory for individualized treatment scenarios.
- The paper identifies future needs including alternative confounding methods, treatment combinations over time, dosage effects, and stronger theoretical error guarantees.
- CRN estimates individualized treatment effects over time by handling time-dependent confounding through adversarially trained representations.
- CRN learns counterfactual outcomes under all treatment options, a harder task than selecting only the optimal treatment but one that can support policy selection.
- Treatment-invariant representations equalize representation distributions across possible treatments, breaking the association between patient history and current treatment assignment.
D PROOF OF THEOREM 1
The proof analyzes the adversarial training objective for a fixed representation and derives the discriminator’s optimal treatment-prediction probabilities. Substituting this optimum transforms the minimax objective into a divergence-based representation-learning objective.
- For a fixed representation Φ, Proposition 1 characterizes the optimal prediction probabilities of the treatment discriminator.
- The proof obtains the discriminator optimum by maximizing the value function pointwise under a probability constraint using Lagrange multipliers.
- Theorem 1 follows by substituting Proposition 1’s expression into the minimax game defined in Equation 6.
- The resulting objective involves the multi-distribution Jensen-Shannon divergence, which is nonnegative and zero exactly when the treatment-conditional representation distributions are equal.
E TRAINING PROCEDURE FOR CRN
CRN training separates encoder optimization from decoder optimization. The encoder learns treatment-invariant history representations and one-step predictions, then initializes a decoder trained on shorter trajectory sequences.
- The encoder is trained first to build treatment-invariant patient-history representations and perform one-step-ahead prediction.
- After encoder optimization, its balancing representation is computed at each timestep and used to initialize the decoder hidden state.
- Patient trajectories are split into shorter sequences of τmax timesteps, producing T(i) − τmax decoder-training examples per patient.
- Both encoder and decoder training use Adam and run for 100 epochs in the reported implementation.
- The decoder is optimized using its own parameters and learning rate within the training procedure.
F PHARMACOKINETIC-PHARMACODYNAMIC MODEL OF TUMOUR GROWTH
The evaluation uses a tumour-growth PK-PD simulator with chemotherapy and radiotherapy whose assignment depends on tumour history, then compares CRN with MSM, RMSN, and baseline models. CRN’s advantage over RMSN increases for longer-horizon counterfactual prediction.
- Pharmacokinetic-pharmacodynamic model: The PK-PD simulator models non-small cell lung cancer and generates outcomes under all possible treatment options for counterfactual evaluation.
- Time-varying confounding: Chemotherapy and radiotherapy assignments are Bernoulli variables whose probabilities depend on recent average tumour diameter.
- Time-varying confounding: Higher γ increases the influence of tumour-diameter history on treatment assignment; γ = 0 corresponds to random treatment assignment.
- Time-varying confounding: For D̄(t) = 3Dmax/4, chemotherapy probability is pc(t) = σ(0.25γc), giving 0.56 at γ = 1, 0.77 at γ = 5, and 0.92 at γ = 10.
- Baselines: MSMs remove confounding with stabilized inverse-probability weights, while RMSNs use recurrent networks to estimate propensity scores and outcomes.
- Results: CRN’s performance gains over RMSN increase as the number of future timesteps in counterfactual prediction grows.
K.2 DETAILED RESULTS FOR THE COUNTERFACTUAL PREDICTIONS
Tables 8 and 9 report detailed counterfactual-prediction results, evaluated across one-step and τ-step horizons under varying time-dependent confounding. The experiment also assesses treatment and treatment-timing selection from simulated tumour-growth counterfactuals.
- Counterfactual prediction results: Tables 8 and 9 report normalized RMSE for one-step and τ-step-ahead counterfactual predictions under varying degrees of time-dependent confounding.The parameter γ measures the degree of time-dependent confounding.
- Evaluation setup: The evaluation simulates 1000 test samples from a tumour-growth model with four treatment options at each future timestep.The options are no treatment, chemotherapy, radiotherapy, and combined chemotherapy and radiotherapy.
- Evaluation setup: For each patient history and future horizon, the procedure generates counterfactual outcomes for treatment applications at each possible future time.A patient with 50 timesteps and τ = 3 yields 2 · 3 · 50 = 300 counterfactuals.
- Treatment and timing selection: Treatment selection chooses the option with the lowest predicted tumour volume, then selects the application time that produces the lowest outcome.Model selections are compared with the corresponding choices from the true generated counterfactual data.
M RESULTS ON FACTUAL PREDICTION ON MIMIC III
On MIMIC III, CRN is evaluated for factual white-blood-cell-count prediction after longitudinal antibiotic treatment. Because counterfactual outcomes are unavailable, the study reports factual RMSE and requires domain-expert validation for counterfactual alternatives.
- Dataset: The MIMIC III cohort contains 3487 ICU patients with antibiotic trajectories of up to 30 timesteps and 25 longitudinal patient covariates.The covariates include laboratory tests and vital signs measured over time.
- Dataset: Antibiotic administration is represented as a binary treatment at each timestep, with longitudinal covariates aggregated daily because treatment decisions are made daily.The dataset is split into 2826/313/348 patients for training, validation, and testing.
- Prediction target: The task estimates individualized antibiotic effects on white blood cell count, a marker associated with severe illness and poor ICU outcomes.Antibiotic treatment aims to reduce white blood cell count, with effectiveness depending on when treatment is administered relative to patient history.
- Results: Table 10 reports RMSE for τ-step-ahead factual white blood cell count predictions across multiple prediction horizons.Counterfactual data are unavailable for this dataset, so evaluation uses factual predictions.
- Results: CRN achieves better factual-outcome performance than RMSN on this real-world electronic-health-record dataset.Domain-expert knowledge is required to validate counterfactual predictions under alternative antibiotic treatments, and that validation is outside the paper’s scope.