Effective and Efficient Time-Varying Counterfactual Prediction with State-Space Models
Authors: Haotian Wang
Haoxuan Li
Hao Zou
Haoang Chi
Long Lan
Wanrong Huang
Wenjing Yang
Venue: ICLR 2025 Poster (2025)
Keywords: Time Series; State-space Models; Treatment Effect Estimation
Primary area: causal reasoning
TL;DR: This paper presents a novel method for counterfactual prediction over time series with state-space models.

Abstract: Time-varying counterfactual prediction (TCP) from observational data supports the answer of when and how to assign multiple sequential treatments, yielding importance in various applications. Despite the progress achieved by recent advances, e.g., LSTM or Transformer based causal approaches, their capability of capturing interactions in long sequences remains to be improved in both prediction performance and running efficiency. In parallel with the development of TCP, the success of the state-space models (SSMs) has achieved remarkable progress toward long-sequence modeling with saved running time. Consequently, studying how Mamba simultaneously benefits the effectiveness and efficiency of TCP  becomes a compelling research direction. In this paper, we propose to exploit advantages of the SSMs to tackle the TCP task, by introducing a counterfactual Mamba model with Covariate-based Decorrelation towards Selective Parameters (Mamba-CDSP). Motivated by the over-balancing problem in TCP of the direct covariate balancing methods, we propose to de-correlate between the current treatment and the representation of historical covariates, treatments, and outcomes, which can mitigate the confounding bias while preserve more covariate information. In addition, we show that the overall de-correlation in TCP is equivalent to regularizing the selective parameters of Mamba over each time step, which leads our approach to be effective and lightweight. We conducted extensive experiments on both synthetic and real-world datasets, demonstrating that Mamba-CDSP not only outperforms baselines by a large margin, but also exhibits prominent running efficiency.

Published as a conference paper at ICLR 2025
EFFECTIVE AND EFFICIENT TIME- COUNTER- FACTUAL PREDICTION WITH STATE-SPACE MODELS
Haotian Wang'
Haoxuan Li
Hao Zou
Haoang Chi
Long Lan'
Wanrong Huang
Wenjing Yang!
'College of Computer Science and Technology. National University of Defense Technology
Center for Data Science, Peking University
Tsinggua University
*Intelligent Game and Decision Lab
wanghaot innottiannnll 13 3@nudt edu. hxlisstu.pku edu.
ABSTRACT
Time-varying counterfactual prediction (TCP) from observational data supports the answer of when and how to assign multiple sequential treatments, yielding importance in various applications. Despite the progress achieved by recent ad- vances, e.g., LSTM or Transformer based causal approaches, their capability of capturing interactions in long sequences remains to be improved in both prediction performance and running efficiency. In parallel with the development of TCP, the success of the state state- models (SSMs) has achieved remarkable progress toward Lon-ssequenc modeling with saved running time. Consequently, studying how Mamba simultaneously benefits the effectiveness and efficiency of TCP becomes a compelling research direction. In this paper, we propose to exploit advantages the SSMs to tackle the TCP task, by introducing a counterfactual Mamba model with Covariate-barn Motivated by the over-balancing problem in TCP of the direct covariate balancing methods, we propose to de-correlate between the current treatment and the repre- sentation of historical covariates, treatments, and outcomes, which can mitigate the confounding bias while preserve more covariate information. In addition, we show that the overall de-correlation in TCP equivalent regularizing the selec- tive parameters of Mamba over each time step. which leads our approach to be effective and lightweight. We conducted extensive experiments on both synthetic and real-world datasets, demonstrating that Mamba-CDSP not only outperforms baselines by large margin, but also exhibits prominent running efficiency.
Decorrelation towards Selective Parameters Maamba-CDSP)
INTRODUCTION
Inferring counterfactual outcomes time series data is of critical importance across broad range of domains (Van der Klaauw| 2002 Heidari & Krause 2018 Li et al.) [022]) In particular, time-varying counterfactual prediction (TCP) focuses on estimating counterfactual outcomes over various possible
sequences of interventions (Melnychuk al. 2022), supports the answer of when and how to assign
multiple sequential treatments (Bica et al., 2020 Melnychuk et al. 2022 Huang et al. 2024).
Prior research has demonstrated that TCP under dynamic treatment regimes introduces significantly more challenges than in static settings, primarily due to increasing generalization errors (Alaa & Schaar 2018) and potential biases arising from multiple prediction steps (Frauen et al. 2023). To be specific, the overall confounding bias could accumulate over time if the bias correction is not sufficient in previous time steps (Austin al.) 2006 Huang al.] 2024). To alleviate such issues, recent advances have focused on integrating sequential debiasing techniques into various models. Notable examples include recurrent marginal structural networks (RMSNs) (Lim 2018).
counterfactual recurrent networks (CRN) (Bica et al. 2020), G-net (Li et al. 2020). and the Causal
Transformer (CT) (Melnychuk et al. 2022
"Corresponding author.

---

Published as a conference paper at ICLR 2025
18
160
Mamba-CDSP (Ours) CT
Running Time (minules)
(Ours)
RMSE (normalizedd
120
10
80
6
40
2
4
6
8
10
12
14
16
18
100
200
Prediction
Step
Sequence Length
(a) RMSE with varying prediction time step
(b) Running time with varying sequence length
Figure 1: Effectiveness and efficiency comparison between Mamba-CDSP and Causal Transformer on the Tumor simulator dataset. The time complexity (right) and the estimation error (left, with sequence length as 60) of CT grows exponentially with sequence length and prediction time steps.
However, we observe that existing methods suffer from both the efficiency and the effectiveness in the TCP task. As shown in Fig. (a), regarding the effectiveness, increasing prediction time step from 0 to 12 leads to dramatically increasing estimation error of the CT method (from less than 10 to over 15). On the other hand, when the sequence length increases, the running efficiency of CT scales with nearly quadratic complexity (from 26 min to over 120 min) when increasing the length of imput sequence from 5 to 200. We further inform two critical reasons lying behind such phenomenon: (1) Impact of the backbone. Transformer, i.e., the backbone of CT, is capable of modeling long sequences outside the attention window, while capturing all time stamps in the attention results in quadratic complexity w.r.t. the sequence length Gu & Dao (2023). (2) Drawback of concurrent balancing strategies. First, previous TCP methods Melnnchuk et al. (2022); Bica et al. (2020) often control the overall confounding bias at the final time step. Following exactly the debiasing protocols in static setting, such operations ignore the accumulation of confounding bias over time series (Austin al. 2006; Frauen al. 2023). Second, directly balancing representations across treatment groups could raise the over-balancing problem, especially in the case of time-varying prediction (Huang (et al. 2024), which will further corrupt the representations of covariates and degrade overall estimation performance.
To tackle the backbone issue of TCP, we note that success of Mamba (Gu & Dao 2023). a data- selective state-space model (SSM) architecture (Gu et al. 202 b), offers possibilities due to its capability of mapping across continuous sequences with nearly linear complexity Gu Dao (2023). Meanwhile, address problems of balancing in the TCP task, propose efficient bias correction method called Covariate-based Decorrelation towards Selective Parameters (CDSP) for building TCP- oriented Mamba. Specifically, designed CDSP mechanism introduces de-correlation strategy to address sequential confounding bias by removing the cross-covariance between current treatments and representation of historical covariates, treatments, and outcomes. On the one hand, our CDSP addresses the drawback of overall balancing in the TCP task with a step-to-step bias control manner. We show that our proposed CDSP can be decomposed into orthogonal regularization applied to selective parameters at each time step, concerning the linear properties of Mamba. On the other hand, empirical evidence show that our CDSP implements the bias correction in a covariate-rreservation style, which also alleviates the over-balancing problem. Besides, adapt the Mamba model for TCP, we modify its architecture by replacing the convolutional layer with dropout layer (Gu & Dao 2023) to mitigate overfitting, In summary, main contributions
follows:
(I) We develop an efficient and effective Mamba-batt framework named Mamba-CDSP for the task of TCP. To the best of our knowledge, this is the pioneer Mamba model tailored to counterfactual prediction.
(2) We design novel covariance de-correlation- based mechanism to migrate the sequential confounding bias. By accounting for the trade-off between prediction and bias correction, our CDSP mechanism overcomes the critical drawhack of existing sequential debiasing methods.
2

---

Published as a conference paper at ICLR 2025
(3) We validate the proposed Mamba-CDSP model across synthetic, semi-synthetic, and real- world data, demonstrating that our framework outperforms existing baselines by a large margin in performance with much more efficient training and inference phases.
2 RELATED WORK
Balancing Covariates in Static Counterfactual Prediction. Static methods focus correcting the explicit confounding bias across different groups via diverse strategies, including reweighting (Kuang al. 2020), matching (Stuart 2010) or covariate balancing (Athey et al. (018)). Specifically, previous work have accounted for the fimite-samlle degrading effect of confounding bias raised by observed confounders (Shalit et al. 2017 Alaa & Schaar 2018), even when the ignorability principle holds.
Development on TCP with Statistical Approaches. Earlier Studies on estimating treatment effects focuses in the area of epidemiology, where the g-estimation, Structural Nested Models (SNM) and Marginal Structural Models (MSM) are developed in the regime of statistical analysis (Robins 1986
1999: Robins & Hernan 2008). Meanwhile, bunch of methods has introduced non-parametric
models with uncertainty estimation into treatment effect estimation over the whole population (Xu al. 2016 Roy et al.) 2017, Schulam & Saria 2017). To be specific, family of approaches built
on Bayesian nonparametric models (primarily GP) have been proposed to better encode structure in temporal trends and treatment effects (Shi et al. [012 Arbour et al.) [021] Xu et al. (016). In
recent, multi-task Gaussian process model has been established by decoupling the response trend into individual-leva and unit-level ones Chen et al. (2023).
Development on TCP with Deep Models. Recent advances construct the bias correction on top of a series of deep sequential models, and have achieved state-of-the-ard (SOTA) estimation performance. For instance, RMSMs (Lim) 2018) incorporates the propensity score-based reweighting into recurrent neural networks (RNNs). Subsequently, counterfactual recurrent network (CRN) (Bica al. 2020) Wang al.] 2024a Berrevoets et al.] 2021) combines the RNNs with adversarial covariate balancing techniques. Besides, the G-net (Li ct al. 2020) performs deep G-computation with the simultancous parameterization of the outcome and covariates. To further capture the long-range sequences, the Causal Transformer (Melnychuk et al.) 2022) designed a transformer-tailore architecture with domain-confusion--aseee balancing capture long-range and complex sequences.
Distributional Counterfactual Estimation over Time. Several recent studies have emerged to estimate the entire counterfactual distribution rather than the averaged outcome (Chernozhukov et al., 2013 Kennedy et al.) [023] Wang ct al., 2018). To realize the generation of counterfactual density nigh-dimeniional over time series, Wu et al. (2024) approximates the counterfactual distribution with a generative model without explicitly estimating its density, which enables wider range of application scenarios including continuous treatments.
State-space models (SSMs) for long sequence modeling, Stemming from the signal transformation theory. the state-space model aims to turn a continuous input signal into an output signal (Gu et al.) 2021b Gu & Dao 2023 Gu et al. 20)1a). In recent, (Gu et al. (221aa) has developed an efficient SSM named S4 by constructing the HiPPO operator with convolutional acceleration. The $4 model is further improved by incorporating the parallel scan into S4 layer (Smith et al., 2022). Regrading the expressivity of S4, Mehta et al. (2022) developed the Gated state-space layer by introducing more gating units. In recent, data-dependent framework named Mamba builds generic long- sequence backbone by replacing fixed parameters in S4 into selective parameters (Gu & Dao 2023). As Mamba outperforms Transformers long-sequence, we aim to empower the capability of Mamba for estimating counterfactual outcomes over time series, i.e., constructing new backbone for the area of causal inference.
variety of realistic applications with linear scaling towards
3 PROBLEM DEFINITION
Notations. This paper uses upper-case let- ters (c.g., X) to denote random variables, with lower-case letters (e.g.. æ) denoting the corre- sponding realizations. Besides, the bold let- ters (e.g., V) refer vectors/matrices, and
T=t-1
T
X
A
A
3
Intor Time-Step Relationship Cross Time-Step Relationship

---

Published as a conference paper at ICLR 2025
non-bold letters (e.g., .. V) refer to scalars. Fol- lowing previous protocols (Melnychuk et al.) 2022 Huang et al., 2024), we let i refer to th individual (e.g., a patient) and with histor- ical treatment trajectories over multiple time steps (c.g., a health plan with one month) from a = 1,... T()) At each time step t, one can observe i-th individual's features from four aspects: (1) Treatment Assignment A with da categories (e.g,, a health plan with de kinds drugs to be taken in turn); (2) The time- varying covariates X E Rdd with du dimensions (e. g., some disease-dependent factors of the patient); (3) The static features (*() that keeps invariant over time (e. g., gender and age of the patient); (4) The outcomes V,()) Rdd with dimensions dy. Then the observed dataset can be defined as D = {(= (1) 2()) a, ,y; Ja=1 } where N is the sample number. For the sake of clarity, we omit the superscript (v) for each individual unless needed and use + to represent the overall features,
fai...., Odg
(T
i.e. both (ii and x(1), at time step t.
Problem Setup Throughout this paper, we define historical information at time step as H. - (X,, At-1, Y,, V). where X, (X..... X.). Y, (Y,,..., Y,). and A;- (Ai...,,A... (Melnychuk et al. 2022; Allam et al. 2021). As shown in Fig. [ we follow standards in timevarrying counterfatual prediction (Melnychuk et al. [2022] Huang et al. 2024 Li
ct al. 2021; Bica et al. 2020 Robins Hernan (208). and model the causal relationship among variables follows: (a) The treatment assignment A, is affected by the time-varying covariates X, as well as the history, including Y, X,-1 and Ā₁-1, indicating the presence of time-varying confounders, (b) The outcome Y, is affected by the current covariates, and past treatment A,-1- outcomes Y.-1. and covariates X, (c) The time-varying covariates X, is affected by all past information, including At-1. Y,-1, and X,-1-
Estimation Target. We define the 2 1 the projection horizon for r-step-ahead prediction, with 0 = (0e, Ote,, (+---1 is the sequence of the (imagined) assigned treatments in the future time steps. Following the potential-outcome framework (Splawa-Neyman al. 1990. Rubin 1978), we define the estimand of our problem as E (y++7 [dutrr-]] H:), where H.
(X,, 1,,,A-1 A,-1, Y,, V) is the historical observations.
Basic Assumptions. We present basic assumptions to support unbiased identification for time-varying counterfactual prediction using observational data (Robins & Hernan 2008):
(1) Consistency. If A, ā is a given sequence of treatments. then Y | = Y This means that the potential outcome under treatment sequence ā, coincides for the patient with the observed (factual) outcome, conditional on A, ac.
(2) Sequential Overlap. There is always non-zero probability of receiving/not receiving any treatment for all the history space over time: 0 < P (A: = =a H = h.) < if H. = h1 > 0,
where h is some realization of a patient history.
(3) Sequential Ignorability. The current treatment independent of the potential outcome, conditioning on the observed history: A,Y [a,] H,, Va,. This implies that there are no unobserved confounders that affect both treatment and outcome.
4 METHOD
4.1 PRELIMIARIES
Inspired by the continuous signal mapping (Gu et al. 2021b a Gu & Dao 2023 the family of state- space models (SSM) is derived from solving the tareingying differential equations. To be specific, SSM aims to estimate a mapping from an input signal x to a output signal y: axt) E R + T(a) e R through the transition of hidden state h(t) R.N. The continuous form of SSM models can be
4

---

Published as a conference paper at ICLR 2025
Counterfactual Prediction
MAMBA- CDSP
Projection Layer
Projection Layer
Deblasing SSM
Input
Activation
X
x
Dropout
hg Output
Projection Layer
Projection Layer
C
RMS Norm
Embedding
ÜRefr
CRegr=t-
Historical Oiservations
SSM Equation
Debiasing SSM
Residual Aggregation
Figure 3: The Step-by-step Framework our proposed Mamba-CDSP
written follows:
(t()) = Ch(t) + Br(t), T(a) = Rh(t),
(1)
where C e N is the evolution parameter, B R R R I are projection parameters. In recent, the S4 model (Gu et al., 202 provides discrete form of the equation in equation by discretizing C and into C and as follows:
h4 Chu- Bz, V1 Rhg,
(2)
where C = AAC, B (ep) AB and A represents the time step. Although
$4 has the advantage in its fast computation with convolution operators, the lack of sequential-sensitie attention on different time steps becomes a performance bottleneck. To overcome this issue, the
Mamba Gu & Dao 2023) model incorporates the selective mechanism into S4 by generating A, C and B from the input & with linear-projection layer, i.e., the selective parameters.
4.2 MODELING
TREATMENT EFFECTS USING MAMBA
Input and Output. Our Counterfactual Mamba (C-MAMBA) is designed as an integral block. During the inference phase, treat the historical information H. (X, A,-1, Y, V) with a as the input signal x(( of Mamba, and Y the output signal T(a). As shown in Fig. 3 the input signals are embedded into the representation space with linear embedding. Following (G Dao 2023). , Root Mean Square Layer Normalization (RMS-Norm) is deployed on embedded X, and two projection layers split X into two branches. Inside the SSM, the input at time step 4, i.e.,
(X,, ae, Y to V). are fused with historical information h and produces the hidden state hr.
Adaptation of Mamba Structure. On the one hand, the branch of the left side in Fig. remains the same as in (Gu & Dao 2023). where silu activation layer deployed transform the projected branch. On the other hand, the right-side branch is first filted with a silu activation layer, and then fed into a dropout layer. We note that in this main branch, the original convolutional layer (Gu & Dao, 2023) is replaced with dropout layer. The empirical motivation behind such operations is that we observe that the original Mamba model tends to overfit the training data, and the root-cause of such phenomenon is the 1-d convolution layer /Wang et al. 2024b). Originally, the convolutional layer is adapted to improve the mixture of neighboring tokens in language modeling (Gu & Dao 2023), while such token- token-mixing mechanism might lead to the overfitting on the interaction among X. Hence, remove the convolutional layer with dropout layer to alleviate the overfitting phenomenon of Mamba structure on temporal data. Finally, the two branches are fused a residual manner, i.c., an element- wise multiplication operation and projection layer.
5

---

Published as a conference paper at ICLR 2025
4.3 COVARIANCE BASED DECORRELATION TOWARDS SELECTIVE PARAMETERS
Previous empirical studies (Huang et al., 2024) have pointed out that typical balancing methods for correcting confounding bias deteriorates the representations of the covariates itself. To overcome this issue, we start from the structure of Mamba models and design a novel Covariate-based Decorrelation towards Selective Parameters (CDSP) method to cut-off the correlation between current treatment at and historical information ht-1-
Recalling the transition equation in equation 2 follows:
we first expand the expression of Cov(h4-1, a+) as
Cou(h--1, Cov
Σ
K2 K.X?))
as
Σ Cov (K
E K,Cov (x² ai)):
(3)
where the second and the third equation are due to the property of cross-covariance. Meanwhile, due the term K; represents the time- to the expression of the hidden transition equation in equation 2
accumulated multiplication of O, and B as K, B, By denoting the term Cov { at
as Σ Ext we then formally write the constraint, i.e., ||Cou(hy-1, a;) as follows:
[-1
[[Cou(he-1, at)l -
K
(4)
We further derive the following proposition: Proposition 1. Finding K; to minimize equation equals to minimizing K;Dxt ET
0.
Based on the above proposition, we design our proposed CSDP regularization term as follows:
ET 2xpa LCSDP (C, = Σ B.
(5)
where the term Exp ET "xp.e .a. only concerns w.r.t. the representational versions of Y, A, X and V. and can be pre-computed for each batch of samples. Besides, we fit the outcome prediction by minimizing the factual loss of the next outcome. Such operation can be done via the mean squared emoor (MSE):
MMSE Y
by (h+)
where By represent the parameter of the projection layer after the residual fusion of two branches of our CDSP model (shown in Fig. B}.
Remark. We note that the objective LMSE updates the whole CDSP, while the objective LCSDP only updates the selective parameters C' and B. The overall objective function is LMSE aLCCSDP where a is the regularization parameter.
Remark (Computational Complexity Analysis). We perform computational complexity analysis and compare our proposed CDSP with previous adversarial balancing methods. Denote the overall length of the time horizon as T', and the representation dimension of A and X as d', and d& respectively. For adversarial balancing modules, previous practice shows that the discriminator has usually 2 layers with d' for each layer. Then the overall complexity of our proposed CDSP method is O O(B(dd') + )T), where as the overall complexity of ADB can be derived as O(B((d)) )3 + (d' 2a + [A))T) (see detailed derivation in Appendix A.4.3).
4.4 THEORETICAL ANALYSIS
To further show the superiority of our proposed method over existing balancing methods, we theoreti- cally derive the upper bounds counterfactual prediction risks, and find our CDSP outperforms ex- isting adversarial balançing methods by preserving tighter counterfactual prediction risk bound (see proofs in Appendix A.3). We assume the Gaussian covariates X|a ~ N (Han G.I) I) and linear outcome
We that C and B differs in different time steps.

---

Published as a conference paper at ICLR 2025
structure (aa) W + c", where c" (N, 1). Meanwhile, use Na to denote sample number of each treatment arm, with W and & denoted as model parameters of representation layers and prediction head, respectively. We omit the time-indes. (superscript) for convenience.
Definition 1. The expected precision in Estimation of Counterfaciual Prediction (ECP) is defined as:
CЕCP (W, ④) = -fx (fy(x)) rxx)) p(z)dx.
(7)
Theorem 1. (1) The following risk bound in the finite sample regime holds for the vanilla Empirical Risk Minimization (ERM) model without any balancing modules with probability 1 2m:
vanilla (W, 4b) 2 vril((m mo) 12 u(vor voo)r C ECC
(8y)
(8)
where a
1 NI DE of and (k₂)}=1
denotes some constants w.r.f. the moments of and
(2) The following risk bound the finite sample regime holds for the vanilla prediction cases with ad adversarial balancing (ADB) modules with probability 1 2n7
(2 + oa + +1)((1ra)) ||m u2 +C
+
adb ECP (W, ②) 2
V
-
4
1=1
(3) The following risk bound the finite sample regime holds for the vanilla prediction cases with our proposed CDSP module with probability 1 2n:
( K + 7)
(rir3) 2
cdep FCCP
52 (W,
2
12113 +
(10)
Based on the above results, we demonstrate the superiority of our CDSP method as follows.
Corollary (Over-bale balancig of ADB method), When the following condition holds, the risk bound of vanilla ERM model is more tighter than that of adversarial balancing methods:
- M1 - 10012 - 12+0++0))777
(11)
This illustrates that as the gap between two groups decreases, the harm of over-balancing is much larger than that of observed confounding bias.
In addition, conclude that CDSP method has tighter bound compared with the existing adversarial balancing methods, by noting the following inequality:
(rira)² 2
|| p2l (2 + 00 - (r)(ra)))
11 - 111-32
(12)
Corollary 2 (CDSP outperforms ADB with a tighter bound). Our CDSP method always outperforms ADB methods with a tighter bound, i.e., eday (W, 4 < eadh (W, ④).
5 EXPERIMENTS
Benchmarks. Following common practice in benchmarking for counterfactual inference, all the methods are validated on three datasets, including the synthetic tumor growth data (Geng et al. 2017). the MIMIC-IlI-basse semi-synthetic data (Melnychuk et al. 2022 SShhllam & Saria 2017). the MIMIC-III real-world data (Johnson et al. 2016). To further validate our method with high- dimensional, long-range sequences, follow (Melnychuk et al. 2022 by generating patient observations with outcomes under endogeneous and exogeneous dependencies while considering treatment effects. Details of MIMIC-III data present in Appendix A.4.1
7

---

Published as a conference paper at ICLR 2025
3.5
MSMe RMSNs GRN
CT
MSWs
CT
RMSE (normalized)
2.0
Mamba CDSP
Mamba-HSIC
RMSE (normalized)
RMBNS
CRN
Mamba CDGP
Mamba-HBIC
G-Net
2.5
G-Net
1.5
1.0
2
Strength of time-varying confounding y
Strength of time-varying confounding
(a) 1-step-ahear Prediction
2-step-alie ahead Prediction
MSMe
CT
MSW
RMSE (normalized)
CRN
RMSNs
Mamba-CDSF
RMSE (normalized)
3.0
RMSN
CRN
Mamba-HBBC
Net
G Net
2.0
4
0.5
0.0
Sirength time-varying confounding y
Strength of time-varying confounding
3-step-ahead Prediction
6-step-ahead Prediction
Figure 4: Results for fully synthetic data based on tumor growth simulator, which are reported as the mean performance averaged over five runs with different seeds.
Baselines. Although a bunch of temporal counterfactual estimation methods are present, this paper focuses on several state-of-the-arr (SOTA) models for examination, including the Causal Transformer (Melnychuk et al.) (022), G-Net eLi at. al. 202I). Recurrent Marginal Structural Networks (RMSNs) (018). Counterfactual Recurrent Network (CRN) (Bica et al. 2020) and Marginal Structural Models (MSMs) (Robins et al. 2000). We provide two realizations of Mamba model, including the Mamba-CDST method with covariance de-correlation and the Mamba-HSI method by introducing the covariate balaneing with Hilbert-Schmidt independence criterion (Gretton et al.)
2005). (See details of baselines in Appendix (A..2)
Evaluation Protocols. We evaluate the proposed Mamba-CISS in variety of scenarios, including:
Supervised Prediction. This scenario involves using historical data to predict future counterfactual outcomes over multiple time steps, following standard train-test procedure. For the tumor-growth synthetic dataset, the evaluation protocols are different for one-step prediction and multiple-step prediction. For the semi-synthetic dataset, account for all possible combinations for three binary treatments during the one-step prediction task; and account for the random trajectories with Tmax 10. Finally, for the well-world MIMIC-III dataset, we test all patient observations for one-step prediction, while selecting subset of observations during multiple-step prediction (Melnychuk et al. 2022). (Details of evaluation protocols on three benchmarks are illustrated Appendix A.4.1). All metrics are reported in the form of Root Mean Square Error (RMSE).
Covariate Reconstruction. This scenario involves using learned representations from baseline methods to reconstruct the original covariates. The evaluation metric is the reconstruction loss during the training process. All metrics are reported in the form Root Mean Square Error (RMSE).
Visualization of learned representations. This scenario is to examine whether learned represen- tations are contaminated by balancing strategies. The T-SNE technique is adopted for visualiza- tion (Van der Maaten & Hinton 2008).
The empirical experiments are performed around the following three questions: (1) Can our proposed Hamba-CDSP outperform other baselines across various benchmarks? (2) Can our proposed Mamba-CDs preserves the historical information of covariates when performing confounding bias correction?
Questions.
8

---

Published as a conference paper at ICLR 2025
Table Results for semi-synthetic data for r-step-ahead prediction based on real-world MIMIC-III data, which are the average performance over five runs with different seeds, and means statistically significant results (p-value S 0.01) using the paired-t-test compared with the best baseline.
==1
==2 2
3
r=
==
6
7=7
==
10
MSMs
a.37
a.57
0.74gaus
2.88
L.4,0.20
348
> 10.0
> 10.0
>- 10.0
RMSNs
0.24
0.7
0.60lamm
0.702
0.78
084
0.89
094
4.7.0.00
1.000011
CRN
0.30
0.484n02
0.59
0.65
0.68
0.7
0.72am0:
0.7
0.76e0.01
0,784002
G-Net
0.34
a.67
0.83
400cs
1,10200t
1.16
1.2
1.25
1.29,0.s
EDCT
0.29:0.01
0.46210:
0,56
0.62s0.01
06720011
0.70z00:
0.72sno:
0.74
0.76e0.01
0.78
CT
1.20
0.38
0:46
0.50
1.53
0.51
0.55
a.57
0.58
1.59
po.5
oa.ee
1,51
jam
1.59
Mamba-CDSP
0.19
0.25
0.30
0.34
0.37
0.42
0.43
0.44
0.46
0.47:0.1
Table Results for experiments with the real-world MIMIC-III data, which are the average perfor- mance over five runs with different seeds, and # means statistically significant results (p-value S 0.01) using the paired-t-test compared with the best baseline.
==2
== 3
7 4
5
6.37-0.00
2.06
11. 21.89
13.12
14.44=1.12
RMSNs
5.20_0.15
9.79
10.52:0a0
11.0920.1
11.6420.02
4.84_ace
9.15
10.15_a.1
10.40...m
G-Net
5.13 5.13-mco
11.88
12.91aa
13.57
14.08
4.60=0.s
9.01
9.58
9.89.0.:
10. 12.122
Manba-HSIC.
4.72-0.06
5.19
7.24
8.30
9.25
Mamba CDSP
4.412o.05
5.042o.1a
5.14
5.201o.15
5.25
5.1 COUNTERFACTUAL PREDICTION: ANALYSIS AND EXAMINATION
Performance Analysis on the Tumor-growth simulation. We report the RMSE of the estimation performance (RMSE) across each baseline by controlling the confounding parameter y. As shown in Fig. the analysis conducted w.r.t. different levels of y and T:
(I) With relatively small 7 and e.g., 0 in Fig. (b), the counterfactual prediction results of each baseline, i..., CT, Mamba- Mamba-CDSP, CRN, are close to each other. However, with fixed y and increasing y, e.g., 7 = 0 in Fig. (e) and (f), the counterfactual estimation of Mamba-based methods, including Mamba-CDSP and Mamba-HSIC, outperforms other baselines.
(2) With increasing 7 for each e.g., y from 0 to 4 in Fig. (e) and (f), proposed Mamba- CDSP outperforms other baselines with an obvious margin. Such a phenomenon informs that our designed CDSP debiasing techniques achieve a better balance between prediction and preservation of the covariate information.
Performance Analysis the semi-synthetic and real-world MIMIC-III Dataset. In similar to the Tumor-grow simulation. results on the semi-synthetic and real-world data further validate the superiority of our proposed Mamba-CDSP especially on the real-world prediction (nearly reduces 50% of the RMSE compared to previous baselines). Especially, we note that due to the lack of realistic counterfatual outcomes in real-world data, the corresponding task becomes indeed a factual outcome prediction (Huang et al. (202)). Subsequently, ERM baselines without balancing modules should perform the best, and methods with over-balancing. by contrast, will worse the fitting of factual outcome. Hence, the superiority of our methods real-world data reflects that our proposed CDSP regularization indeed preserves most of the covariate information such that the performance of factual prediction is not hurt. Furthermore, the stability of our Mamba-CDSP the semi cemi-synthetic data with long-sequence prediction results (r = 10) also verifies the effectiveness our method on long-sequence prediction. In addition, we obtain similar trends with empirical analysis the M5 dataset (See Appendix 6))
5.2 FURTHER EMPIRICAL INSIGHT
Reconstruction and Visualization: Better Trade-off Achieved by Our CDSP. We follow Huang et al. (2024) to perform analysis on the Tumor- groooh simulation by constructing and training LSTM- based decoder learned representations from CT. Mamba-HSIC and our Mamba-CDSP, As shown in Fig. (b)) balancing the representations with HSIC leads to non-convergence with y = 8, while the vanilla mamba converges well. By construct, our proposed CDSP support obvious better

---

Published as a conference paper at ICLR 2025
10
50
10
8
40
Loss
Loss
20
2
10
100
150
200
200
50
100
150
200
Training Epoch
Training Epoch
Training Epoch
(a) 7 8, Vanilla Mamha
(b) 7 = 8, Mamba with HSIC
(c) 7 = 8, Mamba-CDSP
Figure 5: Reconstruction loss curves for Mamba with different debiasing techniques.
convergence of reconstruction from learned representations. Besides, we further check the effect of each component of our Mamba-CDSP with ablation study on the semi-synthetic MIMIC-III data Table
Extreme Long Sequence, Sensitivity Analysis, and Failure Cases. Besides, we also test our pro- posed Mamba-CDSP method by comparing with CT on extreme sequences in Appendix (Table 13). We note that the contextual window of CT is set to 200 due to the limit of computational complexity. We also test the sensitivity our proposed by tuning the de-correlation hyper-parameter a in the range of j0, 0.0001, 0.001, 0.01, 0.1, 1.0, 10.0] in Appendix [12). We also note that proposed Mamba-CDii achieves near performance on the TG simulator and semi-synthetic MIMIC-II data when the predictio step is small, e.g.,
1. The underlying reason contains two aspects: (a) confounding effect accumulates along with the time step; (2) the contextual window Transformer baseline, i.c. CT. is the length of the whole sequence most cases.
Efficiency Analysis. We analyze the running efficiency by reporting the parameter amount and running time per module of each baseline in Table g and Table in the appendix for saving space. Results align with the original property of Mamba (Gu & Dao 2023) that the parallel scanning mechanism guarantees that Mamba architecture enjoys superiority over the running efficiency.
6 CONCLUSION
This paper presents pioneer work on studying the SSM-architecturr-base estimator towards counterfactual prediction over time series. By introducing the covariance decorrelation into the Mamba model, our proposed Mamba- CDSP achieves superior empirical performance with better balance between confounding bias correction and covariate information preservation.
Limitations. In view of empirical studies in Huang al. (2024). is necessary and important to consider the design of paradigm that balances bias correction and covariate protection. The covariance decorrelation mechanism. However, our CDSP is designed specifically for linear SSMs, and more effective and general methods require further effort.
Future Directions. (1) Theoretical guidance on the design the balancing technique counter- factual prediction (not limited to time series prediction). As the covariate balancing and outcome prediction serves as two sub-tasks for counterfactual prediction, can we adapt learning theories from multi-task learning (Maurer et al. 2016 Royer et al. (2024) to inspire more reliable and theoretical- complete debiasing methods? (2) Practical guidance on debiasing for sequential counterfactual prediction. When transferring static prediction sequential prediction, especially in multi-step prediction, how will the risk of confounding bias change?
ACKNOWLEDGEMENT
This work was partially supported by the National Natural Science Foundation of China (No. 62372459, 62376282, 623B2002).
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Ahmed Alaa and Mihaela Schaar. Limits of estimating hetcrogeneous treatment effects: Guidelines for practical algorithm design. In International Conference on Machine Learning, pp. 129-138. PMLR, 2018.
Ahmed Allam, Stefan Feuerriegel, Michael Rebhan, and Michael Krauthammer. Analyzing patient trajectories with artificial intelligence. Journal of medical internet research, 23(()): 12, 2021.
David Arbour, Eli Ben-Michael, Avi Feller, Alex Franks, and Steven Raphael. Using multitask gaussian processes estimate the effect of targeted effort to remove firearms. arXiv preprint wrXiv:2100.07000 2021.
Susan Athey, Guido W Imbens, and Stefan Wager. Approximate residual balancing: debiased inference of average treatment effects high dimensions. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 80(4):597- ((4)::97-623,, 2018.
Peter C Austin, Muhammad M. Mamdani, Carl Van Walraven, and Jack V Tu. Quantifying the impact of survivor treatment bias in observational studies. Journal of evaluation in clinical practice, 12 (6):601-612. 2006.
Jeroen Berrevoets, Alicia Curth, Ioana Bica, Eoin McKinney, and Mihaela van der Schaar. Disen- tangled counterfactual recurrent networks for treatment effect inference over time. arXiv preprint arXiv:21 2021.
Ioana Bica, Ahmed M Alaa, James Jordon, and Mihaela van der Schaar. Estimating counterfactual treatment outcomes over time through adversarially balanced representations. arXiv preprint arXiv:2002.04083, 2020.
Ychu Chen, Annamaria Prati, Jacob Montgomery, and Roman Garnett, A multi-task gaussian process model for inferring time-varying treatment effects in panel data. In International Conference Artificial Intelligence and Statistics, pp. 4068 4088. PMLR, 2023.
Victor Cheroozuukov, Iván Fernández- Val, and Blaise Melly. Inference on counterfactual distributions. Econometrica, 81(6):2205-226. 2013.
Dennis Frauen, Tobias Hatt, Valentyn Melnychuk, and Stefan Feuerriegel. Estimating average causal effects from patient trajectories. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7586-7594, 2023.
Changran Geng, Harald Paganetti, and Clemens Grassberger. Prediction of treatment response for combined chemo-and radiation therapy for non-small cell lung cancer patients using a bio- mathematical model. Scientific reports, 7(1):13542, 2017.
Arthur Gretton, Olivier Bousquet, Alex Smola, and Bernhard Schölkopf. Measuring statistical dependence with hilbert-schmidt norms. International conference on algorithmic learning theory, pp. 63-77. Springer, 2005.
Arthur Gretton, Kenji Fukumizu, Choon Teo, Le Song, Bernhard Schölkopf, and Alex Smola. A kernel statistical test of independence. Advances in neural information processing systems, 20, 2007.
Albert Gu and Tri Dao. Mamba: Linear-timg time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00755 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.00330 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572-585, 202
Hoda Heidari and Andreas Krause. Preventing disparate treatment in sequential decision making, In IJCAI, pp. 2248-2255 2018.
11

---

Published as a conference paper at ICLR 2025
Samuel Holt, Jeroen Berrevoets, Krzysztof Kacprzyk, Zhaozhi Qian, and Mihaela van der Schaar. Ode discovery for longitudinal heterogeneous treatment effects inference. In The Twelfth International Conference Learning Representations, 2024.
Qiang Huang, Chuizheng Meng, Defu Cao, Biwei Huang, Yi Chang, and Yan Liu, An empirical examination of balancing strategy for counterfactual estimation on time series. In Forty-first International Conference on Machine Learning, 2024.
Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark. Mimic-iii, freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
Edward H Kennedy, Sivaraman Balakrishnan, and LA Wasserman. Semiparametric counterfactual density estimation. Biometrika, 110(4)8855-86, 2023.
Kun Kuang, Peng Cui, Hao Zou, Bo Li, Jianrong Tao, Fei Wu, and Shiqiang Yang. Data-driven variable decomposition for treatment effect estimation. IEEE Transactions on Knowledge and Data Engineering. 2020.
Rui Li, Zach Shahn, Jun Li, Mingyu Lu, Prithwish Chakraborty, Daby Sow, Mohamed Ghalwash, and Li-wei H Lehman. G-net: deep learning approach to g-computation for counterfactual outcome prediction under dynamic treatment regimes. arXiv preprint arXiv:2003. 2020.
Rui Li, Stephanie Hu, Mingyu Lu, Yuria Utsumi, Prithwish Chakraborty, Daby M Sow, Piyush Madan, Jun Li, Mohamed Ghalwash. Zach Shahn, et al. G-net: a recurrent network approach to g g-computation for counterfactual prediction under a dynamic treatment regime. In Machine Learning for Health, pp. 282-292222222 299, PMLR, 2021.
Bryan Lim. Forecasting treatment responses over time using recurrent marginal structural networks. advances in neural information processing systems, 31, 2018
Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes The benefit of multitask representation learning. Journal of Machine Learning Research, 17(81):1-32, 2016.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206. 12::: 2022.
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Causal transformer for estimating counterfactual outcomes. In International Conference on Machine Learning, pp. 15293-15329, PMLR, 2022.
James Robins. A new approach to causal inference in mortality studies with a sustained exposure period application to control of the healthy worker survivor effect. Mathematical modelling, 7 (9-12):1393-151) 1986.
James Robins and Miguel Hernan. Estimation of the causal effects of time- varying exposures. Chapman & Hall/CRC Handbooks of Modern Statistical Methods, pp. 553-599, 2008
James M Robins. Association, causation, and marginal structural models. Synthese, 121(1/2): 151-179, 1999.
James M Robins, Miguel Angel Hernan, and Babette Brumback. Marginal structural models and causal inference in epidemiology, 2000.
Jason Roy, Kirsten J Lum, and Michael Daniels. bayesian nonparametric approach to marginal structural models for point treatments and continuous or survival outcome. Biostatistics, 18(1): 32-47, 2017.
Amelie Royer, Tijmen Blankevoort, and Babak Ehteshami Bejnordi. Scalarization for multi-task and multi-domain learning scale. Advances in Neural Information Processing Systems, 36, 2024.
Donald B Rubin. Bayesian inference for causal effects: The role of randomization. The Annals of statistics, pp. 34-58, 1978.
12

---

Published as a conference paper at ICLR 2025
Peter Schulam and Suchi Saria. Reliable decision support using counterfactual models. Advances in neural information processing systems, 30, 2017.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: general- ization bounds and algorithms. In International Conference on Machine Learning, pp. 3076-3085. PMLR, 2017.
JQ Shi, Wang, EJ Will, and RM West. Mixed-effects gaussian process functional regression models with application to dose-response curve prediction. Statistics in medicine, 31(66)3115-3117, 2012.
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint urXiv:2208.04993 2022.
Jerzy Splawa-Neyman, Dorota M Dabrowska, and Terrence P Speed. On the application of probability theory to agricultural experiments, essay on principles. section 9. Statistical Science, pp. 465-472, 1990.
Elizabeth A Stuart. Matching methods for causal inference: A review and look forward. Statistical science: a review journal of the Institute of Mathematical Statistics, 25(1):1, 2010.
Wilbert Van der Klaauw. Estimating the effect of financial aid offers on college enrollment: -iipooiiiini approach. International Economic Review, 43(4):1249-1288 2002.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp. 35151-35174. PMLR, 2023.
Lan Wang, Yu Zhou, Rui Song, and Ben Sherwood. Quantilc-optimal treatment regimes. Journal of the American Statistical Association, 1133523):12433-1255 2018.
Shirly Wang, Matthew BA McDermott, Geeticka Chauban, Marzyeh Ghassemi, Michael C Hughes, and Tristan Naumann. Mimic-extract data extraction, preprocessing, and representation pipeline for mimic-iii. In Proceedings of the ACM conference health, inference, and learning, pp.
222-235, 2020.
Xin Wang, Shengfei Lyu, Lishan Yang. Yibing Zhan, and Huanhuan Chen. A dual-module framework for counterfactual estimation over time. In Forty-first International Conference on Machine Learning, 2024a.
Zihan Wang. Fanheng Kong, Shi Feng, Ming Wang, Han Zhao, Daling Wang, and Yifei Zhang. Is mamba effective for time series forecasting? arXiv preprint arXiv:2403.11144 2024b.
Shenghao Wu, Wenbin Zhou, Minshuo Chen, and Shixiang Zhu. Counterfactual generative models for varying treatments. In Proceedings of the 30th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 3402-3413, 2024.
Yanbo Xu, Yanxun Xu, and Suchi Saria. A bayesian nonparametric approach for estimating individu- alized treatment-responss curves. In Machine learning for healthcare conference, pp. 282-300. PMLR, 2016.
Michael Zhang, Khaled K Saab, Michael Poli, Tri Dao, Karan Goel, and Christopher Ré. Effectively modeling time series with simple discrete state spaces. arXiv preprint arXiv:2303.09488 2023.
A APPENDIX
A.1 NOTATION SUMMARY
For of reading, summarize the correspondence between the abbreviation and full notation, as shown in Table In addition, we named our proposed Covariate- based Decorrelation towards Selective Parameters method CDSP.
13

---

Published as a conference paper at ICLR 2025
Table 3: Notation table with detailed correspondence for concept, abbreviation, and full notation.
Concept
Abbreviation
Full Notation
Treatment
A
Outcome
Y
Y X V B, C
Time- Vene-Varying Covariates
X
Time-invariant Covariates
V
Selective Parameters
B, C,C
Projection Parameters
R
R
Discrete Time Stamp
A
A
Covariance Matrix
E
Σ
Hidden State of Mamba
h
Time Horizon
T
A.2 DERIVATION OF Loss FUNCTION
Proposition 1. Finding K; minimize equation 4 is equivalent minimizing K for each 1 <ist->
ET 8h 0
> Proof. Recalling the expression | sa 12.. we only need to prove that Ex .a. X .a is the condition achieve minimization of X, ar First, show that, due to the fact that DT E.L.A = Σ Ex *x is semi-positive, the term Tr (KE,,, is convex function such that deriving the derivative w.t.t Tr ,K coincides with the minimization of Tr (KE,,,,1
K,
aTr (K, Ei KT) AT E.t, Tr (aaK, BK2) Eieh (aK1 + +K²K)
asTr (K2 (KI-K2)
K2) Enen (K1 K2)"),
(13)
0.
where final equation is due to the fact that Di,,, is semi-positive. Since the term Tr (K (22g,hh K) is convex w.r.t K, we then derive the derivative w.r.t. K as follows:
ÔTr (K KKE, K")
K
(14)
0K
Then let K = Ki, and our claim follows.
A.3 THEORETICAL ANALYSIS
We consider the historical covariates X,-1, treatment A₁ and outcome Y, in the last timestamp first introduces
but omit the time index for clarity. Assuming the binary treatment with A several basic definitions with lemmas as follows:
Definition 2 (Expected Prediction Error). Taking the common protocols Shalit al. (007 Bica
et al. (2020); Melnychuk et al. (2022) that the counterfactual prediction models can be decomposed into representation layers & and prediction layers W, we introduce the definition of prediction error in each treatment arm as follows:
/ (2, (),(X AA= 1)dr
CF (W,
(15)
f (n, ())=
0=0 (W,
where fy. L(y, (((()), a)) (r(T(a) = X = x)dy refers to the individualieed fulividuaizzeduucuulll factual prediction error, and L refers to the prediction loss per sample, e.g., MSE loss.
Then introduce existing fundamental lemma follows:
14

---

Published as a conference paper at ICLR 2025
Lemma (Risk Bounds from Shalit et al. (2017)).
(EC 2) - 2 (eF
a= (W, 2) + ep a= (W, Φ) + v/PMG (po .P -p + C)
(16)
where and C are constants w.r. model parameters, the integral probability metric (IPM), i.e. IPMa(u, supmeG Ig ((r)( v(r)] dr Shalit et al. (2017) defines the probability diver- gence between two probability measures u, with The and of a famuly of functions G.
Then we first illustrate an immediate lemma which is necessary for our following derivation:
Lemma (Empirical Prediction Error). We first define empirical prediction error for predicting the factual outcomes in each treatment arm a follows:
1
cF (W, ()
Σ
L(m a = =1)
Na=1
t=1
(17)
No
f (W, 3)
Σ
L(yi, (f(;;),
No=0
Assuming that the variance of the loss function is bounded, ie., Var(lu=o) S Ko and Var(I==1) S Ki. then have the following inequalities hold with probability at least 7, respectively:
c (W, 4)
CF (W, 1) S
"
(18)
lep 2 (W, x) CJ - +) -
Proof. The resulting inequalities is immediately obtained by invoking the Chebyshev's Inequality w.r.t. each empirical prediction error term.
Lemma 3 (Equivalence between De-correlation and First-moment Matching). We state that, our proposed CDSP method, i.e., Cov(h, A) = 0, is equivalent to matching first-order moment of historical covariates across the treatment arms.
Proof. We first have:
Cov( pz. A) E A| - PAA 1).
1, A 1)P(A= P(A = 1)E[
(E [A = 1] - E(),,) P(A=1)
(IE
A = 1] (E( AA= 1A=1)P(A 1) 1) + E AA= 0 ()) P(A
(EE A = 1] 2 A= P(A = ()
(19)
Then our claim follows.
Lemma 4 (Construction of Global Minima of Adversarial Balancing). Based on the sufficient and necessary condition in Theorem 4.1 in (Melnychuk et af.) 22022), a treatment-invaiian representation D, that &.A A uchieves the global minima of Traditional adversarial balancing (ADB) modules.
Lemma (Construction of Global Minima of Our CDSP). By contrast, any representation 4 that matches the first-order of historical covariates of our proposed CDSP regularization term achieves the global minima.
Let W and d denotes model parameters of prediction models without any balancing modules, Wa and & denotes model parameters of prediction models with adversarial balancing modules, and aa and denotes model parameters of prediction models with our CDSP modules. We further assume that:
Theorem 1.
(a) Gaussian covariates per group such that X N(pa; for 0,
15

---

Published as a conference paper at ICLR 2025
(b) Linear structure of i.e., Y(a)= W + where ~ (0,, Bich 1 R2 and Elc" 1 I Ki., Such an assumption is rationale, as introducing the model mrissppecification error into prediction models with and without balancing modules equally will not affect our analysis.
(c) The representation layers embed historical covariates are linear, e.g., linear trans- former (Von Oswald et al. 2023) and linear state-space models (Zhang et al. 2023).
(d) The prediction layers of Y are implemented with linear regression based on representation, i.e., the OLS regression.
(e) The scaled loss, i.e., L, belongs to the set of all couplings between p and F
(f) The representation matrix t has bounded operator norm, i.e., ||| S 11 and bounded F-norm, ie.,
| T2-
(g) The prediction purameters Wa, Wa Wa has bounded 2-norm 13-
Then: (1) the following risk bound in the finite sample regime holds for the vanilla prediction cases without any balancing modules with probability 1 2n:
2 2 j=1
ECP
panilla (W, 2
v79]((1 - po) + u(vo- - +
Kgj +
(20)
where o Ni EE 42.
Proof. As linear transformations preserves the Gaussian distributions, the representations of X still follows the Gaussian distributions: X Ja ~ N( () Then by invoking the Lemma
with the fact that the I-Wasserstein distance is governed by the 2- Wassersstein distance (Jensen Inequality), we have that:
<FCP [, ②)
22 + (W, + e Wasss (pg P4 p2-1) CC
(W,9) +)++)++ + (100(11 1o)11 + +[[[o-- - Voo) (TT 1/2 113,) + +C)
(W + 6) (rīll(11 - mo0113 (Voi /oo) + +C)
< 22 +) + + ( b) + (r7ll(m2 - -mo)ll + (Vo1 voo)rz) +C),
(21) where the first equality is due to the closed-form expression of 2- Wasserstein distance between Gaussian distributions, and the second inequality due to the fact that Φ symmetric and semi comi-positive definite. We further invoke Lemma ( and substitute the expected prediction error terms with empirical prediction error terms, and the following inequality holds with probability at least
1 2n::
CECP (W, () 2 (W,
(W, Φ) + (rillH4 mo)1l + (Vo1 voo)r + +C)
(v
$ 2
cF 4) c (W, ②) + (rill(m) po) (VoT -o + +
C
Kg
KI
n
7
(2) the following risk bound in the finite sample regime holds for the vanilla prediction cases with adversarial balancing modules with probability - 21y:
( mg +o -0000 1) (22)
(2 do + (r)(3)))) 1/ Hella +
adb ECP W, $) 2
4
where 1. EN4, No Ei=10 622
16

---

Published as a conference paper at ICLR 2025
Proof. Based on Lemma further construct representation mapping. i.e., 4, as global minima
of adversarial balancing methods including (Melnychuk et al. 2022 Bica et al. 2020 Wang et al. 2024a) follows:
① ( (o
M1
Poy
2
(23)
140
(1)
A=)) = (z
2
such that the resulting distribution are treatment-invariant and the global minima achieved. We then check the factual prediction term follows:
N,
(W,
W DT,jI + +.
VooW
111
140
)12
N1
2
N2
Σ
(c- (c=1)2+ 11loooW
11 - 1213
(24)
N1
i=l
(1 (1+00) (0)
(c)=)2+
| - 1111-
4
On the other hand, the divergence term equals to zero:
wWasss (pg 0=1 Pb -55 = 0,
(25)
then our claim follows.
(3) the following risk bound in the finite sample regime holds for the vanilla prediction cases with our proposed CDSP module with probability 1 21y:
Rg + ))
(rira)² 2
2 1
cdep EECP (W, ②) 2
111 - +C
(26)
j=1
where J
N1 NN
Ni
Proof. Based on Lemma we further construct a representation mapping, i.e., 4, as global minima of our proposed CDSP method as follows:
1/1
M4
(3)AA1) (Z
2
(27)
(<)A=o) = (
H40
2
We note that in order to achieve global minima, our CDSP only requires to align the mean across different treatment groups without aligning the variance terms. We then derive the divergence term as follows: Wass2 (pg ..--) S ((Vo1 Voo)rr)
(28)
and the prediction error terms as derived as follows:
N1
W, ①) =
Σ
W Wi + + -WiD(a
#11 - 140
)12
N
2
1e]=1 + 1 W 11113111111111 - m2l1
(29)
N
i=1
(c) 1c=1)2 2 (r1ra)2 1| 12113
N
and our claim follows.
17

---

Published as a conference paper at ICLR 2025
Proposition 2. [Over-balancing of Adversarial Learning] When the following condition holds, the risk bound of vanilla Empirical Risk Minimization (ERM) model more tighter than that of adversarial balancing methods:
-
| moll 2 S or (2 on)rf13/
(30)
Remark. The above conclusion very intuitive, as the gap between two groups decreases, the harm of over-balancing is much larger than that of observed confounding bias.
Proposition 3. [Why CDSP outperforms ADB] Our CDSP method always outperforms ABD methods with a tighter bound.
Proof. The result is immediate from the fact that:
(rira)²
(2 oo + (1)(r)))) |//9 12112
111 22211
(31)
always holds.
L
Remark. The above conclusion is also very intuitive. Especially, on the factual outcome prediction tasks, i.e., real-world datasets, our CDSP outperforms CT with large margin.
Remark. Our first conclusion intuitive, as the gap between two groups decreases, the harm of over- wer-bbalnncing is much larger than that of observed confounding bias. The second conclusion reflects an important idea hidden behind our CDSP: compared with optimal representations achieved by adb methods, e.g., N I) and ((u, ag/) are aligned via representation layers at both means and covariances, our proposed CDSP only aligns the means of two Gaussian representations with each other, i.e., Ao and μ4 after some transformations, Subsequently, when performing prediction, the extent of the loss of predictive information embodied in covariates for our method will be smaller than that of adb methods. On the other hand, the effect of debiasing observed confounding remains nearly the same when comparing our methods with adb methods. Hence, our de-correltaion method achieves lower risk with tighter bound.
A.4 EXTRA EXPERIMENTAL DETAILS
A.4.1 BENCHMARK DETAILS
Tumor Simulator. As the state-of-the-art cancer effect simulator, the tumor growth (TG) model from (Geng et al. (017)) characterizes the volume of tumor Y++1 for 1++ 1 days after cancer diagnosis. More specifically, the TG simulator contains the following components:
Two binary treatments: (1) radiotherapy (A,) when assigned to patient has an immediate effect d(() on the next outcome; (2) chemotherapy (Aş), which affects future outcomes in exponentially decaying approach scaled by C(();
One-dimensional Outcome, i.e., the volume of tumor Y(+1- More specifically, Yo41 is affected by A& and A& via the following equation:
(K Y,
K
Y 1 plog
B.C. (a,d, + B,d,)) + t Yı:
(32)
where K, Be, ar,, Pr, E (0, 0.012) are simulation parameters, and Be, a,,, 8r describe the individual response of each patient.
varying confounding effect. By treating the past outcomes as the confounders, current treatment will be assigned via a biased approcah:
A A ~ Bernoulli (o(=>(Dis (Fi-(is(YY-D)) - D Draxx D /2))
(33)
where ((-) is the sigmoid function, Dmux is the maximum tumor diameter, D1s (Y-1) and y hyper-parameters.
18

---

Published as a conference paper at ICLR 2025
Overall, the critical parameter in TG simulator, ie., Y, controls the degree of confounding effect such that the confounding bias is controlled. With increasing values of y, the the amount of time-narying confounding bias enlarges.
Finally, during the testing phase, we simulate counterfactual trajectories separately for one-step prediction and 7-step prediction. For the former, simulate all four possible combinations of Y e+1- Following protocols in (Melnychuk et al. (022). we adopt two evaluation protocols to test 7-step prediction, i.e., Single sliding treatment and Random trajectories. To be specific, single sliding treatment simulates a single treatment, where the treatments are iteratively moved over window ranging from to + Tmax 1. Meanwhile, random trajectories are simulated fixed number, i... 2 (Tmax 1). of randomly sampled trajectories.
Besides, for each y, we simulate 10,000 patients for training, 1,000 for validation, and 1,000 for testing. The length for sequence is limited to 60 time steps.
Semi-synthetic Dataset. Following protocols (Meinychuk al. 2022 we construct the semi- synthetic dataset based on the real-world data extracted from the MIMIC-extrrat (Wang et al. 20200) with a standardized preprocessing pipeline in (Johnson et al. 2016). As done in (Melnychuk et al..
2022), we overcome missing values with forward and backward filling, and perform Standardization of all the continuous time-varying features. Following protocols, we extract 25 different vital signs, ie., time-varying covariates with 3 static covariates (gender, ethnicity, and age), and encode all static covariates in the onc-hot approach. Overall, the resulting dimension of features is 44.
By extending the idea in (Schulam & Saria 2017), the semi-synthetic simulation protocols in (Melny chuk et al., 2022) first generate untreated trajectories outcomes under endogeneous and exogeneous dependencies and, then, sequentially apply treatments to the trajectory
1,000 patients are randomly selected, with ICU stays lasting least 20 hours. For these patients, their ICU stays are limited to a maximum of 100 hours such that the length of stay (T()) is between 20 and 100 hours.
dy untreated outcomes Z .j -....... dy is simulated for each patient i from the cohort:
z (.)(0)
a B-spline(t) + a g (t) + ory (2((())
Ee
(34)
endogenous
exogenoos
noise
with (O, 0.005". and o'g- ag. and at are weight parameters, Meanwhile, B-spline(() is sampled from a mixture of three cubic splines, (-)( )) is sampled independently for each patient from Gaussian process with Matérn kernel; and (.) is sampled from a random
Fourier features (RFF) approximation of Gaussian process.
* The da binary treatments A/, 1 = 1,..., dar are simulated in sequential approach:
PAt = σ A' (V,-1) + (X))
(35)
A ~ Bernoulli (PA:),
(36)
where the random function f-((X, are treated the confounding effect, and a(-) refers the sigmoid activation, 7/4 and 7'x are confounding parameters, by is fixed bias.
The outcome is generated via the addition of the untreated outcome Z7 and the simulated treatment effect E (t):
Y Z Et),
(37)
where Z/ characterizes the effect of the treatment bias within a time window:
Bi(1) Σ min d. 1 (AA=
(38)
ww - 1)2
where Btg is the maximum effect size of treatment 1.
We refer detailed introduction of simulation protocols in (Menyychuk 2022 & (-) is sampled from an RFF approximation of a Gaussian process (similar to (-) Melnychuk al.
2022
19

---

Published as a conference paper at ICLR 2025
By setting du 3 and dy 2, the cohort of 1,000 patients is split into train/vallidation/ees subsets via a ratio 60% 20% 20 %.
MIMIC-III Real-world Dataset. Following protocols in (Melnychuk et al. 2022 Huang et al. 2024), we also adopt the Standardized pre-processing pipeline (Johnson et al. 2016) with forward and backward filling for missing values. Keeping the same as in semi-synthetic data, the features, i.e, the potential confounders, contain 25 vital signs (dz = 25) and the 3 static features. Two binary treatments (d,, = 2): vasopressors and mechanical ventilation pressure is the outcome to be estimated. A cohort of 5,000 patients is randomly selected from the patients with intensive care unit (ICU) stays of least 30 hours.
adopted. The diastolic blood
The train/validationn validatiooos// subsets are split with the ratio of 70%\/5%//15% For one-step prediction, all samples are adopted the testing data. For multiple-step prediction with 2 2, we choose patients with observation lengths of at least 6 with rolling origin.
M5 Real-world Dataset. The M5 Forecasting dataset, as cited in (Huang et al. 2024), comprises daily transaction data from Walmart stores across three U.S. states, along with comprehensive details on products, stores, pricing, and significant events. For our analysis, we repurpose this dataset to estimate treatment effects, designating product pricing as the treatment variable and product sales as the outcome variable. All other features are treated as covariates. Since this dataset derived from real-world observations and lacks counterfactual data, evaluate the performance of various models, including their Empirical Risk Minimization (ERM) variants, in predicting factual outcomes. Notably, GENT and MSM models excluded from this analysis due convergence issues with this dataset.
A.4.2 BASELINE DETAILS
We categorize the previous baseline methods based on their architectures:
Transformer:
The Causal Transformer (CT) method is proposed in (Melnychuk et al.) 2022). Three sub-networks are constructed in CT, with cross-attention to capture interactions among predictors. For confounding bias, the domain confusion objective proposed by adversarially training domain predictor such that the learned representations of historical covariates are only predictive of the outcome.
LSTM:
G-Net (Li al. 2021) A deep network version of the G-computation, which estimates the following G-formula (Robins 1986):
E Y c+r 11 H.)
E (Y 14 1 x
Jt XRdo
(39)
L+T-1
P H - 11
J=H+1
where the final estimate is given by first sampling P(X, H 3 - from MonteeCarlo,, and then Y,-r(ār,t is taken as the empirical mean (Melnychuk et al. 2022).
Counterfactual Recurrent Network (CRN) deploys a single LSTM-layer-base based decoder and an encoder. To remove the confounding bias, the gradient adversarial training is adopted to confuse the covariate classifier so that the resulting representations only inform the outcome.
Recurrent Marginal Structural Networks (RMSNs) is also an LSSTM-based deep estimation model by incorporating the inverse probability of treatment weights (IPTW) (Robins et al.) 2000) re-weighting method for bias correction. To be specific, the stabilized weights, i.e., A(An
An-1) and (A, H.), are estimated via L.SM networks.
Shallow Models: Marginal Structural Models (MSMs) (Robins et al. 2000) is the linear realization of the IPTW. Similar to RMSNs, the stabilized weights are learned via the linear-logistic regression.
20

---

Published as a conference paper at ICLR 2025
Mamba-HSIC: When projecting the three-dimensions tensors from Mamba, i.c., he-1 (sample index, time index, feature index), onto the two-dimensional tensor for final prediction, we extract the representations from the final fully-connected prediction layer, denoted as h MSSC to perform the
HSIC regularization. To be specific, we perform HSIC regularization through the following equation:
HSIC(K
HSCC ..)) = Tr ManseJJaay
(40)
where == I- (L is the common dimension shared by husic and az), K is the kernel Gram Matrix (Gretton et al. 2007). We here deploy the standard Gaussian Kernel throughout our experiment.
A.4.3 COMPUTATIONAL COMPLEXITY ANALYSIS
We perform theoretical analysis on computational complexity by comparing our proposed CDSP with previous adversarial balancing methods. Assuming the overall length of the time horizon is T, and the representation dimension of Y, A and X are c, d and d respectively. Moreover, for adversarial balancing modules, previous practice shows that the discriminator has usually 2 layers with d for each layer.
For our proposed CDSP methods, the complexity of the debias process can be divided into four stages:
(1) Pre-computing the covariance terms, i.c., > Xx ΣT "xy before the beginning of the training process. Due to the complexity of matrices multiplication, we derive the complexity of this component (BB((d)) d')));;
(2) Computing the accumulated selective parameters K; = B, which owns the com- plexity as ( ))));
(3) Computing the multiplications between Exp,a, ET - the complexity as (() )7);
and K; = B. which owns
(4) Computing the 2-norm of the matrix as O(B(d))
Hence, the overall complexity of our method can be derived ( + dd))).
For traditional adversarial balancing methods, the complexity can be divided into three parts:
Embedding representations of X and A with two layers matrix multiplications, which owns the complexity of (( )3 + (d'x dz YT) in each round;
(2) Computing the softmax and the cross-entropy loss, which owns the complexity of O(B A(T), where |A] the cardinal number of A.
Hence, the overall complexity of ADB can be derived O(B((d))) (d +A()T))
A.5
IMPLEMENTATION DETAILS
A.5.1
IMPLEMENTATION ON MODEL PARAMETERS
Details method. We detail all the parameters, including the model layers, units of each layer, dropout rate, together with the batch size, EMA, and input-output size in Table
Details of other baselines. We exactly follow the sapee--prreeter tuning protocols in Table 6 & 7 in (Melnychuk et al. 2022) for each baseline we have compared with throughout our experiments.
A.5.2 SENSITIVITY ANALYSIS ON CONFOUNDING PARAMETER
Furthermore, to inform how our proposed method performs towards the correction of confounding bias, we supplement detailed experiments on the TG simulator by tuning the confounding parameter, ie., y, and compared the SOTA method, i.c., CT, with our proposed Mamba-CDST in Table 7 As ? usually to 0 4 in previous empirical studies Melnychuk et al. (2022) Bica al. (2020),
21

---

Published as a conference paper at ICLR 2025
Table Ranges for hyperparameter tuning across experiments. Here, we distinguish (1) data using the tumor growth (TG) simulator (=experiments with fully-synthetic data), (2) data from the semi- synthetic benchmark, and (3) real-world MIMIC-III data. EL refers the embedding layer, and PL refers to the projection layer.
Model
Hyperparameter
TG simulator
Semi-Synthetic Data
Real-world Data
Mamba blocks (B)
1
Learning rate (x)
{0.0005. 0.001, 0.01)
{0.0005, 0.001, 0.013
{0.0005, 0.001, 0.01}
Minibatch size
128
64
De-coreelation Parameter
1
1
EL hidden units L)
32
32
64
Mamaba-CDSP
PL hidden units (drr.)
32
32
64
Dropout rate (p)
0.1
0.1
0.
EMA model weights
0.99
0.99
0.99
Input size
da + d2 + dy + d, dg
d. + dg + dy + de dy
dg dg + dy + d. dy
Output size
increase the value of y to 8,10,16,20 in our analysis such that much higher confounding effect exists in the TG simulation data. Throughout our comparison, we found that our proposed Mamba-CDSP substantially outperforms CT with a large margin.
A.5.3 EXTRA EXPERIMENTAL RESULTS
We also provide extra experimental results including: (1) Extra visualization results at a = 10 in Fig- 6 (2) Performance under single treatment slide the TG simulator in Table T0 Further visualization study on the real-world MIMIC-III dataset also validates the above justification. where the representation learned from Mamba-CDSP enjoys a similar distribution to that from the vanilla Mamba model in Fig. 7
Iree
L-SNE Component 2
20
Grup
5SNE Component
GrgA
SNE Component
20
-20
-20
-40
-40
Grea
-80
-100
-80-60-40-20
USNE Companent
Component
Component
Vanilla Mamba, T 10
(b) Mamba with HSIC, T 10
(c) Mamba-CDSP T 10
Figure 6: Representation visualizations of Mamba, Mamba with HSIC, and Mamba-CDSP.
A.5.4 COMPARISON WITH ODE- BASED ESTIMATOR
Furthermore, we also note the existence of the application of Ordinary differential equations (ODE)
on the problem of TCP, i.e., the Insite method in Holt et al. 2024). However, also note that due to
the limit the overall time-series process can be characterized using an ODE system Holt et al. (2024).
Hence, we only compare it with our proposed Mamba-CDSP
the TG simulation benchmark (the
only overlap between Insite Holt et al. (2024) and previous expreience Meinychuk et al. (022)) Bica
al. (2020): Li al. (020); Huang al. (2024)). Results are present in Table 11
A.5.5 RE-EXAMINATION THE BIAS CORRECTION
In this section, we supplement extra experimental results to understand how our proposed CDSP correct the confounding bias, we compare the training loss curves of the HSIC criteria and our CDSP criteria on the real-workh dataset. To be specific, as HSIC criteria measures how the covariate distribu- tions shift across different treatment groups, use the value of HSIC term inform the confounding
22

---

Published as a conference paper at ICLR 2025
Goop Oxea Ovea
60
Grap
Cragi
1-SNE Comparent
ESE Component
40
Bisa
5 Component
Grepp
20
-20
-20
20
-40
-40
-60
-60
40
-60
40
100
Compoment
Component
Component
(a) 50, Vanilla Mamba
(b) 7' 50, Mamba with HSIC
50, Mamba- Mamba-CDSP
Figure 7: Representation visualizations of Mamba, Mamba with HSIC, and Mamba-CDSP.
Table Ablation study for proposed Mamba-CDD CDSP on semi-synthetic MIMIC-III Data. Reported: normalized RMSE with relative changes. w/ HSIC-Ref refers to the Mamba model debiased by only HSIC.
T= 1
==3
==5
==
T=9
Mamba-CDSP (Our proposed)
0.19
0.30
0.37
0.43
0.48
w/ convolution layer
-0.01
+0.01
+0.01
+0.04
+0.11
w/o dropout
+0.08
+0.09
+0.11
+0.09
+0.10
Ablation Model
w/ selection mechanism
+0.05
+0.07
+0.12
+0.16
+0.19
w/o RMS-norm
+2.13
>10
>10
>10
>0
w/o residual
-0.3
+0.06
+0.06
+0.06
+0.07
w/o Cov-Reg (a 0)*
+0.16
+0.17
+0.16
+0.22
+0.25
Ablation Loss
w/ HSIC-Reg
+0.06
+0.08
0.100
+0.07
-0.07
bias. As shown Fig. 8 the HSIC term also decreases when our CDSP regularization term decreases.
Such phenomena informs that method indeed effectively reduces the confounding bias through performing de-correlation.
Criteria
6x 10-
Independence gap (HSIC)
10
De-correlation gap (CSDP)
2000
4000
6000 training step
8000
10800
12000
14600
16000
Figure 8: Examination of the Bias Correction MIMIC-III Dataset.
23

---

Published as a conference paper at ICLR 2025
Table 6: Results for experiments with the real-world M5 data, which are the average performance averaged over five runs with different seeds, and means statistically significant results (p-value s 0.01) using the paired-t-test compared with the best baseline.
2
3
4
6
RMSN
5.19
9.73
10.48
11.07 20.43
11.63
12.16
CRN
4.
9.1540.17
9.79
10.12.0.17
10.38+0.19
10.59
CRN w/o balancing
4.68
9.05+0.36
9.68
10.00=0.16
10.26+0.18
10.47
CT
4.59
4.994
9.59
20.25
10.1540.25
10.35+0.25
w/o balancing
4.59
8.9840.15
9.58 18
9.90
10.1340.23
10.34
Mamba-CDSP
4.08
4.05
4.59
5.87
6.47
8.39
Table 7: Sensitivity Analysis on Confounding Factors the TG Simulator, where results are reported as the mean of different seeds, and * means statistically significant results (p-value I 0.01) using the paired-t-test compared with the best baseline.
Model
y=
T 3
7 =
=
CT
2.98
3.46
3,56
3.52
3..3
Mamba-CDSP
2.97=0.01
3.36".o.00
3.2910.01
3.02_.o.0
2.69 10.01
CT
10.4610.01
9.38
8.501..01
7.63
6.83
10
Mamba-CDSP
4.25
5.2610..0
5.5810.01
5.55
0..0
CT
15 00.93
13.54_0.01
15. 25.54
16.404a.01
16.58=0.01
16
Mamba-CDSP
8.992n.01
9.062.0.0
8.0910.01
6.932n..1
5.95
CT
16.19
15.84
17.0940.01
16.55
18.32
Mamba-CDSP
14.14%o.oo 10.01
12.73"
10.26
8.21
6.52
Table 8: Runtime of experiments (per training stage per inference stage) on the real-world dataset. Experiments are carried out on x NVIDIA GeForce RTX 3090 GPU, where ITPS refers inference time per sample.
Stages of training & infercace
Total runtino (in min)
training time
inference time
ITPS (in sec)
MSMs
logistic tegressions for IPTW & lincar regression
1,5
100.0
0.5::0.2
0.04
networks for IPTW
38
35
0.24
CRN
encoder & decoder
109 _00.5
84
25
2
G-Net
single network & MC sampling for inference
8945.
0.24
single multi-ine caplti-ipput network
156
98 _03
58
4.64
Ours
Single Mamba Block
1242.5
10
240.1
0.16
Table 9: Total number of trainable parameters of models across synthetic data (Syn); semi-synthetic (SS) data, and (3) real-world (RW) MIMIC-III data.
Syn Data
SS data
RW data
7= 0
1
2
=
3
4
MSMs
<100
3K
RMSNs
20K
4K
23K
21K
22K
477K
947K
CRN
4K
6K
8K
7K
8K
165K
219K
G-Net
3K
2K
3K
4K
3K
151K
310K
CT
11K
11K
10K
10K
10K
45K
69K
Mamba-CDST
12.8K
12.7K
12.8K
12.8
12.9K
26.8K
27.3K
24

---

Published as a conference paper at ICLR 2025
Table 10: Results for experiments with the real-world MIMIC-III data, which are reported as the mean performance averaged over five runs with different seeds, and * means statistically significant results (p-value I 0.01) using the paired-t-test compared with the best baseline.
Methods
2=
7=1
==2
==3
7
1.360111
g0.r
1.9490:7
2.47 g2ss
RMSNs
0.74
g.76
0.3
096g01
1.28
0.87
20.0s
0.74
1672118
1.5
20.00
1.27
1.1g
1.26
10.04
1.70
g.91
10.37
Mambu HSIC
0.73
2.76
0.79
2.84
Mamba-CCDSP
20.04
0.72.20.4
0.78
0.700
MSMs
1.68.0 ia
2.30
2.f2004
2.62
RMSNS
0.78gpos
a.79a..0o
0.892
100
28.89
0.70
2.69
0.82.1005
1962100
1.36
G-Net
1.23
166
1338
1,57 20.39
CT
0.70_noo
0.71 10.00
0.77.10.03
10.45
Mamba HSIC
1..7.
0,68
0.76
1,26
Mamba CDPP
0.68:004
0.70ao.0o
0.73
0.78
0.89 =0.31
MSMs
1.
1.15
2.19
2.79
2,62
RMSNs
0.22000
2.84
0.9662..1
1.22
1.42
CRN
0.73
0.7220.04
0.9022..7
2.20
1.57 2.57
7=
G.Net
1.39
2.35
1.82
1.54
1.77
CT
0.73
0.75
0.8222004
1.09
1. 1,76g0.9
Mamba HSIC
0.69
0.76
0.85 10.01
1.22
1.92 10.35
Mamba
Mamba-CDSP CDSP
0.70_oo
0.74gn.0
0.78
0.85
1.13
MSMs
1.96_o.s
2.55
2.81
2.54
RMSNs
0.86
0.89 Lo.0n
1.03 17
1.35
1.43 10.30
CRN
0.77
0.76
2.98
2.36.100
1,73
G-Net
1.43
1 1.96
1,67
1.91 20.47
CT
0.76
2.79
0.87
1.15
1.92 20.37
Mamba-HSIC
0.782004
0.76g0.04
0.91
1.22
1,78
Mamba CDSP
0.
0.75
0.75_n.0
0.84 on
j..1
L,13
MSMs
1.97_00
2.21 2.21 1...0
20.1
2.73g
2,42 11.00
RMSNs
0.
0.94 Lo.ne
1,08 an
1.47
1.4 10.00
CRN
0.81
2..04
2005
2.8
G-Net
1.54.0m
1.49
2.06
1.76g22
1.99
CT
0.79g00s
g.82
0,91
121ea.00
20.5
Mamba-HSIC
0.80
g085
0,922005
1.02au12
1.65
Mamba-CDDSP
0.77.m005
0.87.005
0.30
Table 11: Comparison with our Mamba-CDSP and Insite on synthetic TG simulator with the random- trajectory evaluation protocol, and means statistically significant results (p-value < 0.01) using the paired-t-test compared with the best baseline.
7
7- 2
7
Insite
1,02z0m
1.01
1,03
1,4290.2 1.15gnms
Mamba-CDSP
0.00
0.02
0.632000
0.68
0.72
Insite
1.00
1.03
1.03
1.17
1.440000
Mamba-CDSS
0.62gam
0.64 nn
0.65
0.70
0.6 me
1.00_oot
1.04
1.04
1.18
1.48
Mambu CDSP
0.64
0.73
0.79 jo
L12" 0..00
Insite
0.99.am
10.5
1o05
1.2
Mambu-CDSP
0.49" j..0
noi
0.83
091 0.15
Insite
0.98_not
10.00
1008
1.2gam
1,53
Mamba
00.00
n.0
0.83
0.11 0.15
25

---

Published as a conference paper at ICLR 2025
Table 12: Sensitivity analysis on the semi-synthetic MIMIC-III dataset by tuning the de-correlation parameter results are the average performance over five runs with different seeds.
1..2
_.m
j.75
1.05
1.02
1.09 .c.m
10.0
0.30
0.37
0.442mm
0.49
0.53
0.57
0.60:mm:
0.62:mmm
0.64:..1
0.65
0:27
0.33
0.30
0:43
0.46
0.49
0.52.00
0.54gool
0.56g
0.01
0.25
0.32g00:
0,35g00:
0:44
0.42
0.46g000
0.47
0.48
0.53ga.
0.53
0.1
0.21
0.28
0.33
0.41
0.39
0.15
0.15
0.45
0.77ccm
1.19
1.0
0.19
10.00
0.34.0.01
0.37.001
0.42.001
0.43
0.44
0.46..0.0
0.47
10.0
0.23g0.01
0.27
0,34gom
0.36g001
0,38
0.44200)
0.57,00
0.59.amm
0.52gc.1
0.50g0.01
Table 13: Comparison with extreme long sequences with a = 1000 on TG simulator with 7 = 3, and * means statistically significant results (p-value < 0.01) using the paired-t-test. Due to computational complexity of CT, we limit the range of the attention window CT 200.
==2 2
3
4
7= 5
7 6
CT
1.82
2.
2.77
2.89
3.32
3.85
Mamha-CDSs CDSP
0.512
0.86
20.8%0.00
1.37
1.54
10.21
26


Reviews:
[{"review_id": "3JkPrjsB0Q", "reviewer": "Reviewer_WoTq", "summary": "This paper presents a Time-shared Heterogeneity Learning from Time Series (THLTS) method which infers the shared part of latent factor across time steps with a variational auto-encoders (VAE), the method could capture the hidden heterogeneity by recovering the hidden factors and incorporate it into the outcome prediction. This method can be a flexible component to be easily inserted into arbitrary counterfactual outcome forecast models. The authors demonstrate the effectiveness of THLTS on (semi-)synthetic data in capturing shared patterns by combining several existing counterfactual outcome forecast methods to improve their performance.", "questions": "Please list up and carefully describe any questions and suggestions for the authors. Think of the things where a response from the author can change your opinion, clarify a confusion or address a limitation. This is important for a productive rebuttal and discussion phase with the authors.\n\n\n\nBesides RMSE, it would be good to add other ablation study such as distribution analysis of the counterfactual prediction from utilizing the proposed method vs. baselines, which would provide more evidence to validate the the effectiveness of introducing the shared latent factor as illustrated in Figure 1\nIn table1 and table2, could the author elaborate on more details of the baseline THLTS(v)? Why author think this would be fair baseline to justify the rationality of learning shared part of latent factors compared to the proposed method\nIn Algorithm 1, what is the difference between forecast model pρ() and gρ()?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "A substantive assessment of the strengths of the paper, touching on each of the following dimensions: originality, quality, clarity, and significance. We encourage reviewers to be broad in their definitions of originality and significance. For example, originality may arise from a new definition or problem formulation, creative combinations of existing ideas, application to a new domain, or removing limitations from prior results. You can incorporate Markdown and Latex into your review. See https://openreview.net/faq.\n\nOriginality\nThe paper proposes a novel approach to capturing hidden heterogeneity in time series based counterfactual prediction, which is a significant domain problem in causal learning. The proposed Time-shared Heterogeneity Learning from Time Series method is a novel method that addresses this specific challenge by encoding the shared time-aware latent confounder and then utilizing them for counterfactual outcome forecasting.\n\nQuality\nThe paper provides a clear and well-structured presentation of the proposed method, including a detailed explanation of the shared latent confounder variable encoding process via VAE and how to adapt to time series data.\nThe experimental results basically demonstrate the effectiveness of the proposed method in improving the performance of mainstream models. \n\nClarity\nThe paper is well-written and easy to follow, with clear explanations of technical concepts and methods. The authors provide an informative context in each section that effectively organizes the story and summarizes the paper contributions.\n\nSignificance\nThe proposed THLTS method has the potential to improve the accuracy of counterfactual outcome in time-series data scenarios. The capture of hidden heterogeneity across time domains is a common challenge in many fields. The proposed method is flexible and can be easily inserted with arbitrary causal modeling framework, making it a valuable contribution to the field.", "weaknesses": "Lack of Novelty in Methodology\nThe proposed THLTS method is based on the use of variational encoders (VAE) to improve counterfactual prediction in time-series data, which has been explored in other works such as\n1. https://doi.org/10.1145/3637528.3671950, \n2. https://doi.org/10.48550/arXiv.2310.18615 . \nWhile the shared latent factor encoding part is new, the underlying methodology is not entirely novel. To strengthen the contribution, the authors could provide a more detailed comparison with existing methods and highlight the specific advantages of their approach.\n\n\nLimited Experimental Evaluation\nThe paper only presents experimental results on (semi) synthetic datasets, which may not accurately reflect real-world scenarios. To demonstrate the practical applicability of the proposed method, it would be beneficial to include experiments on (or the connection to) real-world datasets .", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "do561AvSM1", "reviewer": "Reviewer_TFBs", "summary": "The paper works with a time varying counterfactual prediction method using STATE-SPACE model.  It introduces methods that de-correlate between current treatment and historical covariates. They claimed that their model is effective and lightweight. Finally, they performed experiments on several datasets to highlight the efficacy of their method.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "(1) New problem on TCP on state-space model\n\n(2) Design of novel de-correlation mechanism to reduce confounding bias.", "weaknesses": "(1) My major concern is notation and presentation of the paper: The paper has too many overloading of notations-- for example, \"a\" or the actions are giving variable A_t but the system parameter is also A. This has been quite confusing to me for sometimes.  \n\n(2) Re. experiments: I am not sure, results of Table 2 are statistically significant: I was looking for paired t test to see how well their method is effective with respect to baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lWR5KP2DW1", "reviewer": "Reviewer_FCGz", "summary": "This paper introduces Mamba-CDSP, a novel approach for time-varying counterfactual prediction (TCP) using state-space models. The key innovation lies in combining the Mamba architecture (a recent advance in state-space modeling) with a new Covariate-based Decorrelation towards Selective Parameters (CDSP) mechanism. The method addresses two major challenges in TCP: computational efficiency and the over-balancing problem. The authors demonstrate superior performance over existing methods like Causal Transformer and G-Net on both synthetic and real-world datasets.", "questions": "1. How sensitive is the CDSP mechanism to the choice of decorrelation threshold?\n2. Could the authors provide more insight into the computational complexity trade-offs between CDSP and traditional balancing methods?\n3. How does the method perform on extremely long sequences (e.g., >1000 timesteps)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Technical Innovation:\n\n-Novel combination of Mamba architecture with TCP\n\n-Well-designed CDSP mechanism that addresses known limitations\n\n-Efficient implementation with linear time complexity\n\n2. Practical Value:\n\n-Better handling of long sequences\n\n-Improved computational efficiency\n\n-Real-world applicability demonstrated on MIMIC-III dataset\n\n\n3.Experimental Design:\n\n-Comprehensive ablation studies\n\n-Multiple evaluation scenarios\n\n-Reasonable baseline comparisons", "weaknesses": "1. Theoretical Analysis:\n\n-Limited theoretical justification for why CDSP works better than traditional balancing\n\n-Could benefit from more formal analysis of the bias-variance trade-off\n\n2.Empirical Validation:\n\n-Could benefit from more diverse real-world datasets\n\n-Limited discussion of failure cases\n\n-More detailed hyperparameter sensitivity analysis needed", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yNggOWj2Uq", "reviewer": "Reviewer_1NGR", "summary": "The paper presents a novel approach to TCP using a Counterfactual Mamba model with Covariate-based Decorrelation towards Selective Parameters (Mamba-CDSP). It addresses the limitations of current methods, particularly those based on LSTM and Transformer architectures, in capturing long-sequence interactions effectively and efficiently. The proposed Mamba-CDSP aims to mitigate confounding bias by decorrelating current treatments from historical covariates and outcomes. The authors provide empirical experiment over both synthetic and real-world datasets, demonstrating that Mamba-CDSP significantly improves prediction performance and running efficiency compared to existing methods.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The motivation of this paper is well claimed.\n2. The writing of this paper is good and easy to follow. The preliminaries part is particularly clear.\n3. The experiment part is comprehensive, consisting of synthetic and real-world datasets.", "weaknesses": "1. The novelty and technical contribution of this paper are not enough. This paper claims to address the deconfounding and over-balance issues in the TCP task, but these issues are not first proposed by this paper. Below, I will further explain my opinion.\n2. The authors claim that they are the first to consider the step-by-step deconfounding in the TCP task. However, to my knowledge, there are several existing works that achieve a similar goal, such as [1] and [2]. The authors should illustrate the difference between them.\n3. Why would the linear correlation between $a_t$ and $h_{t-1}$ lead to the over-balancing issue while the non-linear correlation would not, as you remark in line 292? Can you provide a rational analysis regarding this method?\n4. This method seems to simply replace the transformer with the Mamba model by slightly tailoring the backbone, which lacks technical contribution.\n5. The sizes of the datasets used in the experiments are quite small; for example, there are only 5,000 patients in the MIMIC-III real-world dataset, which cannot adequately reveal the efficiency of replacing the Transformer with the Mamba model.\n\n\n\n[1] Estimating Counterfactual Treatment Outcomes over Time through Adversarially Balanced Representations. ICLR 2020.\n\n[2] Counterfactual Generative Models for Time-Varying Treatments. KDD 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k7zUGNim9W", "reviewer": "Reviewer_PLDR", "summary": "This paper proposes Mamba-CDSP, a novel approach for time-varying counterfactual prediction (TCP) based on state-space models (SSMs). The key contribution is adapting the Mamba architecture with a covariate-based decorrelation mechanism to handle sequential confounding bias while preserving covariate information. The authors demonstrate superior performance compared to existing methods on synthetic and real datasets.", "questions": "1. Is CPSD on line 316 a typo?\n\n2. Which dataset was used for Table 3?\n\n3. How sensitive is the method to the choice of decorrelation hyperparameters?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Novel application of SSMs (specifically Mamba) to TCP, showing promising results in both effectiveness and efficiency. The paper leverages state-space models for counterfactual prediction, achieving significant improvements in both prediction accuracy and computational speed compared to existing methods.\n\n2. Well-motivated decorrelation approach that addresses key limitations of existing balancing methods. The proposed CDSP mechanism offers a novel solution to the over-balancing problem in sequential settings, effectively balancing between confounding bias correction and preservation of important covariate information.\n\n3. Comprehensive empirical evaluation across multiple datasets and settings.", "weaknesses": "1. Limited theoretical analysis of why covariance decorrelation works better than traditional balancing approaches.\n\n2. While performance improvements are shown, deeper analysis of where/why the improvements come from would strengthen the paper. For instance, Table 2 shows substantial gains from CDSP on the MIMIC-III real-world dataset, but this is puzzling since we cannot observe counterfactuals in this data and thus confounding bias should have minimal impact on evaluation. The authors should explain why CDSP shows such dramatic improvements if the test metrics don't actually measure counterfactual prediction ability. This suggests the gains might come from other aspects of the method beyond bias correction, which deserves further investigation.\n\n3. A more thorough literature review on temporal counterfactual estimation would enhance the paper by incorporating recent works like,\n   - Chen et al, A Multi-Task Gaussian Process Model for Inferring Time-Varying Treatment Effects in Panel Data\n   - Wu et al, Counterfactual Generative Models for Time-Varying Treatment\n   - Wang et al, A Dual-module Framework for Counterfactual Estimation over Time\n   - Berrevoets et al, Disentangled counterfactual recurrent networks for treatment effect inference over time\n\n4. The paper lacks sufficient implementation details for reproducibility. While the model architecture is described, key details such as hyperparameters (hidden dimensions, number of layers), the decorrelation coefficient, and dropout rates are not specified. These details are crucial for reproducing the reported results.\n\nThe reference for domain adversarial learning on line 269 is incorrect, for example Lim 2018 did not use domain adversarial learning strategy", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)