Causal State-Space Model for Causal Inference: Estimating Longitudinal Individual Treatment Effects

AI summary

Overview

Research area: Causal machine learning for longitudinal observational health data — specifically, estimating individual treatment effects (ITEs) over time using deep sequence models.

Technical level: Advanced. The paper combines selective state-space models (Mamba), adversarial domain confusion, contrastive learning (InfoNCE), and information-theoretic proofs (Jensen–Shannon divergence bounds, data processing inequality).

Scope: The paper formalises an information-theoretic conflict between treatment-balancing and outcome prediction in deep causal-inference models, and proposes a state-space architecture with two auxiliary training objectives that resolves it.

What This Paper Is About

Deep models for estimating individual treatment effects from longitudinal data (patient histories of treatments, outcomes, and covariates) typically use "domain confusion" — adversarial training that forces the learned representation to be invariant to treatment assignment. The authors argue this invariance is self-defeating: features correlated with treatment, such as current vital signs, are also independently predictive of future outcomes, so forcing invariance throws away information the outcome predictor needs. They formalise this tension as a mutual information (MI) conflict, derive a Jensen–Shannon divergence bound on the resulting counterfactual prediction error, and propose two models designed to resolve it.

Key Contributions

1. CSSD (Causal State-Space model with Direct decoder): a selective state-space model encoder combined with a parallel multi-step decoder that produces all τ counterfactual predictions in a single forward pass, eliminating the accumulative rollout error of autoregressive decoding.
2. CSSPD (Causal State-Space model with Predictive regularisation and Direct decoder): CSSD augmented with Contrastive Predictive Coding (CPC) and Local Information Maximisation (LIM) objectives, which the authors state resolves the MI conflict without sacrificing treatment invariance.
3. Formalisation of the balancing–prediction MI conflict, with a derived Jensen–Shannon bound on the counterfactual prediction error induced by domain confusion, plus Propositions 1 and 2 showing that the conflict is an information-theoretic inevitability and that CPC and LIM provide tractable lower bounds on the suppressed MI quantities.
4. Empirical evaluation on the real-world MIMIC-III dataset and the synthetic Cancer Simulation benchmark against RMSN, CRN, G-Net, and the Causal Transformer (CT).

Main Findings

• MI conflict is structurally inevitable: Proposition 1 shows that when covariates are correlated with treatment (confounding) and covariates predict outcomes, any encoder reducing I(BR_t; A_t) via domain confusion also reduces I(BR_t; X_t), and consequently reduces the upper bound on I(BR_t; Y_{t+τ}). The result holds conditionally on the regime where X_t is the dominant confounder.
• A bound exposes the trade-off: Theorem 1 bounds the sum of factual and counterfactual error by the two treatment-group factual losses plus a prevalence-weighted Jensen–Shannon penalty W·sqrt(JS_π(P^(0)‖P^(1))). Domain confusion only targets the JS penalty, so it "need not reduce the full bound at all."
• Autoregressive rollout compounds error: single-step error ε becomes O(ε^τ) at horizon τ, which the parallel decoder removes by reading a shared stop-gradient BR_t for every horizon.
• MIMIC-III result: CSSPD achieves lower counterfactual RMSE than the Causal Transformer at every horizon τ ≥ 2, with gains growing from 0.02 (2-step) to 0.07 (6-step). The dataset section notes that because ground-truth counterfactuals are unobservable in real data, normalised RMSE on factual held-out trajectories is what is reported.
• Cancer Simulation result: across confounding strengths γ ∈ {0,1,2,3,4}, CSSPD outperforms CT at γ ≤ 3 with margins of 25.9%–37.0%, while CSSD achieves the lowest overall average RMSE, a 12.7% reduction over CT.
• Variance caveat: Table 1 notes that the CHSD and CHSPD variants exhibit high variance at γ ≥ 2.
• Encoder efficiency: the SSM encoder runs in O(T) time versus the O(T²) cost attributed to the Causal Transformer.
• Prior corroboration: the authors cite Huang et al. (2024) as empirical analysis confirming covariate information loss in CRN and CT.

Methodology in Plain English

The model splits each patient record into three parallel streams — treatments, outcomes, and covariates (plus static baseline attributes) — and embeds each one. Each stream is processed by its own stack of selective state-space layers (Mamba-style), which scan the sequence in linear time while deciding what to remember. A learned "CausalGatedMixer" then fuses the three streams with scalar gates whose starting values encode the assumed causal structure: the treatment→outcome gate starts at 1 (σ(1) ≈ 0.73, near-open) and the outcome→treatment gate starts at −3 (σ(−3) ≈ 0.05, near-closed), with all gates remaining learnable.

The fused output is the balancing representation BR_t. Rather than decoding future outcomes one step at a time, a parallel multi-step decoder uses τ_max separate prediction heads, each reading a stop-gradient copy of BR_t plus an embedding of the hypothetical future treatment sequence. That treatment embedding has two parts: a cumulative mean over future steps and a step-specific term, which distinguishes sequences with the same average dose but different ordering.

Training adds adversarial domain confusion (a discriminator tries to predict treatment from BR_t, and a gradient reversal layer makes the encoder fight it) plus two contrastive objectives applied after a warm-up period: CPC, which uses InfoNCE to predict future representations BR_{t+k} and thereby preserve temporal structure, and LIM, which maximises the MI between the pre-mixer covariate embedding and BR_t to recover covariate signal that domain confusion suppresses. The CSSD loss is L_pred + α·L_DC + L_MS; CSSPD adds λ_CPC·L_CPC + λ_LIM·L_LIM.

Experiments: Cancer Simulation is a PK-PD simulation of NSCLC tumour dynamics under binary chemotherapy/radiotherapy with 10,000 training, 1,000 validation, and 1,000 test samples; MIMIC-III uses 5,000 adult ICU patients with two binary treatments (vasopressor and ventilation), 25 time-varying vital covariates, and diastolic blood pressure as the scalar outcome, split 3,500/750/750. Both use T_max = 60 and τ_max = 5. Models were implemented in PyTorch Lightning and trained on a single CPU with d_model = 32, d_BR = 24, d_state = 16, L = 2 SSM layers, Adam at lr = 10⁻⁴ (MIMIC-III, 64 batches, 300 epochs) or 10⁻³ (Cancer Simulation, 128 batches, 200 epochs), early stopping with patience of 20 epochs, CPC/LIM activated after 120 warm-up epochs, and loss weights λ_MS = 3.5, λ_CPC = 0.05, λ_LIM = 0.1 selected by grid search on the MIMIC-III validation set. All results are reported as mean ± std across 5 seeds.

Why This Matters

Impact on research: The paper shifts attention from "how do we make representations more treatment-invariant" to "what does invariance cost us." By framing balancing and prediction as a mutual information trade-off with a provable error bound, it gives the causal-inference community a diagnostic lens for why adversarial de-confounding models plateau, and it supplies two concrete regularisers as a proposed remedy. The claim of being the first work to formalise the balancing–prediction MI conflict also sets a target for follow-up theory.

Real-world applications:

• Intensive care decision support, such as choosing vasopressor and ventilation strategies, where the evaluation here is grounded.
• Oncology treatment sequencing, where cumulative versus step-specific dose ordering matters (the Cancer Simulation setting).
• Any longitudinal clinical registry where treatments are assigned non-randomly and the goal is to forecast outcomes under hypothetical regimens.
• General observational time-series settings with time-dependent confounding, since the authors state the setup "should apply to other longitudinal data in general."

Industry relevance: Healthcare systems and clinical software vendors evaluating causal forecasting tools need models whose accuracy claims survive confounding. The O(T) encoder cost versus O(T²) attention matters for deployment on long patient histories; the parallel decoder removes iterative rollout at inference, which simplifies serving.

Future Directions

• Extending the theoretical guarantee beyond the explicit assumption that X_t is the dominant confounder (the residual term I(BR_t; Y_{t+τ} | X_t) is assumed small).
• Understanding the high-variance behaviour the paper reports for the CHSD and CHSPD variants at γ ≥ 2.
• Evaluating true counterfactual accuracy on real data, since MIMIC-III only permits factual held-out RMSE as a necessary condition — the paper acknowledges ground-truth counterfactuals are unobservable there.
• Reconciling the reported horizon range: the abstract describes gains at the 6-step horizon and Table 1 is captioned as averaging over τ = 1–6 steps, while the dataset section states τ_max = 5.
• Determining whether the CPC/LIM warm-up schedule (120 epochs) and the chosen loss weights generalise across datasets, given they were grid-searched only on MIMIC-III validation.

Target Audience

Researchers and graduate students in causal machine learning, causal inference, and information-theoretic representation learning; clinical data scientists and ML engineers building longitudinal treatment-effect models; and practitioners familiar with sequence models who want to understand the representational cost of adversarial de-confounding. Readers will need comfort with mutual information, InfoNCE, and state-space sequence models.

Authors’ abstract

Estimating counterfactual outcomes over time from longitudinal observational data is central to clinical decision support. Existing methods rely on domain confusion -- adversarial training that renders representations invariant to treatment assignment -- yet this invariance creates a mutual information conflict: it suppresses treatment-correlated covariate signals necessary for accurate outcome prediction. We formalise this tension via a Jensen-Shannon divergence bound on counterfactual prediction error and develop two complementary models. CSSD (Causal State-Space model with Direct decoder) adapts selective State Space Models with a parallel multi-step decoder that eliminates accumulated rollout error by producing all prediction horizons simultaneously in a single forward pass. CSSPD (Causal State-Space model with Predictive regularisation and Direct decoder) augments CSSD with Contrastive Predictive Coding and Local Information Maximisation to reinforce temporal predictability in the balancing representation and recover local covariate information destroyed by domain confusion. On MIMIC-III, CSSPD achieves lower counterfactual RMSE than the Causal Transformer at every horizon tau >= 2 at O(T) encoder cost, with gains from 0.02 (2-step) to 0.07 (6-step). On Cancer Simulation across confounding strengths gamma in {0,1,2,3,4}, CSSPD outperforms CT at gamma <= 3 (margins 25.9%--37.0%), and CSSD achieves the lowest overall average RMSE (12.7% reduction over CT), confirming the MI conflict analysis. To our knowledge, this is the first work to formalise the balancing-prediction MI conflict and propose a structured resolution through complementary predictive and information-theoretic training objectives.