G-Transformer: Counterfactual Outcome Prediction under Dynamic and Time-varying Treatment Regimes

TL;DR

G-Transformer introduces a novel Transformer-based model for accurate counterfactual outcome prediction in medical settings with dynamic, time-varying treatments, leveraging g-computation to simulate patient trajectories and outperform existing methods.

Contribution

This work presents the first Transformer architecture supporting g-computation for counterfactual prediction under complex, dynamic treatment regimes in longitudinal healthcare data.

Findings

G-Transformer outperforms classical models in simulated datasets.

G-Transformer achieves superior accuracy on real-world ICU data.

The model effectively captures long-range dependencies in time-varying covariates.

Abstract

In the context of medical decision making, counterfactual prediction enables clinicians to predict treatment outcomes of interest under alternative courses of therapeutic actions given observed patient history. In this work, we present G-Transformer for counterfactual outcome prediction under dynamic and time-varying treatment strategies. Our approach leverages a Transformer architecture to capture complex, long-range dependencies in time-varying covariates while enabling g-computation, a causal inference method for estimating the effects of dynamic treatment regimes. Specifically, we use a Transformer-based encoder architecture to estimate the conditional distribution of relevant covariates given covariate and treatment history at each time point, then produces Monte Carlo estimates of counterfactual outcomes by simulating forward patient trajectories under treatment strategies of interest. We evaluate G-Transformer extensively using two simulated longitudinal datasets from mechanistic models, and a real-world sepsis ICU dataset from MIMIC-IV. G-Transformer outperforms both classical and state-of-the-art counterfactual prediction models in these settings. To the best of our knowledge, this is the first Transformer-based architecture that supports g-computation for counterfactual outcome prediction under dynamic and time-varying treatment strategies.