Effective and Efficient Time-Varying Counterfactual Prediction with State-Space Models

We sincerely appreciate the reviewer’s great efforts and insightful comments to improve our manuscript. We thank the reviewer for recognizing our novel application of SSMs, well-motivated decorrelation approach, as well as comprehensive empirical evaluation. In below, we address these concerns point by point and try our best to update the manuscript accordingly.

W1: Limited theoretical analysis of why covariance decorrelation works better than traditional balancing approaches.

Response: We thank the reviewer for pointing out this issue, and we have added many theoretical analysis to show the superiority of the proposed CDSP method than traditional balancing approaches.

[Main Theoretical Result] CDSP has a tighter counterfactual prediction risk bound than traditional balancing approaches.

• To measure the counterfactual prediction accuracy, we adopt expected precision in Estimation of Counterfactual Prediction (ECP) as evaluation metric defined as $\epsilon_\mathrm{ECP}(W, \Phi)=\int_{\mathcal{X}}\left(\hat{\tau}_f(x)-\tau(x)\right)^2 p(x) dx$.

• Consider the Gaussian covariates $X_{\mid a} \sim N(\mu_{a}, \sigma_{a}I)$ and linear outcome structure $Y(a) = W^{a}\Phi X_{\mid a} + \epsilon^{a}$, where $\epsilon^{a} \sim N(0,1)$. Let $N_{a}$ be the sample number of each treatment arm, $W$ and $\Phi$ be the model parameters of representation layers and prediction head, respectively.

• The risk bounds of (1) vanilla Empirical Risk Minimization (ERM) method without any balancing modules, (2) adversarial balancing (ADB) methods, and (3) our proposed CDSP method are derived as:

  • For (1) vanilla ERM method without any balancing modules, with probability at least $1-2\eta$, we have $\epsilon_{ECP}^{vanilla}(W, \Phi) \leq 2(vr_1^2 |\mu_1 - \mu_0|^2_2 + v (\sqrt{\sigma_1} - \sqrt{\sigma_0})r_2^{2} + C - (\sqrt{\frac{1}{\eta}}(\sum_ {j=1}^{2} \sqrt{\kappa_j}) + \overline{\sigma})),$ where $\overline{\sigma} = \frac{1}{N_1} \sum_{i=1}^{N_1}{\epsilon^a_{i}}^2$.
  • For (2) adversarial balancing (ADB) methods, with probability at least $1-2\eta$, we have $\epsilon_{\mathrm{ECP}}^{adb}(\tilde{W}, \tilde{\Phi}) \leq 2(\frac{(2+ \sigma_0 + \sigma_1)(r_1r_3)^{2}}{4} \|\mu_1-\mu_0\|^2_2 + C - (\sqrt{\frac{1}{\eta}}( \sum_{j=1}^{2}\sqrt{\kappa_j}) + \overline{\sigma} ) ).$
  • For (3) our proposed CDSP method, with probability at least $1-2\eta$, we have $\epsilon_{\mathrm{ECP}}^{cdsp}(\tilde{W}, \tilde{\Phi}) \leq 2(\frac{(r_1r_3)^{2}}{2} \|\mu_1-\mu_0\|^2_2 + C - (\sqrt{\frac{1}{\eta}}( \sum_{j=1}^{2}\sqrt{\kappa_j})+ \overline{\sigma} ) ).$

• With the above risk bounds, we then derive two main results:

  • (a) Over-balancing of ADB methods: When the following condition holds, the risk bound of vanilla ERM model is more tighter than that of adversarial balancing methods: $\|\mu_1 - \mu_0\|_{2}^{2} \leq \frac{v(\sqrt{\sigma_1} - \sqrt{\sigma_0})^2 r_2^2}{vr_1^2 - (2+\sigma_1 + \sigma_0)r_1^2r_3^2/4}.$ This illustrates that as the gap between two groups decreases, the harm of over-balancing is much larger than that of observed confounding bias.
  • (b) CDSP outperforms ADB with a tighter bound: Our CDSP method has a tighter bound compared with the existing ADB methods, i.e., $\epsilon_{\mathrm{ECP}}^{cdsp}(\tilde{W}, \tilde{\Phi}) \leq \epsilon_{\mathrm{ECP}}^{adb}(\tilde{W}, \tilde{\Phi})$.

• We have added the above theoretical analysis in a new Section 4.4 with detailed proofs in Appendix A.3.

[Minor Theoretical Result] The computational complexity of our CDSP is smaller than traditional balancing approaches if the representation dimension of $A$ is small.

• We also added 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^h_A$ and $d^h_X$, respectively. For adversarial balancing modules, previous practice shows that the discriminator has usually $2$ layers with $d^{h}_{X}$ for each layer.

• Then the overall complexity of our proposed CDSP method is $\mathcal{O}(B((d_{X}^{h})^3 + d_{A}^{h})T)$, whereas the overall complexity of ADB can be derived as $\mathcal{O}(B((d_{X}^{h})^3 + (d_{X}^{h})^2 d_{Y}^{h} + |A|)T)$ with $B$ as a constant related to parameter $\overline{B}$ (see detailed derivation in Appendix A.4.3).

• From above, the overall complexity of our proposed CDSP method will be smaller than the ADB methods if the representation dimension of $A$ denoted as $d_{A}^{h}$ is small, i.e., $d_{A}^{h} \leq (d_{X}^{h})^2 d_{Y}^{h} + |A|$.

• From the above results, we further validate the effectiveness of our Mamba-CDSP on the real-world dateset.

Add a brief discussion of why our methods achieves significant improvement on real-world datasets

• As refered by As refered by Huang et al. (ICML 2024), prediction on the realistic data equals to factual outcome prediction due to the absence of realistic counterfactual outcomes.
• Hence, improvement on the MIMIC-III and M5 datasets verifies that our method achieves successful preservation on the covariate information.
• We have supplemented a more detailed discussion in our revised paper (P9-10, Lines 482-497).

W3: A more thorough literature review on temporal counterfactual estimation would enhance the paper by incorporating recent works

Response: Thank you for your kind advice! We have re-write the related-work section (P3, Sec 2) with a more thorough and complete review by adding over 15 new recent works including but not limit to the 4 works reviewer mentioned.

• We have supplemented two extra aspects: Development on TCP with Statistical Approaches and Distributional Counterfactual Estimation over Time and grouped the 4 works reviewer mentioned as follows.
• Development on TCP with Statistical Approaches contains the statistical temporal causal methods:
  • Chen et al. proposes to disentangle the trends of unit and group when predicting the ATT.
  • Wang et al. proposes to integrate historical information together based on adversarial balancing modules.
  • Berrevoets et al. considers covariate disentangle for more precise separation between confounders, instruments and precision variables.
• Distributional Counterfactual Estimation aims to model the distribution of outcome rather than the mean:
  • Wu et al. proposes a new counterfactual generation method with the aid of GANs.

W4: 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.

Response: To enhance reproducibility of our paper, we have supplemented the following details parameter configuration in our revised PDF (P22, Table 6).

Thanks for your detailed and through response in time! We have uploaded the whole revised PDF, the revised component are marked in red. In below, we clarify your further questions point by point.

[Q1 Theoretical Insights]

• The bound in case 2 is tighter compared to case 1 when distributional shifts of covariates between different treatment groups decreases: $\|\mu_1 - \mu_0\|_{2}^{2} \leq \frac{v(\sqrt{\sigma_1} - \sqrt{\sigma_0})^2 r_2^2}{vr_1^2 - (2+\sigma_1 + \sigma_0)r_1^2r_3^2/4}.$
• The above inequality informs that the when first-moment of covariates ($\mu_1, \mu_0$) across different treatment groups are close to each other, the profit brought by balancing will be smaller than the harm caused by balancing.
• how can our theory explain why balancing strategies sometimes lead to performance degradation:
  1. The risk bound consists of Wasserstein term and supervised loss term, representing for observed confounding bias and prediction loss respectively;
  2. The profit of balancing corresponds to the Wasserstein term;
  3. In our theory, the process of balancing is formalized as the variation of both $E[X]$ and $Var(X)$ across $T=1$ and $T=0$ at the same time, i.e., distributional alignment at the representations;
  4. At the same time, such variation causes the prediction error to be larger;
  5. When the Wasserstein term is small, i.e., $\|\mu_1 - \mu_0\|_{2}^{2}$ is small, the profit of balancing will be small than the enlarged error of prediction.

[Q2 Adequate Empirical Verification]

• Regarding experiments on real datasets, the loss of covariate information is caused by over-balancing. A recent benchmarking study has conducted extensive empirical studies to inform the phenomena of over-balancing [1].
• In our paper, wo would like to clarify that we have already conducted extensive experimental validations to verify such phenomena:
  1. Reconstruction Experiments:
     • We analyze if models, specifically the our proposed Mamba-CDSP, can accurately reconstruct original covariates with and without the balancing module.
     • Results are shown in Fig.5 in our revised paper. (CT w/o balancing and CT are already shown in Fig.6 in [1]);
     • For CT, without balancing modules, the re-constructed covariates curve converges; while the curve is difficult to converge with adversarial balancing (An existing conclusion in [1]);
     • For Mamba-CDSP, our CDSP ensures a more converged training curve during re-construction.
  2. Visualization Experiments:
     • We perform representation visualization to further inform how distributions of $X$ are changes with and without balancing modules at various time steps in Fig. 6 and Fig.7;
     • Compared to CT with adversarial balancing, our CDSP achieves better preservation of covariate information;

[Q3 There doesn't seem to be much difference.]

• In the M5 real-world table, we rounded each value to two decimal places. We now list the original results to show the risk that balancing may not lead to more superior prediction performance:

• Moreover, we have supplemented experiments on the real-world MIMIC-III dataset by comparing CT with and without balancing modules as follows:

• Compared to the M5 real-world dataset, the balancing on CT on the MIMIC-III dataset has a more significant performance damage compared to pure CT without balancing, due to overbalancing would cause loss of covariate information.

Reference
[1] Huang et al, An Empirical Examination of Balancing Strategy for Counterfactual Estimation on Time Series, ICML 2024.

• For (1) vanilla ERM method without any balancing modules, with probability at least $1-2\eta$, we have $\epsilon_{ECP}^{vanilla}(W, \Phi) \leq 2(vr_1^2 |\mu_1 - \mu_0|^2_2 + v (\sqrt{\sigma_1} - \sqrt{\sigma_0})r_2^{2}+ C - (\sqrt{\frac{1}{\eta}}(\sum_ {j=1}^{2} \sqrt{\kappa_j}) + \overline{\sigma})),$ in which the Wasserstein balancing error term explains the risk due to insufficient balancing in ERM: $2(vr_1^2 |\mu_1 - \mu_0|^2_2 + v (\sqrt{\sigma_1} - \sqrt{\sigma_0})r_2^{2}).$
• For (2) adversarial balancing (ADB) methods, with probability at least $1-2\eta$, we have $\epsilon_{\mathrm{ECP}}^{adb}(\tilde{W}, \tilde{\Phi}) \leq 2(\frac{(2+ \sigma_0 + \sigma_1)(r_1r_3)^{2}}{4} \|\mu_1-\mu_0\|^2_2 + C - (\sqrt{\frac{1}{\eta}}( \sum_{j=1}^{2}\sqrt{\kappa_j}) + \overline{\sigma} ) ),$ in which the supervised loss term explains the risk due to the loss of covariate information by overbalancing: $2\frac{(2+ \sigma_0 + \sigma_1)(r_1r_3)^{2}}{4} \|\mu_1-\mu_0\|^2_2.$
• This shows a trade-off between fitting the factual outcomes (ERM method) and covariate balancing (ADB method).
  • From our theory, the advantage of balancing methods is formulated as the reduction of both the mean difference $|\mu_1 - \mu_0|^2_2$ and the variance difference $\sqrt{\sigma_1} - \sqrt{\sigma_0}$, i.e., distributional alignment at the covariate representations;
  • When the Wasserstein balancing error term is small enough, i.e., $\|\mu_1 - \mu_0\|_{2}^{2}$ is small enough, the profit of balancing will be smaller than the extra supervised prediction error due to the loss of covariate information, causing the performance drop of balancing methods compared to vanilla ERM methods.
• As per your suggestion, we are conducting experiments to validate this trade-off between the two terms, and we will follow up this message once we get the experimental results : )

> W2: 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.

Response: Thank you for the detailed and constructive comments. We would like to address your concerns by (1) recap the synthetic and semi-synthetic experiments with simulated confounding biases, (2) add synthetic experiments with increasing confounding bias, (3) add real-world experiments on the M5 dataset, and (4) add a brief discussion of why our methods achieves significant improvement on real-world datasets.

Recap the synthetic and semi-synthetic experiments we have done with simulated confounding biases

• In our original manuscript, we have evaluated our methods on the synthetic dataset (Figure 4), semi-synthetic dataset (Table 1), and real-world dataset (Table 2).
  • Synthetic dataset and semi-synthetic MIMIC-III dataset: We follow M van der Schaar, et al. (ICLR 2020) to conduct our synthetic experiments, and follow Melnychuk, et al. (ICML 2022) to conduct our synthetic experiments on the MIMIC-III dataset. Both experiments are conducted with the same data generation process (DGP), in which the confounding bias indeed exists due to the simulated counterfactual outcomes;
  • Real-world MIMIC-III dataset: We follow van der Schaar, et al. (ICLR 2020) and Melnychuk, et al. (ICML 2022) to conduct our real-world experiments on the MIMIC-III dataset, in which the counterfactual outcomes are absent.

Add synthetic experiments with increasing confounding bias

• To further validate that our method can correct for confounding bias, we added synthetic experiments with increasing confounding bias on the tumor growth (TG) simulator.
• We follow Melnychuk, et al. (ICML 2022) to control the level of confounding via tuning $\gamma$ in Eq. (33) $\mathbf{A}^c_t, \mathbf{A}^r_t \sim \operatorname{Bernoulli}( \sigma ( \frac{\gamma}{D_{max}})(D_{15}(Y_{t-1}) - D_{max} / 2))$. For $\gamma=0$, the treatment assignment is fully randomized. For increasing values, the amount of time-varying confounding becomes also larger.

• From above results, we conclude that our proposed method significantly outperform with $\text{p-value} \leq 0.01$ using the paired-t-test compared with Causal Transformer for correcting the confounding bias. We have supplemented such results in our revised paper (Table 9, P24).

Add real-world experiments on the M5 dataset

• We also follow Huang et al. (ICML 2024) to add experiments on a new real-world time-series prediction dataset, i.e., the M5 dataset, which provided by Walmart, captures the unit sales data of a diverse range of products sold across the United States, structured as grouped time series.

To verify the current theoretical findings regarding the over-balancing issue as well as the tighter risk bound of our CDSP method, we add synthetic experiments with varying confounding strength $\gamma$.

• We adopt the following three metrics:
  • Wasserstein balancing error (named Wasserstein in the Table) in our derived bound arises from insufficient balancing $2(vr_1^2 |\mu_1 - \mu_0|^2_2 + v (\sqrt{\sigma_1} - \sqrt{\sigma_0})r_2^{2})$;
  • Supervised error (named Supervised in the Table) in our derived bound arises from the loss of covariate information by over-balancing $2\frac{(2+ \sigma_0 + \sigma_1)(r_1r_3)^{2}}{4} \|\mu_1-\mu_0\|^2_2$;
  • Overall TCP performance (named Overall in the Table) containing both factual and counterfactual prediction error.
• We compare the following four methods:
  • CT: adversarial balanced Causal Transformer, in which $\alpha$ are tuned with best hyper-parameter search as in the original CT paper (Melnychuk et al, ICML 2022);
  • CT w/o B: CT without adversarial balancing by setting $\alpha_{CT}=0$;
  • Mamba-CDSP (ours): Our proposed methods with de-correlation;
  • Mamba-CDSP w/o B: Vanilla Mamba model without de-correltaion by setting $\alpha_{Mamba}=0$.

• From above results, our theoretical finding of the over-balancing issue can be verified by comparing performance between CT and CT w/o B:

  • Traditional CT with balancing always has larger supervised error but smaller Wasserstein balancing error than CT w/o balancing;
  • When confounding strength is large, i.e., $\gamma=3$, the reduction of Wasserstein balancing error in CT is larger than the increase of supervised error, making the balancing approach useful;
  • When confounding strength is small, i.e., $\gamma=0,1,2$, the reduction of Wasserstein balancing error in CT is smaller than the increase of supervised error, resulting in the over-balancing issue.
  • For example, when $\gamma=2$, the balancing approach in CT reduces the Wasserstein balancing error from 1.32 to 1.23 (decrease 0.09), but increases the supervised error from 0.93 to 1.24 (increase 0.31), i.e., the following equation holds: $\|\mu_1 - \mu_0\|_{2}^{2} \leq \frac{v(\sqrt{\sigma_1} - \sqrt{\sigma_0})^2 r_2^2}{vr_1^2 - (2+\sigma_1 + \sigma_0)r_1^2r_3^2/4}.$

• Our theoretical finding of the superiority of the proposed Mamba-CDSP can be verified by comparing all 4 methods:

  • Notably, Mamba-CDSP always has a summation of the Wasserstein balancing error and supervised error compared to CT, which leads to the optimal overall performance;
  • Besides, the over-balancing phenomenon of Mamba-CDSP slightly exists in the case of $\gamma=0,1$.

I am curious whether removing CDSP would mean better preservation?

Yes, but removing CDSP would also lead to larger error when correcting the confounding bias at the same time (see our detailed discussion above).

Besides, for CT experiments on MIMIC-III, does w/o balancing simply mean keeping only the experimental results with α=0?

Yes, w/o balancing simply means setting $\alpha=0$ in the loss $(\hat \theta_Y, \hat \theta_R)=\arg \min_{\theta_Y, \theta_R} \mathcal L_{G_Y}(\theta_Y, \theta_R)+\alpha \mathcal L_\mathrm{conf}(\hat \theta_A, \theta_R)$ (see Eq. (18), Page 6, Section 4.3 in the original CT paper (Melnychuk et al, ICML 2022)).

Dear Reviewer PLDR,

As the discussion deadline approaches, we are wondering whether our responses have properly addressed your concerns? Your feedback would be extremely helpful to us. If you have further comments or questions, we hope for the opportunity to respond to them.

Many thanks,

7314 Authors

We sincerely appreciate the reviewer’s great efforts and insightful comments to improve our manuscript. We know that we are now approaching the end of the author-reviewer discussion and apologize for our late rebuttal. During the rebuttal period, we have been focusing on these beneficial suggestions from the reviewers and doing our best to add several theoretical analysis and experiments and revise our manuscript. We believe our current carefully revised manuscript can address all the reviewers’ concerns.

W1: 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.

Response: Thank you for the comment, we agree that these issues are not first proposed by this paper, but we cannot fully agree with the comment regarding novelty. To enhance technical contribution of our paper, we add extensive theoretical analysis with experimental validation as summarized below.

The novelty of this paper

• To the best of our knowledge, this is the first work that adopts Mamba to address time-varying counterfactual prediction. It should be noted that we are not simply replacing the transformer with the Mamba model.
• From a deconfounding prospective, we propose a novel de-correlation approach benefiting from Mamba for better trade-off between covariate balancing and covariate information preservation, which distinguishes from the previous studies enforcing the independence operation in each time step.
• Following the pioneer work focusing on the over-balancing issue in Huang et al. (ICML 2024), we are the first work tackling the over-balancing issue in the task of time-varying counterfactual prediction.

The added technical contribution during rebuttal

• For the over-balancing issue in TCP (W3), we added theoretical analysis on the counterfactual prediction risk bounds, revealing that the intrinsic reason for the over-balancing issue is the profit of balancing in reducing the Wasserstein balancing error term will be smaller than the extra supervised prediction error due to the loss of covariate information (see our response to W3).

• For the benefits of replace the transformer with the Mamba model by slightly tailoring the backbone (W4), we conduct more theoretical analysis showing that our CDSP method always has a tighter risk bound compared with the traditional adversarial balancing methods (see our response to W4).

• For the efficiency of replacing the Transformer with the Mamba model (W5), we added theoretical computational complexity efficiency analysis on both CDSP and traditional adversarial balancing methods, showing that the overall complexity of our proposed CDSP method will be smaller than the previous methods. We also added experiments on the traning and inference efficiency to support our theoretical results (see our response to W5).

W2: 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.

Response: Thank you for referring us two inspiring existing works and we found these works are beneficial to us! We would like to clarify that we never claim that we are the first to consider the step-by-step deconfounding in the TCP task, but have claimed that this is the pioneer Mamba model tailored to counterfactual prediction in Line 94 in our original draft.

• For the referred two existing works, we first summarize the core idea of them as follows:
  • [1] proposes to adapt the domain-invariant adversarial learning for TCP task, which enforces the adversarial loss to optimize each time-step.
  • [2] is proposed for estimating the distribution of $Y$ rather than its mean, and mainly designed towards high-dimensional counterfactual generation.
• As suggested by the reviewer, we illustrate the difference between our method and them mainly in two aspects:
  • Our method indeed performs the de-correlation operation in the last time step, while both [1] and [2] enforce the independence operation in each time step. Considering the over-balancing phenomenons, both [1] and [2] will lead to more serious deterioration of covariate information.
  • De-confounding contains many diverse techniques, e.g., matching, covariate balancing in [1], and re-weighting in [2], etc.. Different from [1] and [2], our paper proposes a novel de-correlation approach benefiting from Mamba, which leads to efficient implementation with linear time complexity.
• Benefiting from the reviewer's suggestion, we re-write the related-work section with a more thorough and complete review by adding over 15 new recent works in Section 2.

Please understand and forgive our delayed response to the updated review, as we have been working to supplement the extensive technical contributions and experiments to fully address your concerns. We would really appreciate it if you would like to kindly consider adjusting your score. We look forward to your insightful and constructive responses to further help us improve the quality of our work. Thank you!

References

[1] Bica et al, Estimating counterfactual treatment outcomes over time through adversarially balanced representations, ICLR 2020.

[2] Wu et al, Counterfactual Generative Models for Time-Varying Treatments. KDD 2024.

[3] Huang et al, An Empirical Examination of Balancing Strategy for Counterfactual Estimation on Time Series, ICML 2024.

[4] Melnychuk et al, Causal transformer for estimating counterfactual outcomes, ICML 2022.

[5] Ahmed M. Alaa and Mihaela van der Schaar, Limits of Estimating Heterogeneous Treatment Effects, ICML 2018.

[W4. This method seems to simply replace the transformer with the Mamba model by slightly tailoring the backbone, which lacks technical contribution.]

Response: We would like to emphasize that we are not simply replace the transformer with the Mamba model, but also propose a novel de-correlation approach benefiting from Mamba for more effective and efficient trade-off between deconfounding and covariate information preservation, which differentiates from the previous studies enforcing the independence operation in each time step.

To enhance the technical contribution of our paper, we conduct more theoretical analysis showing the superiority of our propose Mamba-CDSP method.

[Main Technical Conclusion] CDSP has a tighter counterfactual prediction risk bound than traditional balancing approaches.

• Recap that for previous adversarial balancing (ADB) methods, with probability at least $1-2\eta$, we have $\epsilon_{\mathrm{ECP}}^{adb}(\tilde{W}, \tilde{\Phi}) \leq 2(\frac{(2+ \sigma_0 + \sigma_1)(r_1r_3)^{2}}{4} \|\mu_1-\mu_0\|^2_2 + C - (\sqrt{\frac{1}{\eta}}( \sum_{j=1}^{2}\sqrt{\kappa_j}) + \overline{\sigma} ) ).$
• For our proposed CDSP method, with probability at least $1-2\eta$, we have $\epsilon_{\mathrm{ECP}}^{cdsp}(\tilde{W}, \tilde{\Phi}) \leq 2(\frac{(r_1r_3)^{2}}{2} \|\mu_1-\mu_0\|^2_2 + C - (\sqrt{\frac{1}{\eta}}( \sum_{j=1}^{2}\sqrt{\kappa_j})+ \overline{\sigma} ) ).$
• With the above risk bounds, by noting

we conclude that our CDSP method has a tighter bound compared with the existing ADB methods, i.e., $\epsilon_{\mathrm{ECP}}^{cdsp}(\tilde{W}, \tilde{\Phi}) \leq \epsilon_{\mathrm{ECP}}^{adb}(\tilde{W}, \tilde{\Phi})$.

• We have added the above theoretical analysis in a new Section 4.4 with detailed proofs in Appendix A.3.

[W5. 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.]

Response: Thank you for the constructive comments. In below, we first clarify a misunderstanding that the efficiency of Mamba compared to Transformer should be sequence length, instead of the sample size. Next, we analyze the computational complexity efficiency of our proposed CDSP method and traditional adversarial balancing methods from both theoretical and empirical prospectives, showing the superiority of the CDSP method. Finally, we add experiments on a new real-world time-series prediction dataset, i.e., the M5 dataset, to further validate the effectiveness of our method on the real-world datesets.

Efficiency of replacing the Transformer with the Mamba model on the extremely long sequences

• We would like to clarify that the original Mamba paper as well as its follow up works demonstrate the efficiency advantage of Mamba compared to Transformer is not sample size, but sequence length.
• In counterfactual prediction tasks, Mihaela van der Schaar et al. (ICML 2018) revealed that the confounding effect will vanish with increasing sample size in the absence of unmeasured confounders (see formal theoretical results in Section 3.2 and Figure 1 in Mihaela van der Schaar et al. (ICML 2018)). That is, in the case of large sample regime, the counterfactual prediction will reduce to the trivial temporal forecasting problem.
• We highlight that Mamba-based methods perform better than Transformer-based methods on extremely long sequences, because the computational complexity of Transformer-based methods is squared-level $\mathcal{O}(L^2)$, which prevents setting the contextual window as the whole sequence length (1000 timesteps), but only as a subsequence length (200 timesteps) in practice due to the high computational cost.
• To illustrate this point, we have added experiments on extremely long sequences using the TG simulator data with $L = 1000$ in Table 4, Page 10 in our revised paper.

• From above results, our proposed methods stably outperforms Causal Transformer at the significance level $\text{p-value}\leq 0.01$ at varying time-steps.

W3: 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?