Source-linked AI summary
Causal Transformer for Estimating Counterfactual Outcomes
Valentyn Melnychuk, Dennis Frauen, Stefan Feuerriegel
TL;DR
Estimating individualized counterfactual outcomes over time from observational data matters for personalized medicine, but traditional approaches can be biased and struggle with complex, long-range dependencies. The paper proposes a Causal Transformer with tailored subnetworks and counterfactual domain confusion training, achieving state-of-the-art performance across synthetic and real-world experiments.
Problem
Estimating counterfactual outcomes over time from observational data is important for personalized medicine, but traditional estimation can be biased and existing approaches struggle with complex, long-range dependencies.
Method
The Causal Transformer combines a tailored transformer-based architecture with a counterfactual domain confusion loss for end-to-end training.
Results
Across extensive synthetic and real-world experiments, the Causal Transformer achieves state-of-the-art performance for one- and multi-step-ahead predictions.
Takeaways & Limitations
The proposed method provides a transformer-based approach for estimating counterfactual outcomes over time while capturing complex, long-range patient trajectories.
Takeaways & Limitations
The counterfactual outcomes are identifiable from factual observational data only under consistency, sequential ignorability, and sequential overlap assumptions.
Abstract
from arXiv · showhide
Estimating counterfactual outcomes over time from observational data is relevant for many applications (e.g., personalized medicine). Yet, state-of-the-art methods build upon simple long short-term memory (LSTM) networks, thus rendering inferences for complex, long-range dependencies challenging. In this paper, we develop a novel Causal Transformer for estimating counterfactual outcomes over time. Our model is specifically designed to capture complex, long-range dependencies among time-varying confounders. For this, we combine three transformer subnetworks with separate inputs for time-varying covariates, previous treatments, and previous outcomes into a joint network with in-between cross-attentions. We further develop a custom, end-to-end training procedure for our Causal Transformer. Specifically, we propose a novel counterfactual domain confusion loss to address confounding bias: it aims to learn adversarial balanced representations, so that they are predictive of the next outcome but non-predictive of the current treatment assignment. We evaluate our Causal Transformer based on synthetic and real-world datasets, where it achieves superior performance over current baselines. To the best of our knowledge, this is the first work proposing transformer-based architecture for estimating counterfactual outcomes from longitudinal data.
1. Introduction
The paper targets counterfactual outcome estimation over time from observational data, where long-range dependencies and time-varying confounding challenge existing LSTM-based methods. It introduces a Causal Transformer combining three transformer subnetworks with a CDC-based training procedure and reports state-of-the-art performance.
- Observational data are increasingly used because randomized controlled trials are costly, impractical, or unethical.
- Existing longitudinal methods use simple LSTM networks, limiting their ability to model complex, long-range dependencies in medical data.
- The paper proposes a Causal Transformer for one- and multi-step-ahead counterfactual outcome predictions.
- The architecture combines separate transformer subnetworks for time-varying covariates, past treatments, and past outcomes through a joint network with cross-attentions.
- The CDC loss learns representations predictive of outcomes but non-predictive of current treatment assignment, addressing confounding bias.
- Experiments on synthetic and real-world data report state-of-the-art performance, while ablations indicate that both the architecture and CDC loss are needed.
2. Related Work
Related work covers causal outcome estimation in static and longitudinal settings, transformer modeling for sequential data, and the gap between these areas. Existing longitudinal methods address time-varying confounding, while transformers had not been developed specifically for causal inference.
- Estimating counterfactual outcomes in static setting: Static counterfactual estimation uses cross-sectional data and includes several deep-learning methods.
- Estimating counterfactual outcomes over time: Longitudinal counterfactual methods include epidemiological models, Bayesian non-parametric approaches, and recurrent neural networks.
- Estimating counterfactual outcomes over time: RMSNs, CRN, and G-Net address time-varying confounding through different mechanisms within the potential outcomes framework.
- Transformers: Transformers use self-attention and are powerful for modeling long-range associative dependencies in sequential data.
- Transformers: The paper identifies a novelty gap: no prior work had developed transformers specifically for causal inference.
3. Problem Formulation
The task is to estimate future counterfactual outcomes for a patient given their history and a specified treatment intervention. Identification relies on consistency, sequential ignorability, and sequential overlap, while ordinary machine learning can be biased because treatments affect future covariates.
- Patient trajectories contain time-varying covariates, categorical treatments, outcomes, and static covariates.
- The potential outcome Yt+τ[at:t+τ−1] represents the outcome under a specified non-random treatment intervention over the projection horizon.
- Potential counterfactual outcomes are identifiable from observational data under consistency, sequential ignorability, and sequential overlap.
- The goal is to estimate future counterfactual outcomes after applying a treatment intervention for a given patient history.
- Directly estimating the target function with traditional machine learning can be biased because treatments influence both outcomes and future covariates.
4. Causal Transformer
The Causal Transformer combines three sequence-specific transformer subnetworks with cross-attention and trains balanced representations using a counterfactual domain confusion loss.
- Model architecture: The CT processes past covariates, outcomes, and treatments through three parallel transformer subnetworks, while incorporating future treatment assignments and autoregressive outcome predictions.Static covariates are supplied to all subnetworks.
- Model architecture: Each transformer block combines multi-head self- and cross-attention, feed-forward processing, layer normalization, and residual connections.Cross-attention exchanges information between the parallel subnetworks, while masking preserves causal information flow.
- Balanced representations: The CT constructs balanced representations by pooling parallel hidden states and applying a fully connected layer with an ELU nonlinearity.The resulting representations are intended to be treatment-invariant.
- Positional encoding: Relative positional encodings represent distances between positions rather than absolute locations and restrict attention to past information or the current position.This supports ordered clinical trajectories while avoiding dependence on absolute time steps.
- Training objective: The CDC objective jointly makes representations predictive of the next outcome and non-predictive of the current treatment assignment to address confounding bias.The method differs from gradient reversal by minimizing a reversed KL-divergence between treatment-conditional representation distributions and their mixture.
- Training procedure: Mini-batch masking enables one model to perform both one- and multi-step-ahead prediction through end-to-end training.The paper reports that the CDC loss is effective and theoretically leads to balanced representations, while the combined architecture and loss are crucial to performance.
5. Experiments
Experiments on synthetic, semi-synthetic, and real-world data evaluate the Causal Transformer against established baselines and its architectural components. CT consistently performs better, with the CDC loss and three-subnetwork design jointly contributing to performance.
- Experimental setup: Synthetic and semi-synthetic experiments compare CT with established baselines under controlled counterfactual settings.Synthetic data permit evaluation against true counterfactual outcomes, while semi-synthetic data provide high-dimensional patient trajectories.
- Fully-synthetic data: CT shows a notable performance gain over state-of-the-art baselines, especially with stronger confounding and longer prediction horizons.The reported overall comparison describes CT as superior by a large margin.
- Component and runtime analysis: Removing the CDC loss reduces performance, particularly for multi-step-ahead prediction, while CT also provides a significant runtime speedup over other neural methods.The speedup is attributed mainly to self- and cross-attention processing and single-stage end-to-end training.
- Semi-synthetic data: On semi-synthetic data, CT improves over baselines by 38.5% on average across projection horizons.Comparisons with CT (α = 0) and EDCT show benefits from both the CDC loss and the three-subnetwork architecture.
- Real-world data: On real MIMIC-III data, CT outperforms all state-of-the-art baselines for factual-outcome prediction.True counterfactual outcomes are unavailable in this setting, so the evaluation reports factual prediction performance.
- Ablation study: Ablations show that combining the three-subnetwork architecture with the CDC loss is crucial, particularly at the long horizon τ = 6.Replacing CDC with gradient reversal produces unstable or worse results, and the CDC loss also improves CRN RMSE.
6. Conclusion
The paper targets counterfactual outcome estimation over time for personalized medicine and relies on potential-outcomes identification assumptions. It concludes that the Causal Transformer achieves state-of-the-art performance, while empirical G-computation is avoided because conditional-distribution estimation can be difficult in high-dimensional, finite data.
- Conclusion: Counterfactual patient-trajectory estimates are needed to support personalized medicine.The paper frames these estimates as relevant to individualized treatment outcomes.
- Conclusion: The Causal Transformer combines a custom three-subnetwork architecture with a counterfactual domain confusion loss for end-to-end training.The architecture is designed for complex, long-range patient trajectories.
- Assumptions for causal identification: The identification framework uses consistency, sequential overlap, and sequential ignorability assumptions for time-varying treatments and outcomes.Sequential ignorability excludes unobserved confounders affecting both treatment and outcome.
- Assumptions for causal identification: The data-generating mechanism is represented by the causal diagram in Figure 3, where unobserved noise affects covariates and outcomes but not treatments.Observed history includes all time-varying confounders up to time t.
- Methodological boundary: Empirical G-computation is not used because estimating conditional distributions of time-varying covariates can be challenging with finite, high-dimensional datasets.The paper identifies this as a practical reason to refrain from explicit G-computation.
B. Methods for Estimating Counterfactual Outcomes over Time
Methods for longitudinal counterfactual estimation span linear, recurrent, and transformer-based approaches under differing assumptions and data settings. The paper’s Causal Transformer uses intertwined attention subnetworks and clipped relative positional encodings to model time-dependent inputs and generalize across sequence lengths.
- Existing methods: The benchmark includes MSMs, RMSNs, CRN, and G-Net, selected as representative methods for longitudinal counterfactual estimation.The baselines use assumptions aligned with the paper’s setting where applicable.
- Existing methods: Some alternatives are excluded because they support only single-time binary treatments or binary-treatment sequences, or require stronger ignorability assumptions.These setting mismatches limit direct applicability to the paper’s problem.
- Problem setting: The paper distinguishes longitudinal counterfactual estimation from reinforcement learning because its data-generating mechanism is non-Markovian.The relevant history state typically grows with time, increasing the state-space size.
- Causal Transformer architecture: The multi-input transformer block uses three intertwined subnetworks with separate self-attentions, cross-attentions, pooling, and parallel feed-forward layers.Treatment representations are left-shifted relative to covariate and outcome representations.
- Positional encoding: Absolute positional encodings identify time steps with fixed sine and cosine functions, whereas trainable absolute encodings require T × d_h parameters.The paper uses clipped relative positional encodings because trainable absolute encodings limit generalization to unseen sequence lengths.
E. Details on Adversarial Training
The adversarial training procedure optimizes the Causal Transformer iteratively rather than globally and uses exponential moving averages for stabilization.
- Adversarial training: Adversarial training updates the Causal Transformer through iterative gradient descent rather than global optimization.The procedure includes exponential moving average stabilization.
- Adversarial training: Algorithm 1 specifies the adversarial training procedure for the Causal Transformer.Its inputs include the iteration count, smoothing parameter, CDC coefficient, and learning rate.
F. Proof of Theorem F.2
The proof derives the theorem by optimizing the adversarial objective for fixed representations and then showing that global optimality is equivalent to treatment-balanced representations.
- Lemma F.1 provides the optimal prediction probabilities for a fixed representation network Φ.The proof then substitutes these probabilities into the objective from Eq. (45).
- The proof uses Lagrange multipliers to obtain the pointwise optimum of the prediction objective.
- Theorem F.2 defines distributions over histories and representations conditional on each treatment before characterizing the optimal pair (Φ*, G*A).
- At the optimum, Φ* satisfies Eq. (45) if and only if its representation distributions are balanced across treatments.
- The contradiction argument rules out an optimal representation that remains unbalanced, establishing the theorem’s global-optimum characterization.
G. Baseline Methods
The baselines cover weighting, recurrent representation learning, adversarial balancing, and G-computation approaches for longitudinal counterfactual prediction, with implementation details chosen for comparability.
- Baseline selection: The evaluation compares MSMs, RMSNs, CRN, and G-Net under the same longitudinal setting as the proposed method.
- Weighting-based baselines: MSMs estimate counterfactual outcomes with inverse probability of treatment weights and linear modeling, using stabilized weights to address time-varying confounding.
- Recurrent baselines: RMSNs use four LSTM subnetworks for propensity estimation, encoding, and decoding, with IPTW-based reweighting during training.
- Recurrent baselines: CRN learns representations predictive of the next outcome but non-predictive of treatment assignment through gradient-reversal adversarial training.
- G-computation baseline: G-Net estimates conditional covariate distributions and uses Monte Carlo simulation, averaging M = 50 samples to predict counterfactual outcomes.
- Implementation: For comparability, the experiments use one- or two-layer LSTMs with representation and prediction layers, while hyperparameters are tuned by random grid search.
J.3. Additional Results
Additional experiments examine one-step and multi-step prediction under different treatment trajectories and report lower RMSE at longer horizons in the random-trajectory setting, where the Causal Transformer remains superior.
- Experimental settings: The additional experiments cover one-step prediction, multi-step prediction with a single sliding treatment, and multi-step prediction with random trajectories.
- Random trajectories: In random trajectories, RMSE decreases as the projection horizon increases because treatment application should reduce tumor volume.
- Random trajectories: The random-trajectory results confirm the superiority of the Causal Transformer.
- Evaluation metric: The reported tables use normalized RMSE, average performance over five runs, and treat lower values as better.
K.1. Data
The data section combines preprocessed ICU records with a semi-synthetic simulator that generates longitudinal covariates, treatments, outcomes, and treatment effects for counterfactual evaluation.
- Real-world data: MIMIC-extract supplies ICU data aggregated hourly, with missing values forward- and backward-filled and continuous time-varying features standardized.
- Semi-synthetic cohort: The semi-synthetic simulator samples 1,000 ICU patients, clips stays to 20–100 hours, and uses 60%/20%/20% train, validation, and test splits.
- Outcome simulation: Untreated outcomes combine endogenous trends, exogenous dependencies on time-varying covariates, and independent noise.
- Treatment simulation: Synthetic binary treatments are assigned sequentially with confounding introduced through current covariates and previous treated outcomes.
- Treatment effects: Treatment effects are modeled as long-lasting, inverse-square-decaying effects within treatment-specific time windows and are aggregated across treatments.
- Evaluation data: The simulator produces three binary treatments, two outcomes, and counterfactual trajectories for one-step and multi-step prediction, with RMSE averaged over five seeded datasets.
L.1. Data
The study describes the real-world MIMIC-III setup, evaluation protocol, and analyses of Causal Transformer representations, runtime, and subnetwork importance.
- Data: MIMIC-III measurements were aggregated hourly, with missing values forward- and backward-filled and continuous time-varying features standardized.
- Data: The experiment used 25 vital signs, 3 static features, two binary treatments, and diastolic blood pressure as the factual outcome.
- Data: The real-world cohort contained 5,000 ICU patients with stays of at least 30 hours, split 70%/15%/15% into train, validation, and test subsets.
- Evaluation: Models were retrained on five random subsamples, and performance was reported as averaged, unscaled test-set RMSE.
- Runtime: Among neural models, single-stage CT had the smallest runtime across synthetic-data confounding levels, combining the architecture with CDC-loss training.
- Representation analysis: CT produced balanced representations with regions appearing non-predictive of current treatment while preserving continuous variation in tumor-volume outcomes.
- Subnetwork analysis: The time-varying-covariate subnetwork had the largest importance score, while treatment-subnetwork importance was near zero for short horizons and increased for longer horizons.