GST-UNet: A Neural Framework for Spatiotemporal Causal Inference with Time-Varying Confounding

Strengths:

Thank you for your thorough review. We appreciate that you found the theoretical background clear and the overall task highly relevant. Below we clarify the differences from prior work and address each of your questions; we hope these points will encourage you to reconsider your score.

Re: Novelty and backbone comparison vs. prior neural G-computation

Prior neural G-computation methods assume many i.i.d. temporal trajectories, whereas GST-UNet targets a single, spatially coupled sequence where interference, spatial confounding, and data scarcity are first-class challenges. This makes both the assumptions and the required architecture fundamentally different. The table below summarises the key distinctions.

As the comparison shows, GST-UNet introduces substantial new elements—and considerable care in both design and theory—to deliver an accurate, unbiased, ready-to-use estimator for spatiotemporal causal effect estimation.

Re: Interference between units

In our grid, each spatial location is a unit; interference occurs when $A_{s',t}$​ influences $Y_{s,t}$​. The learned convolutional kernel aggregates neighbour treatments and covariates before the G-comp heads, capturing both direct and spillover effects. For the purposes of this rebuttal, we reran the synthetic study with a 1 × 1 kernel (no neighbour information). The results (RMSE ± SD) are given below for horizon $\tau=5$:

The RMSE increases when we don’t consider spatial relationships confirm the importance of modelling spatial spillovers. We will include this ablation in the revised manuscript.

Re: Baselines tailored to CAPO over time

We do not benchmark G-Transformer, CausalTransformer, G-Net, etc., because their core assumption—many i.i.d. trajectories with no cross-unit interaction—clashes with our single-trajectory, spatially coupled setting; adapting them would require redesigning their causal graphs and losses, which is beyond this submission’s scope.

We already cite [1–5] in the related-work section and Appendix A; we will add a concise comparison (table above) to the main text for clarity.

We appreciate your thoughtful feedback and hope the clarifications provided show the value of our contributions and justify a higher score.

Strengths

Thank you for your positive and encouraging review. We’re glad you appreciated how GST-UNet unites neural spatiotemporal representation learning with iterative G-computation, is grounded in causal-inference principles, and addresses real public-health and environmental-policy challenges. We respond to your questions below and hope these clarifications will prompt you to raise your score.

Re: Embedding-based time invariance assumption

Thank you for raising this important point. While Assumption 3 may appear strong, it is substantially weaker than classical stationarity assumptions. Instead of requiring the marginal distributions of $\mathbf{X}_t$​ and $\mathbf{Y}_t$ to be time-invariant, we assume that—once the history $(**H**\_{1:t}, \text{A}_t)$ is embedded into a latent representation $\phi(\mathbf{H}\_{1:t}, \text{A}_t)$—the conditional distribution of $(\text{X}\_{t+1}, \text{Y}\_{t+1})$ is invariant across time. This reflects the idea that the mechanism generating outcomes is stable once we condition on all relevant confounders, even if the raw distributions evolve over time.

This assumption is motivated by recent work in time-series causal inference (see Assumption 3 in paper), which uses learned embeddings to pool information across time and improve generalization—particularly important in our single-trajectory setting, where time-dependent parameters are hard to estimate reliably. In many applied domains, such as climate, environmental health or epidemiology, transitions conditioned on relevant covariates (e.g., prior exposure, mobility, or weather) often behave consistently, making this a reasonable and practically useful approximation.

We acknowledge that stronger empirical validation of this assumption is an important direction for future work.

Re: Comparisons with Causal Transformers and related methods

We intentionally did not include Causal Transformer [1], G-Net [2], or DeepACE [3] because their foundational assumption—many i.i.d. patient-style trajectories with no cross-unit interaction—conflicts with our single, spatially coupled lattice. Adapting them would require (i) redesigning their causal graphs to include neighbour treatments, (ii) revising their loss functions for one trajectory rather than thousands of i.i.d. ones, and (iii) adding spatial inductive bias so they can learn spill-over effects. These steps are non-trivial and beyond this submission’s scope.

We perform a part-for-part comparison between our model and these methods in the table below:

We already cite and compare these works in the related work and Appendix A; we will move a concise comparison (table above) to Appendix A for clarity.

Re: Higher-dimensional confounders

Thank you for pointing this out. We extended the synthetic study to include 5 time-varying confounders. The revised spatiotemporal DGP is:

We index covariate channels by $c \\in \\{1, \\dots, 5\\}$ and let $\bar X_t = \frac{1}{5}\\sum_{c=1}^5 X^{(c)}_{t}$ enote the channel-average field that drives treatment and outcome. Spatial convolutions are indicated with “*”.

We train our method and the corresponding baselines as in the paper and report the results below ($\tau=5$):

These results confirm GST-UNet’s robustness to increasingly strong time-varying confounding, while plain UNet variants and existing spatial models degrade.

Re: Time complexity of the GST-UNet

Analytical form. A single forward/back-prop pass scales linearly in $N_{xy}$​ (the number of spatial cells), in the horizon $\tau$ (the number of G-heads), and in the channel width $d$, leading to a runtime that scales roughly like $O(N_{xy} \tau d^2)$. Empirical runtimes (A100 GPU on NERSC’s Perlmutter cluster) for each experiment:

We hope these clarifications are helpful and further strengthen your confidence in the work.

Strengths

Thank you for your insightful review. We appreciate your recognition of GST-UNet as the first architecture to blend spatiotemporal representation learning with iterative G-computation, and of its balance of theory, simulations, and a real wildfire-health application. Below we address each of your questions and concerns, and we hope these clarifications will highlight the full strength of our contribution and justify a higher score.

Re: Domain breadth

We agree that an additional domain would strengthen the paper. In parallel with this rebuttal, we are preparing a second real-world study on how short-wave radiation affects Arctic sea-ice concentration (dataset request in progress, following the setup of [1]) and will include those results in the camera-ready—or sooner, if they finish by the end of the discussion period.

Re: Regular grid vs. graph extension

GST-UNet’s encoder is modular and thus it can be swapped for a graph–message-passing variant almost by "find-and-replace":

Moreover, Theorems 1 & 2 need no modification: they rely only on the history embedding $\phi(\mathbf{H}\_{1:t})$ satisfying Assumption 2, not on how $\phi$ is produced. Thus replacing grid convolutions with message passing adapts GST-UNet to irregular networks while leaving the G-computation heads and proofs intact. Full implementation and hyper-parameter tuning of this graph variant is promising future work, but beyond the scope of this first grid-based study.

Re: Curriculum "ease coefficient" & schedule sensitivity

As shown in Table 2 of Appendix D, we tuned the ease coefficient $e_c$​ over the grid $(1, 3, 5, 7)$ for every GST-UNet application. Across all synthetic and wildfire runs, $e_c = 3$ consistently gave the best validation RMSE. Intuitively,

• Too small an interval ($e_c = 1$) activates new heads almost every epoch, forcing them to fit highly noisy pseudo-outcomes before later heads stabilize.

• Too large an interval ($e_c \ge 5$) leaves the backbone frozen for many epochs; when it is finally unfrozen, the network overfits to late-stage heads and struggles to reconcile earlier steps, hurting long-horizon accuracy.

The middle value strikes a balance between denoising pseudo-outcomes and keeping the shared encoder adaptable throughout training.

Re: Head-count ($\tau$) vs. prediction horizon Each G-head $k$ learns the conditional outcome function

Thank you very much for the detailed response, which has addressed most of my questions and concerns. I am willing to update my score accordingly.

Strengths

Thank you for your thoughtful review. We’re glad you found the theoretical exposition clear, the methodology well illustrated, and the problem we tackle broadly relevant. We address your concerns below and hope these clarifications will prompt a reconsideration of your score.

Re: Novelty and backbone comparison vs. prior neural G-computation work

Thank you for bringing this up. In short, prior neural G-computation models tackle many independent, purely temporal trajectories, whereas GST-UNet targets a single, spatially coupled sequence in which interference, spatial confounding, and data scarcity are first-class challenges—so both the assumptions and the required architecture are fundamentally different. We include a side-by-side comparison with previous CAPO-estimation works in the table below:

As the comparison makes clear, GST-UNet introduces substantially novel elements—and considerable care in both design and theory—to deliver an accurate, unbiased, and ready-to-use estimator for spatiotemporal causal effect estimation.

Finally, we do not benchmark against G-Transformer, CausalTransformer, or other temporal CAPO baselines because their core assumption—many i.i.d. trajectories with no cross-unit interaction—fundamentally clashes with our single-trajectory, spatially coupled setting; adapting those methods would entail redesigning their causal graphs and losses, which lies beyond the scope of this submission.

While we already mention works [1-6] in the related work section and Appendix A, we will add a concise comparison like this table to the final version of the paper.

Re: Interference

In our grid, each spatial location is a unit; interference arises when $A_{s',t}$​ influences $Y_{s,t}$​. The learned convolutional kernel aggregates neighbour treatments/covariates before the G-comp heads, capturing both direct and spillover effects. For further validation, we reran the synthetic study with a 1 × 1 kernel (no neighbour information). The results (RMSE ± SD) are given below for horizon $\tau = 5$:

The RMSE increases when we don’t consider spatial relationships confirm the importance of modelling spatial spillovers. We will include this ablation in the revised manuscript.

We hope our responses have addressed your key concerns and would be grateful if you would consider a higher score in light of these clarifications.

Thank you for the follow-up and for raising these points, we appreciate the chance to clarify. Below, we expand on your remaining questions regarding novelty, methodological contributions, and theoretical results.

Re: Why the single-trajectory, spatial setting forces methodological changes (beyond simply swapping the backbone)

Main differences, side-by-side:

Key take-away: in the single-trajectory setting we must (i) reshape the raw data, (ii) couple losses across space, (iii) stabilise training with a curriculum, and (iv) prove identification/consistency under time-invariant representations—changes that go well beyond swapping a backbone.

Re: Why Causal Transformer-style baselines are not applicable here

• Data volume. Causal Transformer consumes $N_\text{traj} (T - \tau)$ independent tokens; we have only $T - \tau$ correlated tokens.

• Spatial dimension. Its input is a length-$\tau$ vector per unit; extending it to $(H \times W)$ grids would require
   (i) stacking 4,096 separate “units” for a $64 \times 64$ grid,
   (ii) giving the model a bespoke exposure kernel to mix them, and
   (iii) rewriting the loss to couple those units inside each prefix.

• Interference. Without these modifications, the baseline would treat every grid cell’s history independently and cannot learn spill-over effects. Thus, any performance gap would reflect a mismatch of assumptions, not algorithmic merit.

These extensions are non-trivial and out of scope for our current work.

Re: What's novel about the identification result?

While our G-formula is structurally similar to the classical i.i.d. case, its validity in our setting relies on Assumption 2, which ensures that conditioning on a learned embedding renders prefixes quasi-i.i.d.. This is what allows us to treat each prefix as a valid unit for regression, even though they come from a single dependent trajectory with interference (loosely analogous to martingale CLTs, where conditional independence is required given past information).

Please let us know if there is anything else we can clarify or expand on during the discussion period to help support your evaluation.

Strengths

Thank you for your detailed and constructive review and for highlighting our work addressing a key gap in spatiotemporal causal inference and GST-UNet’s thoughtful design, training, and balance of theory with simulations. We respond to your points below and hope these clarifications will show the full strength of our work and earn a higher score.

Re w1: "Data-scarce regimes"

Thank you for raising this point. GST-UNet operates in an intrinsically data-scarce setting:

• Single spatiotemporal chain. We observe only one evolving lattice $(A_{s,t},X_{s,t},Y_{s,t})_{t=1}^T$​. After choosing a prediction horizon $\tau$, the effective sample size is at most $T-\tau$ quasi-i.i.d. prefixes. To get a sense of scale differences, non-spatial CAPO methods — usually with a Transformer backbone — can use arbitrarily many (usually tens of thousands) independent time-series (e.g. data from different patients). Space–time interactions preclude treating each grid cell as an i.i.d. series.

• Inductive-bias-driven efficiency. With so few samples, performance hinges on a backbone whose built-in biases match the physics:

  • The ConvLSTM-U-Net shares weights across location and scale, enforcing locality and translation invariance—exactly the properties needed for atmospheric fields for example. Consequently, U-Net-style backbones have become the de facto choice in earth science modelling and related spatial sciences (e.g., [1, 2, 3, 4, 5] below).

  • Architectures with weaker locality priors (e.g., vanilla visual transformers with global attention [3]) require substantially more data or pre-training to learn comparable patterns, as documented in the original ViT literature; preliminary tests in our setting confirm this trend. We are running a ViT ablation and will include the results in the final version.

• Backbone agnostic in principle. Section 4 states that any encoder producing embeddings satisfying Assumption 2 will inherit Theorems 1 & 2. U-Net is simply a strong default for limited data; practitioners can swap in graph-convs or hybrid CNN-ViT blocks as larger datasets become available.

We will add a short paragraph in Section 4 summarising the argument above and include the ViT ablation to make the data-scarce claim explicit.

Re w2: Theorem 2

Thank you for flagging this. Theorem 2 was placed in the appendix purely for space, not because it is secondary. In the camera-ready we will: 1) Move Theorem 2 (consistency) into the main text alongside Theorem 1, and 2) Shift some verbose setup, architectural exposition and training-schedule details to the appendix to stay within the page limit. We note that Theorem 2 is, to our knowledge, the first consistency result for CAPO estimation with iterative G-computation, so we agree that highlighting it in the main body will make the paper’s theoretical contribution clearer.

Re w3: Framework vs tool

Our goal is both to close a theory gap and to deliver a practical recipe—but the two pieces are tightly coupled, not competing claims.

• Theory gap. Outside of STCI-Net [2] (which ignores time-varying confounding and offers no identification guarantees) there is no work that tells practitioners when CAPOs can be recovered from a single, spatially coupled chain that is common in Earth science applications, epidemiology, etc. Assumption 2 (representation-based time invariance), Theorem 1 (identification) and Theorem 2 (consistency) fill that gap by stating minimal conditions under which unbiased estimation is even possible. We note that Assumption 2 subsumes (and thus relaxes) the strict stationarity/mixing conditions usually imposed in single-series time-series work, and Theorem 2 is, to our knowledge, the first consistency result for CAPO‐over‐time estimation.

• Methodology gap. Once those conditions are known, a practitioner still needs a concrete procedure that (i) respects the single-trajectory sample constraint, (ii) reliably accounts for spatial effects (interference, confounding), and (iii) adjusts for time-varying confounding via iterative G-computation. GST-UNet is one ready-to-use instantiation of the general framework: ConvLSTM-UNet + curriculum-stabilised G-computation heads. We emphasise in Section 4 that any encoder satisfying Assumption 2 can replace the ConvLSTM-UNet—our code merely provides a non-trivial, field-tested starting point.

Thus the paper is positioned as theory-guided tool: theory tells the practitioner what must hold and how iterative G-computation should be structured; GST-UNet shows one concrete way to achieve it today, while providing a blueprint for future backbones (e.g., graph U-Net and graph convs for irregular meshes).

Re w4 & m1: Presentation and Experiments

Point well taken re the presentation. Thus, we propose to reorganize the camera-ready version as follows:

• Promote Theorem 2 (consistency) to the main text and move low-level implementation details to the appendix

• Add a concise roadmap that tags each element as a core principle or a replaceable implementation/engineering choice, preventing concept–detail jumps.

While the current paper already ablates curriculum training and the attention mechanism, we commit to strengthen the evidence for the current implementation by:

• adding a TinyViT [7] backbone ablation to test robustness to encoder choice, and

• including a second real-world study on the impact of short-wave radiation on Arctic sea-ice concentration (dataset request under way, following the setup of [2]).

We believe that these additions will reinforce the “ready-to-use tool” contribution that the reviewer expects. We hope the streamlined presentation and the new empirical evidence would address your concerns and make the dual contribution—rigorous foundation and practitioner-friendly recipe—clearer.

Re m2 & Q2: Benchmarks & Network causal inference

Our synthetic study emulates pollutant advection–diffusion on a 2-D lattice:

• Physical realism. Each covariate field $X_{s,t}$ evolves via a discretized advection–diffusion equation with an asymmetric kernel $K_X$ that mimics wind-driven transport; treatments $A_{s,t}$​ inject additional emissions that propagate through the same kernel. This produces both temporal carry-over and spatial spillover—the two phenomena GST-UNet is designed to handle.

• Time-varying confounding. Past emissions alter local $tX_{s,t}$​, which in turn influence future $A_{s,t}$ choices, mirroring real regulatory feedback loops.

• Why this DGP? It is (i) analytically tractable, (ii) tunable for different spillover radii, and (iii) aligned with our real-world wildfire application, allowing direct “toy-to-real” translation. Thus, while simplified, the setup is grounded in a real physical process.

At present, no public semi-synthetic benchmark combines both interference and time-varying confounding. Prior spatial works such as Huang et al. [8] model only static spatial confounding and explicitly list dynamic extensions as future work; contemporaneous CAPO papers [2, 9, 10] likewise rely on fully synthetic or domain-specific real data. Constructing a comprehensive benchmark is a valuable but separate research effort that we plan to undertake in the future.

Additionally, most network-causal methods (see e.g. [11-13]) assume static graphs with a fixed, user-specified exposure mapping (e.g., “fraction of treated neighbours”) and focus on bias from interference. Our setting differs in three ways:

• Dynamic graph: the spatial grid remains fixed but covariates and treatments evolve each time step, creating lagged carry-over that standard network estimators do not usually model.

• Learned exposure function: GST-UNet infers the spillover kernel directly from data rather than prespecifying it.

• Spatial confounding: network methods typically ignore latent spatial fields (e.g., weather) that confound both treatment and outcome.

Adapting a static-network estimator would therefore require (i) extending it to temporal graphs, (ii) specifying or learning a lagged exposure mapping, and (iii) modelling spatial confounders—all non-trivial modifications outside the scope of this paper. We believe our framework is the first to integrate all three elements within a single, consistent estimator.

With respect to irregular network data, we note that the backbone is modular: swapping the grid convolution for a graph-convolution layer would let GST-UNet run on irregular networks; implementing and tuning that variant is promising future work but outside this first grid-based study.

We hope our responses have addressed your key concerns and would be grateful if you would consider a higher score in light of these clarifications.

[1] M. Tec et al. Weather2vec: Representation learning for causal inference with non-local confounding. AAAI, 2023.

[2] S. Ali et al. Estimating direct and indirect causal effects of spatiotemporal interventions with spatial interference. ECML PKDD, 2024.

[3] G. Ayzel et al. RainNet v1.0: A CNN for radar-based precipitation nowcasting. Geosci. Model Dev., 2020.

[4] J. C. Fernández et al. Deep coastal sea forecasting using UNet-based models. Knowl.-Based Syst., 2022.

[5] J. Ring et al. Prediction of sea surface temperature using U-Net. Remote Sens., 2024.

[6] A. Dosovitskiy et al. An image is worth 16x16 words: Transformers for image recognition. arXiv:2010.11929, 2020.

[7] K. Wu et al. TinyViT: Fast pretraining distillation for small vision transformers. ECCV, 2022.

[8] M. Tec et al. SPACE: The spatial confounding environment. ICLR, 2024.

[9] G. Papadogeorgou et al. Causal inference with spatio-temporal data. J. R. Stat. Soc. B, 2022.