Transformer-Based Spatial-Temporal Counterfactual Outcomes Estimation

Response to reviewer Uw2a

We are glad the reviewer found our paper well-written. We would like to respond to each detailed point individually.

1. I failed to be convinced that there is significnat novelty on the theoretical front -- the proof in Appendix B2 is thorough, but Proposition 3 seems to be more like a routine proof. To be fair, the authors did not claim theoretical contributions, so this is fine by me. 

We appreciate that the reviewer found our proofs thorough. Now, we elaborate on our key contributions. One of our key contributions is leveraging deep learning to address the challenge of counterfactual outcome estimation for spatial-temporal data. Previous baselines are unable to achieve this. Specifically, our proposed method enables the estimation of propensity scores for spatial-temporal data. This provides a feasible and efficient technical pathway for causal inference in high-dimensional spatial-temporal settings. Besides, our theoretical analysis provides a guarantee for the proposed deep learning method.

---

2. However, I think the lack of validation for the non-synthetic studies makes the story less convincing. I understand that there is non "true answer" for the non-synthetic studies, but the authors could perhaps perform sensitivity analysis or convergence analysis of their estimators.

We appreciate the reviewer's suggestions on sensitivity analysis and convergence analysis. Following your advice, we provide the theoretical sensitivity analysis in the "sensitivity.pdf" file in this anonymous URL.

With regard to the convergence analysis, Proposition 3 demonstrates that as time progresses, our estimator converges to the estimands (ground truth).

---

3. I think it is helpful that the authors discuss the computation cost but a comparison of computation cost with the baseline might shed more light on the comparison between the proposed method.

We appreciate the suggestion to compare the computational cost with the baselines. All baselines were also run on an NVIDIA RTX 4090 GPU with an Intel Core i7-14700KF processor. For detailed information, please refer to the table below.

Note that these baselines were not designed for spatial-temporal data and cannot process such data. Therefore, we convert all spatial-temporal data into scalar values for processing by these baselines (Appendix K.1). As a result, although our method does not significantly outperform these baselines in computation time, it addresses a problem that they cannot solve.

---

4. I think, despite this paper is generally well-written, the constribution seems to be a bit weak -- there is no significnat contributions on the theory front. However, the empirical studies and analyses are related limited compared with what I would expect from an application paper.

We appreciate that the reviewer found our work well-written. Our main contribution lies in addressing the challenge of counterfactual outcome estimation for spatial-temporal data, which previous baselines are unable to solve. We provide a feasible and efficient technical pathway for causal inference in high-dimensional spatial-temporal data. The theoretical section provides the foundation for the proposed method.

---

5. In the figures/tables of empirical studies, perhaps adds error bars and std/confidence bands.

We appreciate the suggestion to include error bars and confidence intervals. However, in the additional experiments, we trained the neural networks 25 times with different training data, and the results were consistently superior. Therefore, we believe the additional experiments demonstrate the robustness of our method.

---

6. Did you try to relax the Poisson assumption and see how well this method holds up in the synthetic setup?

We acknowledge that the Poisson assumption may be somewhat restrictive. In this work, we use the spatial Poisson point process as a fundamental hypothesis for modeling spatial-temporal data. Therefore, if the Poisson assumption is violated, our theoretical framework would need to be restructured, and its performance may be compromised. In statistics, modeling spatial data distributions without assuming a specific distribution remains an open problem, and the Poisson assumption is widely used to describe spatial events [1-2].

[1] Cressie, N., & Wikle, C. K. (2011). Statistics for spatio-temporal data. John Wiley & Sons.

[2] Cressie, N. (2015). Statistics for spatial data. John Wiley & Sons.

Response to reviewer pyND

We are glad the reviewer found the theoretical framework rigorous. We would like to respond to each detailed point individually.

1. The theoretical evidence is adundant, whereas the empirical results lack sufficient ablation to support the claims such as the improvements from the CNN based propensity score estimation.

We really appreciate that the reviewer found the theoretical evidence abundant. Now we explain why removing the propensity score estimation is inappropriate in our setting.

In our work, the CNN-based propensity score estimation serves as a practical implementation of the propensity score in our estimator framework (Sec. 4.4). Therefore, if we remove it, we effectively eliminate the propensity score from the estimator framework. As a result, the estimator would be theoretically incomplete and unable to perform its intended function properly.

---

2. Methods do make sense, but the evaluation criteria is lacking.

We are glad the reviewer found the methods reasonable. Now we justify our choice of evaluation criteria. In the simulation experiments, we use the relative error rate as a direct measure of the estimator’s accuracy. In real-world data experiments, since the ground truth is unavailable, we assess our approach by comparing our findings with those of existing studies.

---

3. Spatio-temporal causal inference is very important and broadly applicable.

We are glad the reviewer found the studied problem important.

---

4. I would encourage including more references to the temporal counterfactual prediction literature as it is an active research area, including (but not limited to).

We really appreciate the reviewer's suggestions on the references. Following your advice, we have cited and discussed these excellent works in our paper.

Temporal Counterfactual Prediction Temporal counterfactual prediction refers to performing counterfactual outcome prediction under time-varying settings. Seedat, N. et al. [1] propose TE-CDE, a neural-controlled differential equation approach for counterfactual outcome estimation in irregularly sampled time-series data. Wu, S. et al. [2] introduce a conditional generative framework for counterfactual outcome estimation under time-varying treatments, addressing challenges in high-dimensional outcomes and distribution mismatch. El Bouchattaoui, M. et al. [3] propose an RNN-based approach for counterfactual regression over time, focusing on long-term treatment effect estimation.

[1] Seedat, N. et al. Continuous-time modeling of counterfactual outcomes using neural controlled differential equations.

[2] Wu, S. et al. Counterfactual generative models for time-varying treatments.

[3] El Bouchattaoui, M. et al. Causal contrastive learning for counterfactual regression over time.

---

5. The details of the transformer architecture seems to be missing: what positional encoding are you using? what tokenization scheme is used?

We employ the Transformer to capture spatial information within a single time step. As a result, we do not apply fixed positional encoding.

For the tokenization scheme, the Transformer processes discrete data points as input. Therefore, we treat each discrete coordinate as a token and map it to a high-dimensional embedding space.

---

6. More ablation is needed to identify the contribution of each innovation, including the CNN-based propensity score estimator.

We appreciate the suggestion of the ablation of the CNN-based estimator. However, the ablation of the CNN-based estimator is inappropriate in our setting, for detailed reasons, please refer to response 1.

---

7. It is unclear if the performance gain come from model backbone (i.e. transformer) or the proposed spatio-temporal causal inference framework. One way is to keep the model backbone simple for similar to other baselines, and adopt several ablation studies to demonstrate the gains brought by the two major innovations.

Compared to previous baselines, our work primarily tackles a problem that these methods cannot handle. This is because the frameworks and implementations of these baselines are not applicable to spatial-temporal data. Therefore, we attribute the superior performance of our method to both its theoretical framework and model implementation.

We appreciate the reviewer's suggestions on the ablation studies. However, it is not appropriate to maintain the same model backbone as other baselines. For example, we employ a CNN to process the high-dimensional tensor. If we replace the CNN with the simple logistic regression or linear regression used in MSMs, the propensity score in our work would not be estimable.

Response to reviewer p6Zj

We greatly appreciate the reviewer's comments and suggestions. We would like to respond to each detailed point individually.

1. The estimator estimates the expected number of outcomes in a specific region at a specific time. This estimator (Eq. 3) seems to be the main contribution of the paper, but it is unclear why this is the quantity of interest and how the estimator is derived.

Now we explain why the expected number of outcomes in a specific region at a specific time is the quantity of interest. These estimands provide spatial-temporal decision-making evidence for policy makers. For example, during epidemic control, policy makers can evaluate the average number of infections in specific areas at specific times under different isolation strategies, thereby optimizing the allocation of isolation resources.

Next, we briefly introduce how the estimator is derived. The estimator is derived based on the Inverse Probability Weighting (IPW) strategy in Marginal Structural Models (MSMs) [1]. Specifically, the $\lambda_{Y_t^{ob}(z_{\leq t})}(s)$ in Eq.3 is the intensity function of outcomes in time $t$ and location $s$, and $\prod_{j=t-M+1}^{t} \frac{p_{h_j}(z_j)}{e_j(z_j)}$ represents the weights similar to those in the MSMs. In summary, we adopt the weighting strategy of MSMs to the studied spatial-temporal setting.

[1] Robins J M, Hernan M A, Brumback B. Marginal structural models and causal inference in epidemiology[J]. Epidemiology, 2000, 11(5): 550-560.

---

2. The proofs seem to hold, but it seems generally unclear how to interpret these results and why they are important. Prop. 3 seems to be the most important as it seems to indicate that the estimator properly converges to the intended value, but there seems to be some subtleties in the result that could be elaborated.

Prop. 1 and Prop. 2 are two essential properties of the propensity score. Specifically, these propositions demonstrate that the propensity score can make the estimation more accurate and efficient. Therefore, by proving Prop. 1 and Prop. 2, we validate that the propensity score used for spatial-temporal data is "reasonable" and can assist in counterfactual outcome estimation.

Prop. 3 establishes the theoretical reliability of the estimator. It demonstrates that the expected value of the proposed estimator equals the estimands (i.e., unbiasedness) and that the variance of the estimation error converges to a finite value.

---

3. I appreciate that the assumptions are clearly highlighted.

We are glad the highlighted assumptions are helpful.

---

4. I think that the clarity of the paper is its greatest weakness. It would be much clearer if the paper walked through a running example, where each term of the estimand is clearly highlighted.

We appreciate the reviewer's suggestion on the running example. Following your advice, we provide an example that illustrates the estimands.

For simplicity, consider the case of $t=8$ and $M=3$. Then the estimands

$$N^\omega_8(F_H) = \int_{Z^3} |S_{Y_8^{ob}(z_{\leq 8}(F_H))} \cap \omega| dF_H(z_{\left[ \text{6,8} \right]}),$$

represent the expected number of outcomes in $t=8$ and region $\omega$ under distribution $F_H$. Next, we employ a table to elaborate on each term of the estimands.

---

5. Why is the dimensionality reduction in Eq. 5 possible? Are we not losing information?

In this work, we mainly estimate the outcome quantities. Therefore, we believe it is possible to reduce the studied variables to their quantity. We acknowledge that this strategy may result in some information loss; however, it is useful for estimating the number of outcome events. Moreover, this strategy has its own theoretical significance in terms of how to reduce the dimensionality for complicated problems.