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

2026-08-08

AI Summary

This work addresses the loss of treatment-related covariate information in longitudinal observational data caused by domain confounding. To tackle this issue, the authors propose a Causal State Space Model that employs a parallel multi-step direct decoder to circumvent error accumulation from autoregressive forecasting. By integrating contrastive predictive coding with local mutual information maximization, the model simultaneously recovers confounding-corrupted covariate information and enforces representation balance. The study formally characterizes, for the first time, an inherent mutual information trade-off between balanced representations and predictive performance, and resolves it through a complementary information-theoretic and predictive objective. Evaluated on MIMIC-III and Cancer Simulation benchmarks, the method reduces RMSE by up to 37.0% compared to Causal Transformer while maintaining linear encoding complexity O(T).

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.

Problem

• causal inference
• counterfactual estimation
• longitudinal data
• treatment effect
• domain confusion

Innovation

• Causal State-Space Model
• Mutual Information Conflict
• Contrastive Predictive Coding
• Domain Confusion
• Longitudinal Treatment Effects