GST-UNet: A Neural Framework for Spatiotemporal Causal Inference with Time-Varying Confounding
Authors: Miruna Oprescu
David Keetae Park
Xihaier Luo
Shinjae Yoo
Nathan Kallus
Venue: NeurIPS 2025 poster ()
Keywords: spatiotemporal causal inference
time-varying confounding
G-computation
observational studies
neural networks
treatment effects
observational data
Primary area: probabilistic_methods
TL;DR: We introduce GST-UNet, a neural framework for valid causal inference from spatiotemporal observational data with time-varying confounding and complex spatiotemporal dependencies.

Abstract: Estimating causal effects from spatiotemporal observational data is essential in public health, environmental science, and policy evaluation, where randomized experiments are often infeasible. Existing approaches, however, either rely on strong structural assumptions or fail to handle key challenges such as interference, spatial confounding, temporal carryover, and *time-varying confounding*—where covariates are influenced by past treatments and, in turn, affect future ones. We introduce the **GST-UNet** (**G**-computation **S**patio-**T**emporal **UNet**), a theoretically grounded neural framework that combines a U-Net-based spatiotemporal encoder with regression-based iterative G-computation to estimate location-specific potential outcomes under complex intervention sequences. GST-UNet explicitly adjusts for time-varying confounders and captures non-linear spatial and temporal dependencies, enabling valid causal inference from a *single* observed trajectory in data-scarce settings. We validate its effectiveness in synthetic experiments and in a real-world analysis of wildfire smoke exposure and respiratory hospitalizations during the 2018 California Camp Fire. Together, these results position GST-UNet as a **principled and ready-to-use framework** for spatiotemporal causal inference, advancing reliable estimation in policy-relevant and scientific domains.

GST-UNet: A Neural Framework for Spatiotemporal Causal Inference with Time- Time-Varyinn Varying Confounding
Miruna Oprescu Cornell University
David K Park
Xihaier Luo Brookhaven National Laboratory xluo@bnl.go
Brookhaven National Laboratory
am078@cornell.cm edu
dpark1@bn1.. gov
Shinjae Yoo
Nathan Kallus
Brookhaven National Laboratory
Cornell University & Netflix
ajyoo@bal.gov
kallus@cornell.em
Abstract
Estimating causal effects from spatiotemporal observational data is essential in public health, environmental science, and policy evaluation, where randomized experiments are often infeasible. Existing approaches, however, either rely on strong structural assumptions or fail to handle key challenges such as interference, spatial confounding, temporal carryover, and time-varying confounding-wheee covariates are influenced by past treatments and, in turn, affect future ones. We introduce the GST-UNet (G-computation Spatio-Tempone Temporal UNet), theoretically grounded neural framework that combines a J.Neet-aase spatiotemporal encoder with regression-based -based iterative G-computation to estimate location-specifie po- tential outcomes under complex intervention sequences. GST-UNet explicitly adjusts for time-varye varying confounders and captures non-linear spatial and temporal dependencies, enabling valid causal inference from a single observed trajectory data-scarce settings. We validate its effectiveness in synthetic experiments and in a real-world analysis of wildfire smoke exposure and respiratory hospitalizations during the 2018 California Camp Fire. Together, these results position GST-UNet as principled and ready-to-use framework for spatiotemporal causal inference, advancing reliable estimation in policy-relevant and scientific domains.
Introduction
Environmental hazards, public health interventions, and socio-economic policies often require understanding complex cause-and-effee relationships across space and time [30] instance, evaluating the health impacts of air quality regulations requires assessing how interventions influence both immediate outcomes and downstream effects across regions. Such applications demand robust tools for estimating causal effects from observational spatiotemporal data.
For
However, causal inference in spatiotemporal settings poses unique challenges. Outcomes are influ- enced not only by local covariates and interventions but also by those of neighboring regions (spatial confounding and interference). Effects may persist and accumulate over time (temporal carryover), and covariates often evolve in response past interventions while simultaneously affecting future ones dime-varying confounding). For example, air quality regulations are often implemented in reaction to recent pollution levels and hospitalizations, which themselves shape future exposures and health outcomes- crrcctng-creattnn feedback loops that violate standard independence assumptions. These complexities induce bias in naive estimators and especially challenging in single-trajectory settings, where replication across units or time is infeasible.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Existing approaches offer limited solutions: classical methods rely on rigid structural assumptions or user-defined exposure mappings. while recent neural models emphasize predictive accuracy over causal identification. Many assume independent time series or model only spatial correlations, leaving gap methods that can jointly address interference, temporal dependencies, and evolving confounding within principled causal framework (see Section 2
To bridge this gap, we introduce GST-UNet (G-computation Spatio-Tempora UNet), a theoretically grounded neural framework for estimating location-specific potential outcomes in spatiotemporal settings with Lme-varying confounding. GST-UNet builds on formal identification and consistency results derived under representation- based inverinvaeinnce assumption, showing how causal effects can be recovered from a single observed trajectory. We then instantiate this theory in practical neural architecture: a U-Net encoder with ConvLSTM and attention modules coupled to an iterative G-computation procedure that performs recursive causal adjustment over time. To ensure stable estimation over long horizons, we design a curriculum- based training strategy that gradually refines recursive peeudo-uutcomes enabling effective learning even in data-scarce regimes. Unlike existing approaches, GST-UNet requires no user-specified structural models and can be directly deployed in real-world spatiotemporal applications.
Our contributions are threefold: (1) We develop the first unified framework that couples theoretical identification and consistency guarantees with an end-to-end neural implementation for spatiotemporal causal inference; (2) We demonstrate through controlled simulations that GST-UNet robustly handles interference, temporal carryover, and time-varying confounding; and (3) We illustrate its practical value via real-world analysis of wildfire smoke exposure and respiratory hospitalizations during the 2018 California Camp Fire.
In summary, GST-UNet provides a principled and ready-to-use framework for causal inference from spatiotemporal data, combining formal guarantees with flexible neural implementation. By abstracting away model-specific assumptions, GST-UNet makes spatiotemporal causal estimation both accessible and reliable for applied scientific and policy domains.
2 Related Work
We summarize the most relevant prior work here, with more detailed discussion in Appendix A
Classical Spatiotemporal Causal Inference. Early approaches (e.g., spatial econometrics [2] lufference-in-diiee -n-differences (20), synthetic controls (4) rely on strong assumptions such as parallel trends and no interference. More recent methods incorporate time-varying confounding using inverse propensity weighting (IPW) and marginal structural models [31 49) but cannot address interference
unless via user-specified exposure mappings or hyper-local assumptions [ 44 48]. As noted by Zhou et al. [49], the literature remains sparse, particularly in settings with rich feedback dynamics.
Machine Learning for Spatiotemporal Modeling. Deep learning models for prediction-e.g. CNNs and RNNs [40 47), graph-based methods [25] 46). and video transformers [6/p7]ccapture complex spatial-temporal patterns but do not incorporate causal adjustments, and thus cannot estimate counterfactuals or adjust for time-varying confounders.
Time Series Causal Inference. Causal methods for longitudinal data include marginal structural models [36], iterative G-computation 335 and recent ML-based extensions using recurrent networks,
transformers, meta-learners [6] [8
24 28 39). However, these assume access to independent time series (e.g. across patients) and cannot model cross-unit interactions in spatiotemporal settings.
Neural-Based Spatiotemporal Causal Inference. Tec et al. 42 propose UNet-based model that adjusts for non-local spatial confounding but focuses on static exposures and does not address interference or time-varying effects. Most similar to our work. ( presents climate-focused model that shares certain architectural similarities but emphasizes prediction rather than causal adjustment, leaving causal identification under time-varying confounding largely unaddressed.
Positioning of Our Work. Our work bridges these threads by uniting a theoretically grounded G-computation framework with a neural architecture for spatiotemporal data. Unlike prior time-series methods that assume independent units or spatial models that overlook confounding feedback, GST- UNet is the first end-to-end approach that (i) establishes identification and consistency under explicit assumptions for a single spatiotemporal trajectory, and (ii) implements this theory in a practical neural model capable of handling interference, spatial confounding, and time- time-varying dynamics.

---

Observational Data
Interventional Data
H
Vazez
Auysa
Vazzes
Macto
Hyckt
Direct effects
Temporal canryover
Spatial confounding
Interference
Time-varying confouudingg
Unobserved
Figure I: Observational data (left) versus interventional data (right) for a horizon 2 across multi- ple locations ( Under the intervention (right), treatments are set independently of confounders, and the full history is not observed for the entire horizon.
3 Problem Formulation
Spatiotemporal Data. We model observed data as random variables on discrete spatial domain
represented by an Nx x Ny lattice: S = (2. j) [Nxl.j [yy], where [N] = {1.....0 N} denotes the index set. Time indexed by E [T]. each spatial location observe tuple (X,,t. As,t, Ys,t), where As.: E [0,1] represents a binary treatment (or interven- tion), Ys,: E IR is continuous outcome of interest, and X.,: E Rdx is a vector of time-varying covariates (e.g. local weather conditions, pollution levels, or socioeconomic indicators). Additionally, each location is associated with static features V, Rdd (e.g. geographical characteristics and socioeconomic indicators). While we focus on binary interventions for clarity, the methods gener- alize to more complex treatments. Conceptually, each variable forms 3D spatiotemporal tensor of size T Nx Ny, though in practice, observations may be incomplete. Missing data can be accommodated using masking techniques during downstream modeling.
(i.j) at time t. we
To streamline notation, we use boldface symbols for random variables defined over the entire spatial domain. For U E {X, A, }}, let U, denote its value at time t, and let U (U..... U(+r) denote its value over time interval. For specific location s, we write US (Ust..... Us,++) ( The history up to time is denoted by H (X1:, A YI:, V) for the entire spatial domain and H. (X.,,::, 1 1. V.) for specific location s. Specific instantiations of these random variables are denoted using lowercase letters (e.g.. (,,a,y, h)).
Quantities of Interest. Our primary goal is estimate location-specife Conditional Average Potential Outcomes (CAPOs) for a sequence of future spatiotemporal interventions, conditioned on observed history. Our approach builds on Rubin's potential outcomes framework [35] 36 38). which we extend to accommodate spatiotemporal settings. More concretely, we consider a future time horizon of length 7 2 1 and a predetermined interventional sequence a applied across the spatial domain starting at time f. Our goal is to estimate the potential outcomes at time + denoted as (a particular, we aim to compute:
E|Y 147 [am=47-1] H h
which represents the CAPOs at time + J under the given treatment sequence. Given two different interventional sequences am and a related secondary goal is to estimate the location specific Conditional Average Treatment Effect (CATE), given by:
E|Y (+r [acttr1]] Y H1 = hja]
Although we focus primarily on CAPOs, CATEs and other effect measures can be derived similarly.
Prefix Data in a Single Spatiotemporal Chain. The conditional expectations defining the CAPOs in Eq. (I) cannot be directly estimated from a single observed spatiotemporal realization, since the empirical averages would contain only one sample of each future outcome Y t--r [a):t=r-1]. To obtain a workable regression-based estimator, we therefore reorganize the single observed trajectory into overlapping "prefixes" of varying lengths. For each fl...., T), we define
P. = (X AL.-T. Y Y1s+r. V),
which represents the observed history up to time + 7 along with all covariates, treatments, and outcomes. When > т, this construction yields 7 segments that partially overlap in time, providing additional training samples in this intrinsically data-scarce, single-chain setting.
3

---

However, these prefixes are not independent: successive segments share overlapping histories, so standard i.i.d. assumptions do not apply. the next section, we introduce conditions under which these prefixes can be treated as conditionally exchangeable given an appropriate learned embedding. This enables regression-basee estimation of CAPOs by pooling information across time without violating the dependence structure of the original process.
4 Identification and Estimation of CAPOs in Spatiotemporal Settings
Identification of CAPOs from observational data relies on standard causal inference assumptions. In setting, these must be complemented by additional structure to handle the fact that observe only a single spatiotemporal trajectory. Building on the prefix construction introduced above, we impose conditions that render these overlapping segments conditionally exchangeable, enabling principled pooling of information across time.
Assumption 1 (Causal Inference Assumptions). We assume: (Consistency) Y 14 Bra47-1l
whenever the observed sequence of treatments A -1 satisfies A 71: (Positivity) P(Ag. = aa,tt Hit = h1:)) >> 0 for any as,1 E {0,1) and feasible realization of history hia
(Sequential Unconfoundedness) Y ( 1 A, Hist. Vat+1.T E {0, 100,13"-- i.c. at each time step t, the treatment assignment is independent of future potential outcomes.
Assumption 2 (Representation Baass Time Invariance). There exists function (or embedding) 6: x A + 2 S TR that maps (H At) to a finteedimenssonal representation such that once we condition on Formally, for any t,t' E (1....;T] T) and E 2, we have:
6(H1, At). the distribution (X++1, Y++1) does not explicitly depend on t.
p(X+11, Y44 ((H12, Ar)
(X Ye41 (HH Ar) z).
Assumption is a standard set of requirements in longitudinal causal inference settings (c.g.. 7 18
24 28 35 36). Assumption i specific to the single-time series setting, where pooling information across time is essential to enable estimation. We note that the single time-series setting frequently arises in causal inference, where assumptions such as stationarity or strict time homogeneity enable consistent estimation 88, 31 49 In contrast, our representation beeeeseenttio-bbeeed time invariance is weaker: rather than requiring Xt. Y, themselves have time-invariant distribution, we only assume that, once the history is summarized by o ((H1: At), the transition to (X++1, Y++1) follows single shared mechanism. This approach aligns with modern time-series causal inference that learn time-
invariant latent embeddings to pool information across time steps [18 24 26), thus leveraging more data for a single, stable representation rather than time-dependent parameters.
Under Assumption 2
conditioning on 6(H As) removes explicit dependence on f, such that 6 (Hir, At)l
Ep Y 1+7
represents shared conditional expectation across all prefix segments. In this view, t indexes the segment's 's position rather than a distinct distribution. Pooling over thus yields T - approximately exchangeable segments from a single trajectory, enabling regression- bassed estimation of future outcomes from embedded histories.
4.1 Identification via Representation-Basee G-Computation
Given P7, we next show how to identify CAPOs from observational data. For horizons 2 future covariates and outcomes (i.e. X -1++++))) can influence subsequent treatments, inducing vime-varying confounding [3. Such feedback violates standard condition-on-histore adjustments and leads to biased estimates. Figure [| jillustrates these dependencies by contrasting observational data (left) and hypothetical interventions (right) for
By contrast, when 7 conditioning on H is sufficient under standard assumptions, as no future confounders intervene between A, and Y 1+ Formally, the following naive identification fails to hold for 7 >
t-r E Y e+ (1 H hial # Y H
h Atttr-1 a
To correct this bias, we adapt regression-busee iterutive G-computation 3 35 to the spatiotemporal setting, yielding a principled adjustment procedure for evolving confounders and valid CAPO estimation. We formalize this connection in the following result:

---

Theorem 1 (Identification with G-Computation). Assume that Assumption 7 and Assumption 2 hold. Further, let H 1a4k (X1++ks 1 a 11++k) denote the history where observed treatments from time onward are replaced by att+k-1- Define recursively:
QH A = Ep|Y [ Actr-1)]
Qr-1 -H
Q2 (H ac+r-1) ((H Artr-2)]
Q1 A1) = Ep [Q2 (H 1++1 a++1) (HH,,, Ar] ...
Then E (1t--aau---1 H hi:d Q1i at).
We provide proof of Theorem I in Appendix B This result naturally motivates a recursive regression approach for spatiotemporal CAPO estimation, fitting each Qx(+) in reverse order and substituting interventional treatments where required.
4.2 Estimation via Iterative G-Computation
While Theorem motivates a recursive regression algorithm for each Qk only Q can be directly estimated from the prefix data. At the next step, Qr-1 depends Q, at-r- the observed treatments At are replaced by a such substituted outcomes are not observed in the prefix data. Therefore, for k < т. we propose procedure where generate pseudo-oulcomees by predicting with the previously learned QK+1- Going forward, we use F denote any quantity estimated from data. Formally, let embedding satisfying Assumption 2 and let Q be our function class for Qs. We learn the sequence Q.....,,2 from prefix data (P7 ,,...,, T - T], via:
1,.... T),
be an
1. Initialization. Fit @ predict Y,+7 from the prefix embedding of (H 1.:1,,AA4t A1tr-1). Backward recursion. For = T 1,..., 1:
(a) Substitule interventions. For each prefix P7, replace A by the interventional a to form the modified history Ha H 1+t+k
(b) Generate pseudo-outcon outcomes. Let Y44++1 Qu+z H. ac+k). where Q6+1 was learned
in the previous step. These Y act surrogates for Y in the prefix data.
(c) Fil Qx- Regress Y++k41 the current embedding di ((Hi: 1H1 A+tk-1) to learn Qx Q.
3. Final step. Given a new history h and an interventional path a we predict
Ep|Y++ [Bcctr-1] oHH as) 6(b12 at)] Q (br... a.).
The iterative regression procedure yields consistent CAPO estimates provided each stage Qk is estimated consistently from data [22]. Informally, if the learned embedding & converges the true time-invariant representation é, and small perturbations in e or Q lead to proportionally small changes in predictions, then the overall recursive estimator remains consistent. These regularity conditions-formaliii through uniform stochastic equicontinuity-ar detailed Appendix C Formally, we state the following theorem:
Theorem 2 (Consistency of Iterative G G-Computation in Spatiotemporal Settings). Assume Assump- tions 1 and and that (a) the learned embedding T is Lg-consistent for o, and (b) each regression head Qu consistently estimates Qu and is uniformly well- well-behaved on Im & (intuitively, small input perturbations induce small oulpul changes). Let Z₁ = (His+k, Artx) denote the history-action pair at step k. Then
(Zo: Qi 1) (,))
so the recursive estimator Q of the CAPO is probabilistically consistent.
We provide proof of Theorem in Appendix [ In the following section, we instantiate this procedure in our GST-UNet architecture. illustrating how to incorporate spatial dependencies and interference into ó and each Qu, and implement streamlined, end-to-end training strategy that unifies history embeddings and outcome predictions.
We formalize "well- behaved" via uniform stochastic equicontinuity and continuity in Appendix C
5

---

Spatiotemporal Module: U-Net + ConVLSTM Attention
G-Heads
FX x
Legend
di A...e J.N.
R.
ConvISTM
Myp
de
de Raa Nyz
Conv2D+ReLU (x2)
(by
N,8 x
4,24 Nez d LEXMaa Nu
Upsample (by
Attention Gate
d,x K.
Sldp Connection
de
Concatenation
Q1
G-Heads
Figure 2: Overview of the GST-UNet architecture. The spatiotemporal learning module (left) is U-Net augmented with a ConvLSTM layer and attention gates. Its final feature map is passed to set of G-heads (right), where each G-head Qx implements iterative G-computation (see Algorithm T
GST-UNet Implementation
The theoretical results above establish how CAPOs can be identified and consistently estimated from single spatiotemporal trajectory. We now provide concrete neural implementation of this procedure. GST-UNet instantiates the iterative G-computation framework with a spatiotemporal deep architecture that embeds strong inductive biases-locade locality translation invariance, and temporal smoothness-we well suited to data-scarce settings. While alternative backbones could be employed, our U-Net with ConvLSTM and attention offers natural choice for learning stable, history-invarian representations that satisfy Assumption 2 that realizes the GST-UNet (Algorithm
We now describe the architecture and the training procedure
5.1 Model Architecture
The GST-UNet consists of two main components:
1.
Spatiotemporal Learning Module: a U-Net-based network augmented with ConvLSTM and attention gates for spatiotemporal processing.
2. Neural Causal Module: 7 G-computation heads, each mapping the spatiotemporal features to the final outcome predictions in the iterative procedure.
We illustrate the GST-UNet architecture in Figure and describe its main components below.
Spatiotemporal Learning Module. (1) Spatial Module. While our framework is agnostic to the choice of spatiotemporal learning module, we adopt U-Net with ConvLSTM and attention due to its strong performance data-scarce regimes. To efficiently process nigh-dimensional spatial data, we employ U-Net (37), fully convolutional architecture originally developed for biomedical image segmentation, It employs an encoder-decoder design with skip connections: the encoder progressively downsamples the spatial grid through convolution and pooling, while the decoder upsamples it back to the original resolution, merging encoder features at each scale. (2) Temporal Module. U-Net has limitations in capturing temporal information. To address this, we integrate Convolutional Long Short-Term Memory (ConvLSTMM layer 44 to the U-Net encoder. This module captures temporal dependencies by maintaining hidden state across time steps while aggregating spatial information through convolutions. After computing the final ConvLSTM state, we append static (time-invariant) covariates V as additional feature channels, ensuring the subsequent U-Net encoder-decoder has direct access to both temporal dynamics and static -ocation-specific information. In the decoder, we incorporate attention gates [29] to selectively highlight relevant spatial regions, refining skip connections and emphasizing critical global or local patterns. The embedding module ultimately produces dy di-iimennion feature map of size Nx XNy. capturing essential spatiotemporal context- including interference, spatial confounding, and static covariates for downstream G-computation.
Neural Causal Module. We attach G-computation heads to the U-Net's final feature maps, corre- sponding to the Qi estimators in the iterative procedure (see Section 4.2). Each head can be a small convolutional module or simple feed-fome feed-forward network, depending on how much spatial structure

---

Algorithm GST-UNet Training and Inference
1: Input: Horizon 7, prefixes (Pi interventions a curriculum O
(e) total epochs E.
2: Initialize: parameters # (U-Net embedding G-heads).
3: for 1==1... 1...B do
4: for = do
5: (Supervision) For each prefix i, predict outcomes 9 - Qu (o litkk-1 A A (1-1)): t-
6: (Generation (detached)) For each prefix i, generate pseudo-outcomen
Qu+z (
a 0)) h <<T,
40 [+++1
- ltr
k= T,
where the observed 1 were replaced with a in 1i4k
7: (Loss aggregation) Compute the MSE loss C(0;e) e) Ek=1 0 Σ, Ex
(e)
-
-++++1)
8: (Backward pass) Update 0 by backpropagation.
9:
(Inference) Given a h return Q (0(1 0)):
remains to be captured. The information flow at the G-computation heads proceeds as follows: each head Qu 7) receives the da Nx Ny U-Net embedding @H At+k-1) (encompassing spatiotemporal and static context) and outputs an Nx x Ny prediction for that time step. We refer this the supervision slep, since Q, compares its predictions the real observed outcomes Y Yitr, anchoring the model in genuine data, while each Qker compares its predictions to pseudo-outcomes Y1+++1 provided by QK+1. These pseudo-outcomes arise in subsequent gener- ation step, wherein Qx+1 processes the intervened history & (HHf: a(+k) in detached forward pass (so Qk+1 is not updated by Qx 's loss), thereby creating surrogate targets for Qu. This procedure realizes the iterative G-computation logic from Section 4.2, enabling GST-UNet to estimate future outcomes under various counterfactual treatments. By separating the spatiotemporal embedding from the G-heads, we maintain a common representation for all prefix data (see Assumption ( and flexibly capture interference and spatial confounding. Each G-head enforces the proper temporal adjustments to yield bias-free counterfactual inference.
5.2 Training and Inference
While each G-head Qx could be trained sequentially-from Qr down to Qi-by passing pseudo- outcomes backward through time, this creates conflict when all heads share the same U-Net embedding Φ. Specifically, each Qs may push &i toward optimizing its own objective, resulting in misaligned training signals and unstable learning.
Joint Loss and Multi-Task Training. To address this issue, we employ a joint (or multi-task) training approach 99 15] by aggregating the loss terms from all G-heads into a single objective, then backpropagating once per batch. Concretely, for each head Qk. let Y++++1 be the real outcomes if or pseudo-outcomes (generated by Qk+1) if r. Our head-specific loss a mean squared error (MSE) over all prefix samples:
T
(x) Σ ( litt- A ( 1))
[
(1)
2
r+++1]
where & encompasses all model parameters (the shared U-Net embedding ó and the G-heads Q+). (e) Let o denote head-weight for epoch e. We then form the overall training objective at epoch by
C(0;e) Σ Cx(
By summing the losses and performing single backward pass, we learn a common embedding that balances the needs of all G-heads, rather than fitting each head separately.
Curriculum Training. A naive implementation of Eq. 3 where each G-head given equal weight- can be suboptimal: early in training, Q, (which real data) is inaccurate, and the pscudo-outcomes
7

---

generated for Quer are effectively noise. Consequently, Qi...., may overfit to poor targets before Q, has converged, leading suboptimal solutions. To mitigate this, we employ a curriculum training approach [ gradually increasing the loss weight of earlier heads as Q- improves.
While many currieula are possible, adopt simple scheme controlled by single hyperparameter ec (the "curriculum period") so we can readily tune it. Let p(e) = min(T, [e/ea]] which indexes a "phase" based on the current epoch We then define
(n) [1/p(e),
if l E {T, otherwise.
..... p(() + 1}.
a 10. 0,
Hence, during epochs 1 S e S Ee (phase p(e) 1), only Q, is active with af 1; in the next interval 6, < e I 2e, (phase p(e) = 2). Q, and Q--1 each have weight 1/2, and so on until all heads are active with uniform weight 1/T. For Tec. training continues with at all heads. This schedule ensures Q, becomes reasonably accurate before earlier heads rely on its pseudo-outcomes. The hyperparamettr e, controls the pacing, helping prevent early training noise.
(e) 1/T for
We also adopt standard neural network practices, including mini-bateh maii- optimization and early stopping, to stabilize training and mitigate overfitting. At inference time, given a new history h1; and an interventional sequence ag(+r-1, compute Q1 (0(h1.,, a;); We sketch the overall training and inference procedure in Algorithm
() as target CAPO estimate.
Experiments
We evaluate the proposed GST-UNet framework through two applications. First, simulate synthetic data that incorporates key spatiotemporal causal inference challenges: interference, spatial confounding, temporal carryover, and time-narying confounding. Using this synthetic data generation process (DGP), we compare the GST-UNet algorithm against several baselines. Next, we demonstrate real-world dataset analyzing the impact of wildfire smoke respiratory
the utility of GST-UNet hospitalizations during the 2018 California Camp Fire.
Additional details-including exact simulation parameters, model architecture and execution se- tups, hyperparameter selection strategies, and validation procedures-can be found in Appendix D Replication code is available at hhtts:: /tttss:/giin //gi thub. comm/noopresccu/GSTTeeett
6.1 Synthetic Data
We generate T 200 time steps of a 64 64 (Nx x Ny) grid of observational data using the following data generating process (DGP):
X, ag + a,x, az 1-1 + (Kx X.-1) +
A Bern (3363000 + + KA KA* (A-X)l)))).
I (Krx & X + Y
Y. To 7 (KYA AI-1) + 72 Σ
where dx = 1, "*" denotes 3x3 spatial convolution over the Nx x Ny grid, and 6x,6y ~ N(0, 1) are i.i.d. noise. Fach kernel Kx, KA, KY Kyx encodes a local advection-diffusion process that mimics wind-driven pollutant transport, with interventions A, injecting additional emissions that propagate through the same kernel. This physically realistic setup produces interference, spatial confounding, and temporal carryover- the three challenges GST-UNet is designed to address. Each equation is evaluated at every spatial location, X,, A,, and Y, are Nx X Ny matrices. Here, X, acts as a vary-ng confounder: its past influences both A, and Y,. while current interventions A, affect future X,+1- For example, A, may represent regulatory actions, X, air quality, and Y; health outcomes- outcomes-Cappuurng feedback from policy to exposure to outcome, and back to future policy.
We vary 31 to control time-varying confounding: when 3 = 0, X, does not affect A,, eliminating confounding; larger values increase its strength. For each 91, we generate 50 test trajectories from random initial states, fix their histories, and simulate 100 7-step counterfactual futures to estimate true
8

---

Table 1: RMSE + SD across test trajectories. Bold indicates lowest per column; color shows improvement (RMSE decrease or increase) over best baseline (excluding ablations).
Model
$ = 0.0
B1 = 0.5
= 1.0
31 = 1.5
81 = 2.0
UNet+
0.28 0.00
0.36 0.00
0.54
0.01
0.81
STCINet
0.29 - 0.00
0,38
0.02
0.00
0.900.000
IPWUNet
0,60 0.01
0.58
0.58
0.50
0.59
5
GST-UNet w/o Attention
0.50 0.00
0.00
0.00
0.45
0.07
GST-UNet w/o Curriculum
0.69 0.00
0.00
0.00
0.01
0.61
GST-UNet
0.33 0.00
0.35
0.40
0.40 0.00
(+17.9%)
(-2.7%)
(-2166%)
(-25.4%)
(-32.2%)
UNet+
0.28 = 0.00
0.61 0.00
1.18
1.00
1,71 0.01
STCINet
0.31 + 0.00
0.68 0.00
1.25
1.01
1.60 # 0.01
IPWUNet
0.78 0.01
0.00
0.96
1.19
1.08 = 0.01
10
GST-UNet w/o Attention
0.00
0.00
0.00
0.79
1.07
GST-UNet w/o Curriculum
0.00
0.00
1.02
1.12 1.08
GST-UNet
0.38
0.00
0.08
0.73
0.85 # 0.01
(+35.7%)
(-9.8%)
(-29.2%)
(-38.7%)
(-21.3%)
CAPOs, with т E {5, 10). We compare GST-UNet against three baselines: (i) UNet+, which uses U-Net ConvLSTM + Attention backbone with A, as an input channel but performs no iterative adjustment; (ii) STCINet [.] which estimates direct and indirect effects without modeling time- varying confounding; and (iii) IPWUNet, an ewvrse-propennte weeitty--eiggingggg variant that reweights pseudo-outcomes using a UNet-style propensity estimator but cannot correct for spatial interference (details Appendix D). We also test ablations of GST-UNet without curriculum or attention. Table shows that when 81 O, UNet+ performs best G-sopppccmtttaio is unnecessary and adds noise. As 31 increases, UNet+ and STCINet degrade sharply, while GST-UNet remains stable, IPWUNe: shows some benefit but is biased even at 31 due to uncorrected interference. GST-UNet consistently outperforms all baselines, demonstrating the value of iterative G-computation. Curriculum training substantially improves performance across horizons, while attention yields modest gains consistent with our predominantly local dynamics. Additional ablation analyses, including neighbor aggregation experiments, reported in Appendix D
6.2 Impact of Wildfires on Respiratory Health
Wildfire smoke has been linked to short-term respiratory harms [10 33 34 with older adults especially vulnerable 14]. the time this work was conducted (January 2025), a series of 14 destructive wildfires affeectdd the Los Angeles metropolitan area and San Diego County in Californaa underscoring the urgency of understanding the health impacts of such events. In this study, we focus previous large-scale episode: the 2018 California wildfire season (45) which included the Carr Fire (July-August) and the Camp Fire (November) and significantly worsened air quality.
We use daily, county-level data from Letellier et al. 23 (see Appendix D.2) including PM2.s. respi- ratory/cardiovaseuaar hospitalizations, and weather variables (temperature, precipitation, humidity, radiation, wind), along with population estimates from the California Department of Finance. Each of the weather variables can be a time-varying confounder: weather conditions affect future smoke levels and health outcomes, while also being influenced by prior smoke levels.
We focus on weeks 20-48 (May [8-December 2, 2018), covering the Carr and Camp fires. Following standard practice, we label county as "treated" on days with mean PM2.5 > 10 μg/m² and use raw hospitalization counts (rather than per-10,000 incidence, which can be unstable for small counties). We interpolate daily county-level data (treatment, outcome, five covariates) onto 40 x 44 Attiude-longitut grid, discarding cells outside California, yielding spatiotemporal tensor of size 203 x 7 x 40 x 44. Interpolation ensures each grid cell approximates the region it overlaps (area- weighted), enabling the model to capture spatial gradients in PM2.5. weather, and hospitalizations. We train GST-UNet with horizon 10, using the Carr Fire period (June-July) for validation, and generate counterfactual predictions for the Camp Fire peak, November 8- 8-17. See Appendix D.2 for preprocessing and masking details.

---

PM2.5 Levels California Counties
PM
10 puinn During Camp Fire
Camp Respiratory
200
42-
Counterfact Incidence
150
Camp Fire
40-
10,000)
Fire
2.0
100
Latituee
38-
Latitude
36-
1.0
50
34-
0
Jun Jul Aug Sep Oct Nov Dec Months in 2018
-144
-122
1200
-18 -16
-114
-124 -122 120 -118 -114
Longitude (*)
Longitude (*)
Figure 3: (Left) Daily PM2.5 levels across California from May to December 2018, with red lines marking major wildfires. (Center) Counties exposed to average PM2. > 10 pg/mmm during the Camp Fire (red), origin county dark red. (Right) Factual minus CAPO-predicted daily respiratory admissions during peak Camp Fire. Hashed areas indicate small-population counties (< 30,000).
Figure [ (left) shows the rise in PM2.5 during the mid-late 2018 wildfire season; (center) highlights counties with daily PM2.5 > 10, pg/m². Using GST-UNet, we estimate daily CAPOs had the Camp Fire not occurred (i.e., setting PM2.5 S 10, pg/ m /g/m²² statewide). Figure 3 (right) compares these to factual daily incidence (hospitalizations per 10,000 residents). To reduce small-sample variability, we exclude counties with population below 30,000 (ys. > 70,000 for others), marking them with hatching (see Appendix(D22) D.2). Over November 8-17, GST-UNet predicts approximately 4,650 excess respiratory hospitalizations (465/day) attributable to the Camp Fire, with the highest incidence near the fire source. This aligns with 95% bootstrap confidence interval of [1888, 6535]. UNet+ yields lower mean and higher uncertainty (3,981; [- 899, 5202)). STCINet produces highly variable near-zero estimates (88; [-3077, 3281)), and IPWUNet gives implausibly high, near-constant values (~20,500), reflecting limitations of weighting under rare-event support. These results underscore GST- UNet's improved stability and accuracy in counterfactual estimation. Our findings are qualitatively consistent with Letellier et al. (23, who report 259 excess daily cases averaged over longer, lower- intensity window (Nov 8-Dec 5). Overall, the GST-UNet captures spatiotemporal variation in smoke exposure and health outcomes, illustrating its promise for real-world causal inference in domains such environmental health and policy.
Conclusion
We presented GST-UNet, neural framework for spatiotemporal causal inference that combines U-Net-based representation learning with iterative G-computation to adjust for time-varying con- founders. GST-UNet addresses key challenges such as interference, spatial confounding, temporal carryover, and time-varying feedback. We establish theoretical identification and consistency guar- antees, validate performance in synthetic settings with controlled confounding, and demonstrate practical utility in estimating the impact of wildfire smoke exposure during the 2018 Camp Fire. Together, these results position GST-UNet as ready-to-use tool for practitioners, offering reliable, interpretable causal estimates in complex spatiotemporal environments. We discuss limitations and broader impacts in Appendix E
Acknowledgements
We thank the anonymous reviewers for their thoughtful feedback and the constructive dialogue during the review process, which greatly strengthened the final version of this work. Miruna Oprescu and BNL team (D. Park, X. Luo, Yoo) were supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, under Awards DE-SC00231 and DE-5,0011704, respectively. Nathan Kallus was supported by the U.S. National Science Foundation under Grant No. 1846210. Part of this work was conducted while Miruna Oprescu was a research intern at Brookhaven National Laboratory. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the U.S. Department of Energy or the U.S. National Science Foundation.
10

---

References
[1] S. Ali, O. Faruque, and J. Wang. Estimating direct and indirect causal effects of spatiotemporal interventions in presence of spatial interference. Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 213 130. Springer, 2024.
[2] L. Anselin, Spatial econometrics: methods and models, volume Springer Science & Business Media, 2013.
[3] H. Bang and J. M. Robins. Doubly robust estimation in missing data and causal inference models. Biometrics, 61499662,,77, 2005.
[4] E. Ben-Michael, A. Feller, and Rothstein. Synthetic controls with staggered adoption. Journal of the Royal Statistical Society Series B: Statistical Methodology, 84(2):351-385 2022.
[5] Y. Bengio, J. Louradour, R. Collobert, and J. Weston. Curriculum learning. In Proceedings of the 26th annual international conference machine learning, pages 41-48, 2009.
[6] G. Bertasius, H. Wang, and L. Torresani. Is space-time attention all you need for video understanding? In ICML, volume 2, page 4, 2021.
[7] Bica, M. Alaa, J. Jordon, and M. van der Schaar. Estimating counterfactual treat- ment outcomes over time through adversarially balanced representations. arXiv preprint arXiv:2002.04083 040000000403,.... 2020.
[8] I. Bojinov and N. Shephard. Time series experiments and causal estimands: exact randomization tests and trading. Journal of the American Statistical Association, 2019.
[9] R. Caruana. Multitask learning. Machine learning, 28:41-75, 1997.
[10] W. E. Cascio. Wildland fire smoke and human health. Science of the total environment, 624: 586-595, 2018.
[11] R. Christiansen, M. Baumann, Kuemmerle, M. D. Mahecha, and Peters, Toward causal inference for spatio-temporal data: Conflict and forest loss in Colombia. Journal of the American Statistical Association, 11 117(538):591-60) 2022.
[12] S. E. Cleland, M. L. Serre, A. G. Rappold, and J. West. Estimating the acute health impacts of fire-originated PM2.5 exposure during the 2017 California wildfires: Sensitivity choices of inputs. Geohealth, 5(7)x20021GH00444 2021.
[13] A. Coston, Kennedy, and A. Chouldechova. Counterfactual predictions under runtime confounding. Advances in neural information processing systems, 33:4150-4162, 2020.
[14] S. DeFlorio-Barkee J. Crooks, J. Reyes, and A. G. Rappold. Cardiopulmonary effects of fine particulate matter exposure among older adults, during wildfire and non-wildfire periods, in the Environmental health perspectives, 127(3):037006, 2019.
United States
[15] Evgeniou and M. Pontil. Regularized multi-task learning. Proceedings of the tenth ACM SIGKDD international conference Knowledge discovery and data mining, pages 109-117, 2004.
[16] D. Frauen, K. Hess, and S. Feuerriegel. Model-agnostic meta-learners for estimating heteroge- neous treatment effects over time. arXiv preprint PrXXi:2407.002877 2024.
[17] X. Glorot and Y. Bengio. Understanding the difficulty training deep feedforward neural networks. In Proceedings of the 13th International Conference Artificial Intelligence and Statistics (AISTATS), pages 249-256. JMLR Workshop and Conference Proceedings. 2010.
[18] K. Hess, D. Frauen, V. Melnychuk, and Feuerriegel. G-transformer for conditional average potential outcome estimation over time. arXiv preprint arXiv:2405.21022 2024.
11

---

[19] K. Jordahl, J. V. den Bossche, M. Fleischmann, J. Wasserman, J. McBride, J. Gerard. J. Tratner, M. Perry, A. G. Badaracco, C. Farmer, G. A. Hjelle, A. Snow, M. Cochran, S. Gillies, I.. Culbertson, M. Bartos, N. Eubank, maxalbert, A. Bilogur, S. Rey, C. Ren, D. Arribas- Bel, L.. Wasser, J. Wolf, M. Journois, Wilson, Greenhall, C. Holdgraf, Filipe, and F. Leblanc. geopandas/geopandar v0.8.1, July 2020. URL httpss//dde org/10 5281/ .en66o.3946761
[20] I.. J. Keele and R. Titiunik. Geographic boundaries as regression discontinuities. Political Analysis, 23(1):127-155, 2015.
[21] D. Kingma. Adam: method for stochastic optimization. arXiv preprint arXiv: 1412.6980 2014.
[22] M. J. Laan and J. M. Robins. Unified methods for censored longitudinal data and causality. Springer, 2003.
[23] N. Letellier, M. Hale, K. Salim, Y. Ma, F. Rerolle, L. Schwarz, and T. Benmarhnia. Applying two-stage generalized synthetic control approach to quantify the heterogeneous health effects of extreme weather events: A 2018 large wildfire in California event a case study, Environmental Epidemiology, (1)xe336, 2025.
[24] R. Li, S. Hu, M. Lu, Y. Utsumi, P. Chakraborty, D. M. Sow, P. Madan, J. Li, M. Ghalwash, Z. Shahn, al. G-net: a recurrent network approach to g-computation for counterfactual prediction under dynamic treatment regime. In Machine Learning for Health, pages 282-299. PMLR, 2021.
[25] Y. Li, R. Yu, C. Shahabi, and Y. Liu. Diffusion convolutional recurrent neural network: Data- driven traffic forecasting. arXiv preprint arXiv:1707.011966 2017.
[26] B. Lim. Forecasting treatment responses over time using recurrent marginal structural networks. Advances in neural information processing systems, 31, 2018.
[27] Z. Liu, J. Ning, Y. Cao, Y. Wei, Z. Zhang, S. Lin, and H. Hu. Video swin transformer. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 3202 2211, 2022.
[28] V. Melnychuk, D. Frauen, and S. Feuerriegel. Causal transformer for estimating counterfactual outcomes. In International Conference on Machine Learning, pages 15293-15329, PMLR, 2022.
[29] O. Oktay, J. Schlemper, L. L. Folgoe, M. Lee, M. Heinrich, K. Misawa, K. Mori, S. McDonagh, N. Y. Hammerla, B. Kainz, et al. Attention u-net: Learning where to look for the pancreas. arXiv preprint arXiv: /804.039999 2018
[30] G. Papadogeorgou, F. Mealli, and C. M. Zigler. Causal inference with interfering units for cluster and population level treatment allocation programs. Biometrics, 75(3):778-777, 2019.
[31] G. Papadogeorgou, K. Imai, J. Lyall, and F. Li. Causal inference with spatio-temporal data: estimating the effects of airstrikes insurgent violence in iraq. Journal of the Royal Statistical Society Series B: Statistical Methodology, (9(5))1969-1999 2022.
[32] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L.. Antiga, A. Desmaison, A. Kopf, Yang, De Vitioo, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, I.. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, nigh-performance deep learning library. NeurlPS, 2019.
[33] C. E. Reid, M. Brauer, F. H. Johnston, M. Jerrett, J. R. Balmes, and C. T, Elliott. Critical review of health impacts of wildfire smoke exposure. Environmental health perspectives, 124 (9):1334-134) 2016.
[34] C. E. Reid, M. Jerrett, I. B. Tager, M. L. Petersen, J. K. Mann, and J. R. Balmes. Differen- tial respiratory health effects from the 2008 northern California wildfires: A spatiotemporal approach. Environmental research, 150:227-235, 2016.
12

---

[35] J. Robins and M. Hernan. Estimation of the causal effects of tame-varying exposures. Chapman Hall/CRC Handbooks of Modern Statistical Methods, pages 553-599, 2008.
[36] M. Robins, M. A. Hernan, and B. Brumback. Marginal structural models and causal inference in epidemiology, 2000.
[37] O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention-MICCA 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234-241. Springer, 2015.
[38] D. B. Rubin. Bayesian inference for causal effects: The role of randomization. The Anmals of statistics, pages 34-58, 1978.
[39] N. Seedat, Imrie, A. Bellot, 7. Qian, and M. van der Schaar. Continuous-time modeling of counterfactual outcomes using neural controlled differential equations. arXiv preprint arXivv2206.088: 08311, 2022.
[40] X. Shi, Z. Chen, H. Wang. D.-Y. Yeung. W.-K. Wong, and W.-c Woo. Convolutional Istm network: A machine learning approach for precipitation nowcasting. Advances in neural information processing systems, 28, 2015.
[41] C. Song, Y. Wang, X. Yang. Y. Yang, Z. Tang, X. Wang, and J. Pan. Spatial and temporal impacts of socioeconomic and environmental factors on healthcare resources: county- bayesian local spatiotemporal regression modeling study of hospital beds in southwest china. International Journal of Environmental Research and Public Health, 17(16):5890, 2020.
[42] M. Tec, G. Scott, and C. M. Zigler. Weather2vec: Representation learning for causal inference with non-local confounding in air pollution and climate studies. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 14504-14513, 2023.
[43] A. W. Van Der Vaart, A. Wellner, A. W. van der Vaart, and J. A. Wellner. Weak convergence. Springer, 1996.
[44] Y. Wang. Causal inference under temporal and spatial interference. arXiv e-prints, pages arXiv-2106, 2021.
[45] Wikipedia, Camp Fire (2018) Wikipedia, the free encyclopedia. http:// wwwpppiiiiiiipiiiiiaiaaa accessed 29-Jamuary-20222
2025. [Online;
[46] 7. Wu, S. Pan, G. Long, J. Jiang, and C. Zhang. Graph wavenet for deep spatial-temporal graph modeling. arXiv preprint arXiv: /00::990.000 00/20002,100 2019.
[47] J. Zhang, Y. Zheng, and D. Qi. Deep flows prediction. In Proceedings of the AAAI conference artificial intelligence, volume 31, 2017.
residual networks for citywide crowd
[48] W. Zhang and K. Ning. Spatiotemporal heterogeneities in the causal effects of mobility intervention policies during the covid-19 outbreak: A spatially interrupted time-series (sits) analysis. Annals of the American Association of Geographers, 113(5):1122-113 2023.
[49] L. Zhou, K. Imai, J. Lyall, and G. Papadogeorgou. Estimating heterogeneous treatment effects for spatio-temporal causal inference: How economic assistance moderates the effects of airstrikes on insurgent violence. arXiv preprint arXiv.2412.15128, 2024.
13

---

Table Key differences between prior neural G-computation methods and GST-
Aspect
Prior Neural G-Co G-Computation
GST-UNet (ours)
Data structure
Many independent temporal trajectories (e.g.. patient sequences); inter-unit interactions such spatial dependence.
Single spatiotemporal chain where out- comes, covariates, and treatments evolve jointly across lattice; strong spatial cou- pling and interference.
Encoder
RNN/Transformer over time only.
ConYLSTM- UNet encoder aggregates neighbour covariates/treatments before G-heads, capturing interference and spa- tial confounding.
Training
Standard end-to-end due to i.idd trajec- tories; stability arises from large data rather than curriculum or spatial priors.
Murriculum-stabille multi-head train- ing for accurate pseudo-outcome genera- tion under limited samples.
Theory
Classical G-formula under i.i.d trajecto- ries; no single-chain guarantees.
Identification (Theorem and consis- tency (Theorem 3 under representation- based time-invariance for a single chain.
A Extended Literature Review
Classical Spatiotemporal Causal Inference. Early spatiotemporal causal inference methods- [20]. and synthetic controls [4] ---rrvvide useful frameworks for estimating treatment effects across regions but rely strong assumptions such as parallel trends or stable treatment assignment. These approaches struggle with interference, nonlinear dependencies, and time-varying confounders, limiting their applicability in complex settings, More recent approaches for spatiotemporal causal inference handle time-varying confounding through inverse propensity weighting (IPW), typically by extending marginal structural models the spatial or spatiotemporal domain. For instance, Papadogeorgou et al. [31 and Zhou et al. [4) employ IPW- style adjustments to estimate regional average treatment effects across space and time. However, these approaches cannot accommodate interference unless strong assumptions are made-e.g.. defining user-specified exposure mapping or restricting attention to hyper-local interactions (see also [1] 31 44 48). Such simplifications may be ill-suited for weal-world systems with rich spatial dependencies. Moreover, even recent advances in this space remain limited; as noted by Zhou et al. [4). the literature on spatiotemporal causal inference remains sparse, especially in settings with feedback loops or virying confounding.
including spatial econometrics [2), difference-in-diffeeee
Machine Learning for Spatiotemporal Modeling, Spatiotemporal predictive modeling has seen rapid progress with the rise of deep learning. Convolutional and recurrent neural networks are widely used for forecasting spatially indexed time series (c.g., weather or traffic) 40 47), while graph-based methods (c.g., Graph WaveNet (46), Diffusion Convolutional RNN (25 capture non-Euclidean spatial dependencies. Vision transformer variants, including Video Swin Transformers (27 and TimeSformer (6. extend attention- based models to spatiotemporal video data. These architectures can learn complex non-local interactions over space and time. However, such models are typically optimized for prediction tasks and do not include causal adjustments. Without mechanisms like propensity modeling G-computation, they remain ill-equipped to estimate counterfactual outcomes or adjust for varying confounding. Some recent work integrates spatial representations for causal inference-e.g. Tec et al. 42 incorporate non-local confounders using UNet-based model-but these methods do not explicitly model dependencies over time or adjust for time sire-varying confounders.
Time-Series Causal Inference. In the longitudinal domain, time-series causal inference has devel- oped tools for handling temporal confounding using models such marginal structural models (36.. IPW-style estimation 26). and iterative G-computation [35]. Recent ML-based extensions include recurrent networks 7 39). Transformers [8 28 and meta-learners [6)] However, all these meth- ods assume access to independent time series e.g., across units or patients-whi allows for pooling across trajectories. These methods do not consider spatial dependencies, interference, or scenarios with a single observed spatiotemporal realization. As such, while they may handle time-varying confounding, they do not generalize to our setting. 'ablle summarizes the key methodological differences between GST-UNet and prior neural G-computation frameworks.
14

---

Neural-Based Spatiotemporal Causal Inference. There has been limited work neural models that explicitly address spatiotemporal causal inference. Tec et al. 42 use U-Net backbone to learn spatial representations for causal inference in air pollution studies but do not address time-varying confounding or feedback loops. Ali et al. 1 present U-Net-based architecture for predicting direct and indirect effects in climate contexts, but primarily focus on forecasting rather than causal identification. While these works highlight growing interest neural approaches to causal inference in spatiotemporal domains, none incorporate iterative adjustment procedure like G-computation that handles time-varying confounders, leaving identification in these settings largely unaddressed.
Our Contribution. GST-UNet bridges these gaps by combining flexible spatiotemporal neural architectures with theoretically grounded iterative G-computation framework. This allows valid estimation of potential outcomes in the presence of interference, spatial confounding, and time-varying confounding-without requiring practitioners to specify structural models or exposure mappings. To our knowledge, this is the first end-to-end framework to implement G-computation for causal inference over a single spatiotemporal trajectory. We integrate spatiotemporal processing via U-Nets and ConvLSTMs with principled multi-head neural causal module, and we design curriculum-bac -base training strategy to stabilize learning recursive pseudo-outcomes. Together, these components yield ready-to-use tool for practitioners, with consistent identification guarantees and robust empirical performance. By abstracting away the modeling choices typically required in structural spatiotemporal methods, GST-UNet makes spatiotemporal causal estimation more accessible, interpretable, and reliable for real-world applications.
B Proof Theorem
We aim to show that under Assumption | and Assumption 2 the CAPOs in Equation identified recursively from single time series via sequence of conditional expectations.
| can be
Step 1: Recursive decomposition for the intractable expectation We first demonstrate the recursive decomposition of the intractable expectation in the CAPO definition (Equation (T). While this expectation is theoretically well-defined, it cannot be directly estimated in practice due to the limited availability of data. Specifically, one sample of the history at time t + r for each Nevertheless, as we will show, we can convert these expectations into expectations over prefix-based segments that allow us to estimate these quantities from the data.
only observe single time series, meaning have just
Starting from Y [ax+r-11] H1: = ha:d,, we have:
Y &+r -
H = h
E YY+t
Ye+ [1
H = h A, = a.]
(Sequential ignorability and positivity (Assumption) !
E Y c+r
B
Hi Hist his: A. ar
(Law of total probability)
a
Ha 1
A411 a++1] H h A, = a.]
(Sequential ignorability and positivity)
E (a
H H
H
H A4t1 a++1]
H = h1:t, A, = a
(Law of total probability)
E (+ 11177-1
E
H A44 a+4a
H H Att1 a H h A, a
(Sequential ignorability and positivity)
E
Ha E e+r (autr-1l
H14 A a:4r-1]
Ha Hi A+t--2 ar+r-2]
H h A, a.
(Sequential ignorability and positivity)
H. H A = a
15

---

H 2: = 2
His h1:1, A. ai
(Consistency)
Thus, we had multiple spatiotemporal time-series samples, we could directly estimate this nested expression from data, since the right-hand side depends solely on observed quantities, ensuring identifiability.
Step 2: From intractable to prefix-based expectations
We now show how to estimate the nested we can rewrite the inner-most expectation
expectations using the prefix data. First, by Assumption 2
as
E Y
H H A a Rp Y (+
0 H 184--1:
A+T
Q, 1: 1 a(1-1))
(Definition of Q,)
Thus, by using Assumption 1
we can write this expectation over the prefix data which we have many samples of. Now consider the next nested expectation:
E E
H 1. 11+7-1)
H H = hi
A 21+7-2
f
Qr 2e
/0 a 2 2)) Btttr-
(Assumption
Ep(Qr (1 21 = a
Q--1 (h 1++++ a++r-2)
Tracing this argument recursively through the nested expectation in Step obtain:
Y|Y++r [ax+r-1] His: = h1.:] = (h ar);
as desired. Thus, Q1 which can be estimated from the prefix data recovers the CAPOs, under our assumptions, even from single chain.
C Consistency of the Iterative G-Computation Estimator
In this section, we state the conditions under which the iterative G-computation procedure in Sec- tion yields consistent estimator, and show that our implementation of the Qx estimators satisfies these conditions.
We denote the L2 norm of a function f as |lfll2 == (x((())))1/ where the expectation is over the probability distribution P. The notation fn represents the estimated value of a parameter or function learned on n data points, where ʃ is the true value. For a sequence of random variables
Notation:
(Znh221 write Z, (1) if Pr(IZ.] > 0 for every > 0, i.e. Zn 2 0.
To begin, we introduce the following stochastic equicontinuity condition from 43
Definition 1 (Stochastic equicontinuity 43 Def. 1.5.7)). Ler (Z,d) be semi-metric space and (frta2 C (2) sequence of random functions. It is uniformly stochastically equi-continuous if. for every > 0, rj > 0, there exists a 8 > O such that
limsup Pl sup d(
[n(2)) - [a(=') > <1
n-000
Stochastic equicontinuity ensures that, with high probability, each estimator changes only slightly when its input is perturbed by small amount. It is strictly weaker than global Lipschitz continuity any family that is Lipschitz on a bounded domain with constants bounded in probability automatically satisfies Definition We impose this condition in Theorem so that the (1) error in the learned embedding propagates to only (1) errors in the G-heads, making the recursive estimator consistent.
The following theorem restates Theorem from the main text in full detail and provides its proof.
16

---

Theorem 3 (Consistency under Uniform Stochastic Equicontinuity). Suppose the conditions of Theorem hold, and let é be a learned embedding. Define Zu= ;= (H Actk), and recursively define the learned estimators Q& (Z == Ept (Z):d) Ø) [(Zx)] for l L.....T, T, with
terminal condition 1Zz 3) = yt+r. Assume that {Qx are obtained via the iterative G-computation algorithm.
(i) 10-
(1) for all k; (iii) for every k the random maps + Qh, a; 2) are stochastically equicontinuous Im ø (Definition 1. and Qu(.)) is uniformly continuous there,
(i) 1Q - Qu(Za-1;
then
Q(000)))) oy()).
Thus the recursive G-computation estimator is (probabilissicallyy consistent.
Proof. We proceed by reverse induction on k, starting from k = T and working backward 1. For each ki we aim to show:
Qu 11 = (1,)
Base case (k = T). By definition, 1 = y which observed. Thus,
= (- 3) = Ep|yttr [Z,)) and - ((zz))
We decompose the difference:
Q- (Z. -z- Q-. Q- S (Z. Q (Z. + -10 1--
A2
Term A2 is (1) by assumption (ii). Term A₁ converges to zero in probability due to (i) - 1 ø|12 (1) and (iii) stochastic equicontinuity of Q,. Therefore,
Q- (Z. (1).
Inductive step. Suppose for some 1 + 1 S T that
@ktz (Zk Qk+1 (Ze: 10/13 = (1).
We now consider
(Za -E ( 1 3(Zz)), (u(Zn-1id) =E(Q=+ I(Qx+1(Ze;() (Zz)).
Again, decompose:
(A -
- (Zc- 1 S | (2t-1: - (Zk-1:
Qx
+ - - QZ
The second term is (1) by assumption (ii). The first term is also (1) because & + Φ in Lo and the stochastic equicontinuity of Qx ensures that perturbations in yield small changes in predictions uniformly over Im é. Thus,
- o
= By induction, the result holds for all k: T, 1,.... and in particular:
[Z:: Ō) Qi 22 (1).
Thus, the proof is now complete.
17

---

Example (Feed-forward or convolutional heads), Suppose each G-computation head Qx(: 2) is implemented as a depth-d neural network
(2) = W.
where the activations oe are Lipschitz continuous (e.g. ReLU, Leaky RELU, SoftPlus, Tanh, Sigmoid, or ArcTan). If each layer weight satisfies a spectral norm bound ||Well2 S pt so, then Vi is globally Lipschitz on Rh with constant -= II Pt. and thus uniformly continuous any compact subset. This implies the stochastic equicontinuity condition in Definition] 0
In practice, norm control can be enforced via weight decay, spectral normalization, or weight clipping during training. Similarly, the encoder output 6(H, A) can be bounded-e-g e.g. through normalization or clipping-st its image lies in compact subset of RR. Together, these ensure the continuity and equicontinuity conditions required by Theorem 3
The same argument applies to convolutional networks, since 2-D convolutions are linear operators whose induced matrix representations also admit spectral norm bounds controlled via spectral normalization.
D Experimental Details
In this appendix, provide further information the simulation experiments (Section 6.1 and the real-world wildfire application (Section 6.2), including exact parameter settings, model architecture and execution details, hyperparameter selection strätegies, and validation procedures. All code for generating, preprocessing, and analyzing both the synthetic and real-world datasets- and for training and evaluating GST-UNet-iss available https: / /github. com/ Comm/oopresc//GSTUNe with step-by-step replication instructions in the repository's README. md.
For both applications, GST-UNet employs U-Net backbone with single ConvLSTM layer (hidden dimension 32) and a seirraatingexpaddingg path of channel sizes 16 -> 32 > 64 + 128 + 256. The G-computation heads implemented as shallow feed-forward neural networks that operate the U-Net feature maps at each grid cell for G-computation. In practice, to ensure stable ConvLSTM training and reduce computational overhead, neural networks are implemented via the nn module in PyTorch [32]. Experiments were conducted on an NVIDIA A100 (Ampere) GPU using the Perlmutter system the National Energy Research Scientific Computing Center (NERSC). The synthetic experiments required roughly 55 minutes per hyperparameter set, while the wildfire experiment completed in about minutes.
truncate the input history fixed length. All
D.1 Synthetic Experiments
Data Simulation Process. For our primary simulation experiments, we generate T 200 time steps on a 64 x 64 grid. The simulation parameters in the generating equations
X, ao + a1 R,xX,- + 02 A1AA11 as(Kx xx + +
A, Bern (30 Σ 1=0
KA+X(
KA*
1
Y, = 7o 7 Ac-1) + 722 Σ
(Kyx * X(-i) + + +ey,
are given by:
Xx;
0
0.45
0
= 0.5, = 0.5, a2 = -2.0, 0.2, Kx =
0.15
0
0.35
0
0.05
0
where Kx influences how diffuses across neighboring cells, with an asymmetry due to advection.
18

---

Feature (Xyo)
Intervention (Apo)
Outcome (Yial
0.
40
04
40
50
50
60
60
0.0
60
0.8
30
30
30
0.4
00
50
50
600
60
0.0
60
0.8
20
20
20
30
30
40
40
4.4
40
30
30
50
60
0.0
60
45
Figure 4: Samples from the DGP at t = 100, comparing feature X (left), intervention A1o0 (center), and outcome Yior (right) for varying {0.0, 1.0,
Table Hyperparameters and their ranges. We boldface the values that provided the best validation performance.
Hyperparameter
Model(s)
Value Range
Batch size
All models
[2, 4, 8)
Learning rate
All models
[10-4 5 x 10 10-4, 1.-33
-,000,00 10
Scheduler patience
All models
(3, 5, 10)
Early stopping patience
All models
(5, 10)
Curriculum period
GST-UNet
5,
Curriculum learning rate
GST-UNet
(10-1, 5 10 10-4, 10-3)
UNet output dim d,
GST-UNet
(8, 16,
G-head hidden size
GST-UNet
[8, 16)
G-head layers
GST-UNet
[1, 2,
Ar:
1.0
1.0
1.0
1
Go = 110, 81 E (0.0, 0.5, 1.0, 15, 2,0,,,,,,, KA
1.0
8.0
1.0
16
1.0
1.0
1.0
Ye:
= 2.0, = 1.5, 72 = 0.5, 73 0.5
1.0
1.0
1.0
1.0
1.0
1.0
1
Kvx
1.0
8.0
1.0
KYA
1.0
8.0
1.0
16
16
1.0
1.0
1.0
1.0
1.0
1.0
We use L 5 temporal lags for X and Y, seed of 42 for reproducibility. The parameter values were chosen such that the simulation remains stable (i.e., the process does not diverge). See Figure for representative
100 snapshots of X inn, Anno, and Yio under varying 31. For each 3, we first generate factual dataset of length T 200 (i.e. ((X,, A., Y.)) create Thest 50 test histories of length (œ 10. For each test history, simulate 100 traeectories
200 We then
19

---

Table 4: Ablation on spatial kernel size (r = 5). Removing neighbor aggregation (1 1 kernel) degrades performance, confirming the need to model spatial spill-overs.
Kernel size
=100.
0.5
1.0
1.5
2.0
3 x 3
0.33 + 0.004
0.35 + 0.004
0.40 + 0.005
0.44 + 0.004
0.40 = 0.005
1x1
0.53 + 0.004
0.55
0.005
0.54 = 0.005
0.60 + 0.007
0.64 = 0.006
Table 5: Effect of increasing 7 on RMSE for 81 = 2.0. GST-UNet improves with more data, while baselines remain biased.
Model
T=100
T=200
T=400
T=600
T=800
UNet+
0.78
0.81
0.82
0.95
0.87
STCINet 0.80
0.90
1.04
1.02
0.91
GST-UNet 0.69
0.40
0.32
0.32
0.36
under randomly chosen (yet fixed over the test data) counterfactual intervention of length 10, and average the outcomes at each step to approximate the true CAPOs. This procedure yields a final test set of shape Thest X (1 + 1) x 64 x 64, i.e., 50 21 64 x 64.
Neural Architectures. The GST-UNet comprises a single ConvL.STV layer (hidden dimension 32), followed by a U-Net with channel sizes 16 + 32 - 64 + 128 + 256. Its G-computation heads are shallow feed-forward networks operating on the final U-Net feature maps at each grid cell; both the U-Net's output dimension (da) and the G-head architecture (number of layers, hidden size) are treated as hyperparameters. The UNet+ baseline uses the same CONVLSTM+U-Net backbone as GST-UNet but outputs single channel (d, 1), omitting any Computaaiion. For direct comparison, also implement STCINet [T] with an identical ConvLSTM+U-Ner backbone, and retaining their original Latent Factor Model (LFM) details.
IPWUNet Baseline. We adapt the Inverse Propensity Weighting (IPW) estimator from to the spatiotemporal setting. Given estimated propensities # (a; Hıæ), the estimator is defined as:
I[ a:] #(a H1,) )
P- E+
CAPO B H1s
h1g).
We implement the IPWUNet baseline by reusing the UNet+ architecture (U-Net + ConvLSTM + Attention) for both propensity estimation and outcome prediction. Specifically, propensity model with a binary cross-entropy loss to estimate (A,, H.) at each time t. We then freeze this model and use the estimated weights train second instance of the same architecture with a weighted MSE loss, where pseudo-outeomes are reweighted by the estimated inverse propensities along the counterfactual treatment path. While this allows partial adjustment for time-varying confounding, the method does not correct for spatial interference and is sensitive to small propensity values, which can lead high variance.
first train the
Training Details. We randomly initialize all model parameters (GST-UNet and baselines) with Xavier uniform weights [17]. We the Adam optimizer [2 with initial learning rate, halving it whenever the validation loss plateaus for a specified scheduler patience. To mitigate overfitting, we adopt early stopping when the validation loss fails improve for specified early stopping patience epochs. Validation uses 40 of the 190 training prefixes, and the total training is capped at 100 epochs. We tune the following hyperparameters: (i) batch size, learning rate, scheduler patience, and early stopping patience (common to all models); (ii) for GST-UNet, the curriculum period and learning rate for curriculum phases, the U-Net output dimension dn, and the number and width of hidden layers in the feed-forward G-heads. Table 3 lists the hyperparameter ranges considered, with the values yielding the best validation performance in bold.
Evaluation Procedure. We evaluate each model by averaging the root mean square error (RMSE) of the estimated CAPOs against ground truth across 50 test trajectories. Table i the main text reports RMSE + standard deviation for horizon lengths T E {5. 10) and 3, E (0. 0.5, 1.0, 2.0}.
Effect of Varying T. We ran additional simulations varying the trajectory length T E (100, 200, 400, 600, 800) and 19, 2.0 while keeping the grid size fixed (d, = dy 64). Re- sults are shown Appendix D.I GST-UNet consistently improves with more data, while the baselines
20

---

Daily Respiratory Hospitelizations
Weelly Respirstor licepttllzzsttoss
Woekly Respiratory Hesgitalization Rate
PMaa Leveis Camp
Averegg Daily PMC viggemn Laowls
Carr Fire
Camp
42.
Daily Respliatur Illnnss
Campé
Fire
(cases gor 10,0003
1.4
(casze pas 10 10,000)
4.2
Latzude
33
4.1-
20
3.8
4.0-
34-
May
Jun
Jul Aug Sep Oct Nov Dec Months 2018
20
25
30
35
40
45
-124 -122 -120 118--16--11
Week in 2018
Longstude
Figure 5: (Left) Daily respiratory illness incidence (cases per 10,000). (Center) Weekly aggregated incidence. (Right) Average daily 2M2.5 during the Camp Fire.
PMaa Leveis in Counties,
PMaa Levels Unnoegnlltted Grien,
80
40-
40
50
Latmodo
36-
36
30
-124 -122
-120
-118
-16
-114
-124
122
2020
118
135
Longitude
Longttude (9)
Figure 6: An example of county-level (leff) vs. grid-interpolated (right) PM2s levels on November 18 (during the Camp Fire). The grid interpolation produces a 40 x lattice of wreghteighted estimates aligned with our spatiotemporal framework.
remain biased even as increases. This highlights the importance of adjusting for time-varying confounding: without there is persistent asymptotic bias.
Effect of Neighbor Aggregation. To evaluate the importance of spatial spill-over modeling, we ablate the convolutional kernel used in the ConvLSTM encoder. Table [ compares GST-UNet with standard 3 x kernel against variant that removes neighbor aggregation by using a kernel. Across all levels of confounding strength (31), performance deteriorates markedly when neighbor information is excluded, with RMSE increasing by 30-40%. This confirms that explicitly aggregating information from nearby locations is essential for capturing spatial interference and achieving unbiased counterfactual estimates.
D.2 Wildfire Application
Data Preprocessing and Interpolation. We analyze daily, county-level data from Letellier al. 23 that include PM2.s (particulate matter < 2.5 jum) hospitalization counts for respiratory and cardiovascular conditions, and weather variables (temperature, precipitation, humidity, radiation, wind), plus population estimates from the California Department of Finance. Our study period spans weeks 20-48 (May 18-December 2, 2018), covering both the Carr and Camp fires. As illustrated in Figure[3 daily and weekly aggregated respiratory illness rates rise around these events, while PM2.5 levels also surge during the Camp Fire.
To align with our spatiotemporal framework, we use geopandas 19 to interpolate county-level covariates, PM2.s. and hospitalizations onto latiuue-longitude grid from 32" N to 42"N latitude and -125 to -114 longitude, at a resolution of 0.255. Each grid cell's values are an area-weighted average of the counties it intersects, yielding 40 44 spatial lattice. We mask out non-Califee coli-allfornia cells by setting them to zero, thus obtaining a consistent dataset for further analysis. Figure [ illlustraees how the county-level data compare to the interpolated grid for PM25 November T8.
Model Training and Validation. We train GST-UNet, UNet+, STCINet, and IPWUNet with prediction horizon of T 10 days. All models shared set of hyperparameters: batch size 4.
21

---

Table Estimated county-level increases respiratory ED visits attributable to the wildfire event. with 95% bootstrap confidence intervals. Population is reported in units of 10,000. Counties marked with * have smaller populations, which may lead to greater uncertainty.
County
Mean
2.5%
97.5%
Population (x 10*)
Interval Width / Population
Tehama
37
-126
158
6.4
44,4
Butte
168
30
325
23.0
12.8
Glenn
-52
-262
39
2.8
107.6
Colusa*
13
-158
107
2.1
124.0
Sutter
-18
-170
70
9.6
24.9
Napa
81
-41
192
13.9
16.8
Lake
6.4
41.8
Solano
38
-79
173
44.6
5.6
Sacramento
202
-107
484
153.9
3.8
learning rate = 5 x 10-4 scheduler patience = 5, early stopping patience = 10, and curriculum 10 10- For GST-UNeet, we additionally set the
period (with curriculum learning rate U-Net output dimension to d,- = 16, the G-head hidden layer size to 8, and the number of G-head layers to We optimize a mean squared error (MSE) loss with two adjustments:
(1) we mask grid cells outside California to exclude them from the loss computation, and (2) we apply cell-specific weights proportional the number of grid cells per county avoid bias toward geographically larger counties. Hyperparameter tuning and validation are performed using data from the first 50 days of the wildfire season. Using the selected configuration, we generate counterfactual predictions for the Camp Fire peak period (November 8-17) by iteratively applying each trained model with increasing history lengths. We note that counties with populations below 20,000-30,000 can yield unreliable incidence rate estimates (baseline daily rates of approximately 4 cases per 10,000 individuals); in counties with hatched markings.
Figure 3
we denote these
Bootstrap Confidence Intervals. We compute 95% bootstrap confidence intervals for all models using 40 bootstrap samples, balancing statistical rigor with computational load. Counties with populations below 20,000-305 000000,0000 tend to yield unstable incidence rate estimates, driven by low baseline daily counts (approximately cases per 10,000), and are excluded from the analysis. These counties are indicated with hatching in Figure 3 choice further supported by the bootstrap results. In Table we report bootstrap intervals for the counties closest to the Camp Fire. Glenn and Colusa cxhibit disproportionately wide intervals- population sizes-and this further justifies their exclusion from the final analysis.
the uncertainty introduced by their small
E Limitations and Broader Impacts
Linitations. While GST-UNet demonstrates strong empirical performance and theoretical grounding. several limitations should be acknowledged. First, our method relies on standard causal identification assumptions, including no unobserved confounding (Assumption [ which inherently untestable and may not hold in all real-world settings. Second, our framework assumes the existence of a time-invariant representation
the spatiotemporal process (Assumption 2 useful but idealized condition that may be violated in domains with highly non-stationary or regime-shifting dynamics. Finally, GST-UNet designed for gridded spatiotemporal data and assumes regular spatial lattice; while this is common in environmental and health applications, adapting the framework to irregular spatial structures (e.g., graphs or administrative boundaries) is an important direction for future work.
This work advances machine learning by introducing a spatiotemporal causal inference framework for estimating treatment effects in complex real-world settings. The GST-UNG has broad applications in public health, environmental science, and social policy, where understanding interventions supports wvidence-based decisions. For example, can inform pollution control, wildfire response, or health resource allocation. However, like all observational methods, GST-UNet depends on the quality and completeness of the data, as well as the assumptions stated in this work. We caution against uncritical use in high-stake settings, as violations of model assumptions or data biases can lead to misleading conclusions. We encourage responsible setooymenttsspeiilll in contexts affecting vulnerable appulations-and recommend pairing our method with domain expertise, sensitivity checks, and uncertainty quantification.
Broader Impacts.
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes]
Justification: The abstract and introduction state that the paper introduces GST-UNet, a neural framework for estimating causal effects in spatiotemporal settings with time- varying confounding. They accurately summarize the paper's three core contributions: the development the framework with theoretical guarantees, empirical validation in synthetic settings, and real-world application to wildfire smoke exposure. The claims align with both the theoretical results and empirical findings presented in the main text.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: We present the limitations of our work in Appendix Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
section The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms
and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We present two theorems, Theorem D and Theorem D which cover the identifi- cation and consistency guarantees of our estimator. The (complete and correct) proofs are
23

---

included in Appendix B and Appendix C respectively. The assumptions are incorporated in Assumption as well as the text of the theorems.
Assumption II
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented
by formal proofs provided in appendix or supplemental material. Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section 6 and Appendix D provide the information necessary (including data generation processes, model choices, training and validation procedures, hyperparameter choices, etc.) to reproduce the main experimental results.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully specific model and empirical evaluation, may
might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
24

---

Justification:
We provide the replication data and code at https:://ithhb. com/ moprescu/GSTUNee along with instructions for reprodueibility (see README, docu- ment).
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https wwwee//ggu////eeeeeeeeee www.eeeuwumisssooooslsesee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes] Justification: Section 6 and Appendix D Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper to level of detail that is necessary to appreciate the results and make sense of them. The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the standard deviation for the synthetic experiments in Table and boostrap CIs for the real-world case study in Section 6.2
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
25

---

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative
error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: See Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttps:/nneuris....
Answer:
[Yes]
Justification: We have reviewed the ethics guidelines at attttpp:///ourrisccccpbbb/
Ethi cecucceGulelini ines and confirm that our work adheres to them.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid-
eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed? Answer: [Yes]
Justification: See Appendix Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal
impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks,
26

---

mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The models and datasets used in this work do not pose significant risks for misuse. GST-UNet is a causal inference framework for scientific and policy applications using structured spatiotemporal data. Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification; All datasets and code used in this work are publicly available and appropriately cited (e.g., Letellier et al. (23 for wildfire data). License and usage terms are respected as per the original sources.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasete has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
We release code and simulation scripts to reproduce all experiments, along with documentation and usage instructions (see Appendix D). All assets are anonymized for review and available at the provided URL.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
27

---

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs were used solely for writing assistance and code debugging; they were not involved in the development or implementation of the core methodology.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy [ttppe///eerrrppp.c/ooooeeeeerrr http://////oooeeeeee////////// for what should or should not be described.
28


Reviews:
[{"review_id": "nDpHyh982P", "reviewer": "Reviewer_GdrS", "summary": "The paper introduces GST-UNet, an end-to-end framework for estimating potential outcomes and treatment effects in spatiotemporal settings that exhibit interference, spatial confounding, temporal carry-over, and time-varying confounding. GST-UNet uses a U-Net and a ConvLSTM to learn spatiotemporal representations, and a set of G-heads that implement G-computation to adjust for confounders. A multi-task loss lets all heads and the encoder train stably. The authors provide a proof of identification (Theorem 1) under standard assumptions plus a representation-based time-invariance assumption, and they also provide a proof consistency in the appendix (theorem 2). Experiments on synthetic data show superior performance of GST-UNet compared to the baselines used in the paper. Additionally, a case study on the health outcomes of wildfire further demonstrates the utility of the model on an empirical example.", "questions": "Q1. Generally, the weaknesses listed each pose a question, but I would like to repeat the question referring to w1: What is exactly the basis of the claim on the method being suitable for data-scarce regimes?\n\nQ2. Do you think a network causal inference method could be one to compare against, given that the spatial information can be modeled by modeling the input graph of the treatment unit as a grid (please correct me if I’m missing something)?", "limitations": "This is not the shortcoming of the paper, but similar to essentially all causal inference methods, lack of real-world data for evaluation is a significant limitation, which makes a stronger theory or careful experimental design more important.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n- The authors properly identify an important gap in representation learning for spatio-temporal  causal inference, and the paper takes important steps in addressing this. \n- There is significant effort put into a careful design of the framework, to address this gap. The paper shows significant potentials. \n- The paper shows theoretical results, careful methodology, and experiments, including a simulation for testing the models (given the natural limitations of empirical validation in causal inference).\n- The authors utilize a well-suited training strategy using curriculum learning to overcome challenges in training G-computation heads simultaneously on noisy recursively-generated targets.\n\n---\n\n**Major Weaknesses:**\n\nTo my understanding the main shortcoming of the paper is due to w3 and w4 below, which make it difficult to set clear expectations and settle on a criteria to assess whether the paper is ready to be published or needs substantial improvement (hence a later resubmission).\n\nw1. The paper claims the proposed method is particularly suitable for data-scarce regimes, but I find it hard to see an appropriate discussion of that in the main paper. Are the authors claiming this is implied from their model architecture (e.g. using UNet)? If so, I don’t think that is enough to emphasize on the suitability of the method for data-scarce regimes.\n\nw2. Why is theorem 2 in the appendix? I think the main paper contains quite a lot of unnecessary details that could be deferred to the appendix and replaced with theorem 2, especially given that the authors claim that theorem 1 and 2 are the main theoretical contributions (identification and consistency).\n\nw3. Is the paper primarily proposing a rather general framework supported by promising theory that could serve as a first step for developing spatio-temporal causal inference tools, or are they proposing the tool itself? The paper seems to swing between the two. I doubt that claiming both is a suitable strategy, because in that case the results are doubly insufficient. \n\nw4. Related to w3:\n - a. I find the paper’s presentation unclear for a theoretical paper that proposes a method based on the theory as a stepping stone for future tools to advance spatio-temporal causal inference.\n - b. I also find the experiments insufficient for a paper that proposes a ready to use tool. There are several heuristic/engineering choices in the design of the framework, which need a more thorough and better designed set of experiments to evaluate. e.g., one expects an ablation study, or at least experiments that evaluate how would alternative choices (e.g., replacing UNet with an alternative) impacts the performance.\n \nAs I was reading the paper, I initially got the impression that the paper is of the former type (theory-inclined), but later on my impression started to swing towards the latter (ready to use tool) as I read through extensive implementation details (overly elaborate for a theoretical paper), did not find theorem 2 in the main paper, and the elaborate discussion of the experiments.\n\n\n---\n\n**Minor Weaknesses:**\n\nm1. I found the presentation somewhat confusing. The paper jumps between main concepts and remarks or delves into details so frequently that it becomes distracting from the main contributions and choices. Ideally, it should be clear to the reader which choices are fundamental and of conceptual importance, and which ones can be replaced with an acceptable alternative. \n\nm2. I am not familiar with the literature on spatio-temporal causal inference, so I cannot tell if there are other benchmarks or datasets or not. If there are, it would substantially improve the paper to include them. Otherwise, perhaps a more careful discussion of the simulation (why this specific data generating process, what if a different one is used, etc.) and ideally experiments on more simulations could clarify the utility of the proposed method.\n\n---\n\n**Note:** I believe the paper shows significant merits and potentials (see the strengths). I am just not yet convinced that the current draft is ready to properly show and advocate for these merits, due to the reasons discussed above.", "was_revised": true, "final_justification": "As I explained in my initial review and my comment on the rebuttal, I acknowledge the paper's merits in addressing an interesting and relevant problem through a suitable approach and methodology, and with novel contributions. I initially had concerns regarding some of the claims (e.g., why this is suitable for data-scarce regime) and how the contribution was presented (e.g., a main theoretical result being in the appendix, while the main paper thoroughly discussed details of the heuristics). I also believe the initial submission jumps back and forth between discussing the main concepts and presenting conceptual and theoretical arguments, versus extensive implementation details (which again, involve several heuristic choices). The authors' rebuttal however seems convincing to me and their revision plan is reasonable. Of course, to be fully confident about whether the issues will be resolved in the camera-ready version, one needs to see the revised paper, but I believe that is not how this review process works and we need to present an assessment according to the information in hand (that is also the fair way to assess the authors' work; they clearly have put a substantial amount of effort into preparing a convincing rebuttal and a thoughtful revision plan). \n\nTL;DR: I had some concerns and questions, which initially brought me to recommend a weak reject despite recognizing the motivation and the merits of the paper. The rebuttal answered my questions. The authors' promised revision is critical for accepting the paper, but I find their plan reasonable and convincing. In conclusion, I am raising my score to an accept score.", "rebuttal": "**Strengths**\n\nThank 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.\n\n**Re w1: \"Data-scarce regimes\"**\n\nThank you for raising this point. GST-UNet operates in an intrinsically data-scarce setting:\n\n* **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.\n\n* **Inductive-bias-driven efficiency.** With so few samples, performance hinges on a backbone whose built-in biases match the physics:\n\n    * 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).\n\n   * 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.\n\n* **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.\n\nWe will add a short paragraph in Section 4 summarising the argument above and include the ViT ablation to make the data-scarce claim explicit.\n\n**Re w2: Theorem 2**\n\nThank 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.\nWe 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.\n\n**Re w3: Framework vs tool**\n \nOur goal is both to close a *theory gap* and to deliver a *practical recipe*—but the two pieces are tightly coupled, not competing claims.\n\n* **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.\n\n* **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.\n\nThus 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).\n\n**Re w4 & m1: Presentation and Experiments**\n\nPoint well taken re the presentation. Thus, we propose to reorganize the camera-ready version as follows: \n\n* Promote Theorem 2 (consistency) to the main text and move low-level implementation details to the appendix\n\n* Add a concise roadmap that tags each element as a *core principle* or a *replaceable implementation/engineering choice*, preventing concept–detail jumps.\n\nWhile the current paper already ablates curriculum training and the attention mechanism, we commit to strengthen the evidence for the current implementation by:\n\n* adding a TinyViT [7] backbone ablation to test robustness to encoder choice, and\n\n* 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]).\n\nWe 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.\n\n**Re m2 & Q2: Benchmarks & Network causal inference**\n\nOur synthetic study emulates pollutant advection–diffusion on a 2-D lattice:\n\n* **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.\n\n* **Time-varying confounding.** Past emissions alter local $tX_{s,t}$​, which in turn influence future $A_{s,t}$ choices, mirroring real regulatory feedback loops.\n\n* **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.\nThus, while simplified, the setup is grounded in a real physical process.\n\nAt 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.\n\nAdditionally, 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:\n\n* **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.\n\n* **Learned exposure function:** GST-UNet infers the spillover kernel directly from data rather than prespecifying it.\n\n* **Spatial confounding:** network methods typically ignore latent spatial fields (e.g., weather) that confound both treatment and outcome.\n\nAdapting 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.\n\nWith 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.\n\n**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.**\n\n[1] M. Tec et al. Weather2vec: Representation learning for causal inference with non-local confounding. AAAI, 2023.\n\n[2] S. Ali et al. Estimating direct and indirect causal effects of spatiotemporal interventions with spatial interference. ECML PKDD, 2024.\n\n[3] G. Ayzel et al. RainNet v1.0: A CNN for radar-based precipitation nowcasting. Geosci. Model Dev., 2020.\n\n[4] J. C. Fernández et al. Deep coastal sea forecasting using UNet-based models. Knowl.-Based Syst., 2022.\n\n[5] J. Ring et al. Prediction of sea surface temperature using U-Net. Remote Sens., 2024.\n\n[6] A. Dosovitskiy et al. An image is worth 16x16 words: Transformers for image recognition. arXiv:2010.11929, 2020.\n\n[7] K. Wu et al. TinyViT: Fast pretraining distillation for small vision transformers. ECCV, 2022.\n\n[8] M. Tec et al. SPACE: The spatial confounding environment. ICLR, 2024.\n\n[9] G. Papadogeorgou et al. Causal inference with spatio-temporal data. J. R. Stat. Soc. B, 2022."}, {"review_id": "tOH2NfOh9q", "reviewer": "Reviewer_XYsK", "summary": "The paper proposes the GST-UNet, a neural method for estimating conditional average potential outcomes over time. Therein, the authors combine interative G-computation with a neural backbone to model spatial-temporal dependencies, which is different to existing neural methods. The method is validated on synthetic and real world data.", "questions": "Can the authors ellaborate more clearly on their methodological contribution as compared to existing works? Concrete examples, such as specific changes in architecture, learning algorithm, theory etc would be helpful. I am happy to increase my score if the authors can provide convincing arguments.", "limitations": "Well addressed in Appendix E.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 1, "strengths_and_weaknesses": "Strengths:\n\n- The paper is a complete piece of work. Assumptions and proofs are provided, and the theoretical background is well explained.\n- The task of interest is challenging and relevant in a variety of applications, including personalized medicine.\n- The paper is well written and the complex methodology is well illustrated.\n\nWeaknesses:\n\n- My main point of concern is the novelty of the proposed novelty as compared to existing work. The proposed method is very close to existing works [1, 2, 3, 4] that leverage G-computation within neural architecture. Specifically, [3, 4] seems to use the exact same iterative regression scheme (generating pseudo-outcomes and learning them in an iterative manner to estimate CAPO/ CATE). Further, from a high-level perspective, [3] seems to have a similar neural architecture (neural backbone + G-computation heads) and proposes a very similar learning algorithm. [1, 2] leverage very similar approaches for estimating ATEs over time, with an additional targeting step to remove first-order bias. At the very least, the authors of the paper under review should acknowledge these existing works where appropriate (e.g., in the architecture, the learning algorithm, etc.) and point out key differences more clearly.\n- From my understanding, the GST-UNet only differs from [1] in its neural backbone, which is a U-Net consisting of ConvLSTM layers. The proposed backbone itself, however, is also not novel. Hence, the proposed paper appears to be a simple combination of existing works.\n- It remains unclear how dependencies between samples (interference) can be accounted for with the proposed backbone. The only dependencies that are additionally accounted for are spatial, but not the dependencies between units (i=1,...n), which is typically referred to as interference.\n- In the experimental section, no neural state-of-the-art methods for CAPO estimation over time are used as baseline (e.g., [3, 4, 5, 6]), which seems insufficient since one of the main claims is that these baselines cannot handle spatial dependencies.\n\n\nReferences\n\n[1] Frauen et al. Estimating average causal effects from patient trajectories.\n[2] Shirakawa et al. Longitudinal targeted minimum loss-based estimation with temporal-difference heterogeneous transformer.\n[3] Hess et al. G-Transformer for Conditional Average Potential Outcome Estimation over Time.\n[4] Xiong et al. G-Transformer: Counterfactual Outcome Prediction under Dynamic and Time-varying Treatment Regimes.\n[5] Lim et al. Forecasting Treatment Responses Over Time Using Recurrent Marginal Structural Networks.\n[6] Bica et al. Estimating Counterfactual Treatment Outcomes over Time Through Adversarially Balanced Representations.", "was_revised": true, "final_justification": "I believe this is a borderline decision (see my response to the authors for details).", "rebuttal": "**Strengths**\n\nThank 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.\n\n**Re: Novelty and backbone comparison vs. prior neural G-computation work**\n\nThank 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:\n\n| **Dimension** | **Prior neural G-comp** | **GST-UNet (ours)** | \n|-----|-----|-----|\n| Data structure & assumptions |  *Single spatiotemporal trajectory*—e.g., one 64 × 64 grid set of $(A_{s,t}, X_{s,t}, Y_{s,t})$ observed for $T$ days. Every frame is correlated with its neighbours *in space* (shared weather patterns) and with previous frames *in time* (smoke today affects pollution tomorrow). The target CAPOs themselves are *spatially structured surfaces*—we seek $ \\textbf{CAPO}_{s}(a)$ at *every* grid cell, not a single scalar per individual. Because of this tight coupling we cannot break the data into $64 × 64$ independent time-series or assume thousands of separate “videos.” To learn with only $T-\\tau$ usable prefix windows, we introduce **Assumption 2 (representation-based time-invariance)**: if a history embedding $\\phi({H}\\_{1:t})$ captures the latent state, then different prefixes behave *as if* they were i.i.d. samples, enabling consistent G-computation in this data-scarce, highly correlated regime. | Independent patient/user time series; no spatial dimension. Outcomes depend only on each unit’s own history, so interference and spatial confounding are out of scope and only temporal relationships are modeled. Identification (when given) relies on i.i.d. sampling across units. |\n|Backbone & learning algorithm | ConvLSTM-UNet encoder couples space + time and leverages learned convolutional kernels that aggregates neighbour treatments/covariates, capturing both direct/spillover effects and spatial confounding. Must be sample-efficient because only $T-\\tau$ data prefixes (quasi-i.i.d samples) are available, where $T$ is the length of the spatiotemporal sequence. Curriculum training stabilises the iterative pseudo-outcome recursion, needed due to the complex data structure and scarce data. | Architectures assume arbitrary numbers of i.i.d. trajectories, allowing expressive transformers or RNNs that model *temporal* dependence only. No mechanism for neighbour treatments and covariates; no curriculum for iterative G-comp., most likely due to the simpler nature of the data structure and the larger number of i.i.d. samples |  \n| Theoretical results | Unlike previous works (e.g. the STCINet in our experiments), we state the minimally sufficient requirements (Assumption 2) for the identification results (Theorem 1) to hold. We additionally provide consistency of CAPO estimates under the single-trajectory/interference regime (Theorem 2) — new in the G-computation literature.| No prior neural G-comp. work (to our knowledge) provides identification or consistency  guarantees when only one spatiotemporal sequence is observed.| \n\nAs 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.\n\nFinally, 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. \n\nWhile 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.\n\n**Re: Interference** \n\nIn 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$:\n\n| Kernel Size | $\\beta_1 = 0.0$ | $\\beta_1 = 0.5$ | $\\beta_1 = 1.0$ | $\\beta_1 = 1.5$ | $\\beta_1 = 2.0$ |\n|-------------|------------------|-------------------|-------------------|-------------------|-------------------|\n| \\(k = 3\\)    | 0.329 ± 0.004     | 0.348 ± 0.004      | 0.402 ± 0.005      | 0.440 ± 0.004      | 0.404 ± 0.005      | \n| \\(k = 1\\)    | 0.532 ± 0.004     | 0.547 ± 0.005      | 0.537 ± 0.005      | 0.599 ± 0.007      | 0.636 ± 0.006      | \n\nThe  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.\n\n**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.**"}, {"review_id": "V1hDXO4tKs", "reviewer": "Reviewer_sKpM", "summary": "The paper introduces GST-UNet, an end-to-end neural framework for spatiotemporal causal inference with time-varying confounding. A U-Net + ConvLSTM + attention encoder learns a history-invariant representation, while a stack of \\tau “G-heads” implements regression-based iterative G-computation to estimate conditional average potential outcomes (CAPOs). A curriculum schedule gradually activates earlier heads so that pseudo-outcomes become reliable. The authors prove identification (Theorem 1) and consistency (Theorem 2) under standard sequential ignorability plus a weaker “representation-based time-invariance” assumption.", "questions": "1. Curriculum schedule sensitivity: How was the “ease coefficient” chosen? Have you tried other schedules, and does performance vary markedly?\n2. Head-count vs. horizon: Your implementation fixes the number of G-heads \\tau equal to the longest horizon evaluated. What happens if \\tau is over- or under-specified? A small ablation (e.g., \\tau ∈ {4, 6, 8, 10}) would help increase robustness.\n3. Section 5.4 mentions future work on graphs, but GST-UNet currently assumes a regular grid for convolutions and skip connections. Can you outline concretely how you would replace the ConvLSTM encoder with a graph-message-passing block, and whether any of the proofs need modification?\n4. Figure 4 reports training time on 64×64 grids; how does memory usage and wall-clock scale for inputs with more grids (e.g., 256×256 or 512×512), and for longer \\tau?\n5. For deployment in, say, wildfire smoke early-warning, what is the per-timestep inference latency (ms per grid cell) on an A100 and a CPU?\n6. Your synthetic DGP uses a Gaussian smoothing kernel for spill-over. Does GST-UNet require the same spill-over kernel at train and test time?\n7. Have you inspected the temporal-attention weights or produced saliency maps to verify that the model focuses on plausible meteorological or demographic features?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. First architecture that unifies deep spatiotemporal representation learning with iterative G-computation, handling interference, carry-over, and evolving confounders without user-specified exposure mappings.\n2. Application to wildfire smoke quantifies health burden with spatial granularity—aligning with epidemiological literature—showing practical utility.\n\nWeaknesses:\n1. Results depend on one environmental-health example; an additional domain (e.g., mobility interventions, agriculture) or a cross-region generalisation split would strengthen external validity.\n2. GST-UNet currently assumes a regular lattice; the authors mention future work on irregular graphs, but no experiment shows feasibility.\n3. Only one \"ease coefficient\" value is reported; how did you choose this? A hyperparameter sensitivity analysis would be more helpful.\n4. While runtimes are modest on A100, memory/latency for larger \\tau or higher-resolution grids is not reported.\n5. No empirical diagnostics are offered for the representation-invariance or no-unobserved-confounding assumptions.\n6. GST-UNet does not visualise which spatial features drive G-head predictions; attention heat-maps or Shapley scores over the grid could help.", "was_revised": true, "final_justification": "The author's detailed rebuttal has addressed all of my questions and made me understand the paper better; therefore, I decided to raise my rating from 3 to 4.", "rebuttal": "**Strengths**\n\nThank 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.\n\n**Re: Domain breadth** \n\nWe 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.\n\n**Re: Regular grid vs. graph extension**\n\nGST-UNet’s encoder is **modular** and thus it can be swapped for a graph–message-passing variant almost by \"find-and-replace\":\n\n| Grid Component   | Graph Analogue (Irregular Mesh) | Implementation Details|                                                                                        \n|---------------------|-----------------|-------------------|\n| 2-D conv in ConvLSTM gates | Graph convolution (GCN, GraphSAGE, or GAT [2-4]) on the static adjacency $N(v)$                         | Replace every `Conv2d` call with a graph-conv layer from PyTorch Geometric; weight sharing still enforces locality. |\n| ConvLSTM temporal recursion        | Graph-ConvLSTM / Graph-GRU ([5,6])   | Simply swap the 2-D conv inside each gate with the graph conv above; hidden-state shapes remain.|                                                       \n|UNet | Graph U-Net hierarchy with top-k pooling/unpooling and skip connections [7] | Use the Graph U-Net pooling operator in [7]; pool ratio mirrors the $2\\times$ grid stride, and skip tensors are passed as in the original U-Net.|\n\nMoreover, *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.\n\n**Re: Curriculum \"ease coefficient\" & schedule sensitivity**\n\nAs 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,\n\n* **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.\n\n* **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.\n\nThe middle value strikes a balance between denoising pseudo-outcomes and keeping the shared encoder adaptable throughout training.\n\n**Re: Head-count ($\\tau$) vs. prediction horizon**\nEach G-head $k$ learns the conditional outcome function\n\n$$ Q_k \\bigl(h\\_{1:t+\\tau-k}\\bigr) =\n\\mathbb{E}\\left[Y\\_{t+\\tau} \\mid| \\mathbf{A}\\_{t+k:t+\\tau} = \\mathbf{a}\\_{k:\\tau},  \\mathbf{H}\\_{1:t+\\tau-k} = \\mathbf{h}\\_{1:t+\\tau-k} \\right].$$\n\ni.e., the response when exactly $k$ steps remain in the horizon. Consequently, a model intended for horizon length $\\tau$ needs exactly $\\tau$ heads so that every offset $k=1,\\dots,\\tau$ is covered. If we were to train with fewer heads than the test horizon, the recursion would halt prematurely; using more heads than necessary would leave the earliest heads untrained and yield arbitrary outputs.\n\n**Re: Time/ Memory profiles for larger grids/ $\\tau$**\n\nMemory usage and wall-clock for training with $\\tau=5$:\n\n| Grid (H × W)   | Cells   | Peak GPU Memory | Train Time |\n|----------------|---------|------------------|-------------|\n| 64 × 64        | 4,096   | 1.15 GiB         | 1.80 min    |\n| 128 × 128      | 16,384  | 4.43 GiB         | 2.33 min    |\n| 256 × 256      | 65,536  | 17.59 GiB        | 7.59 min    |\n\nA forward/back-prop pass scales as\n $\\mathcal{O}\\bigl(N_{xy}\\,\\tau\\,d^{2}\\bigr)$ —linear in the number of spatial cells $N_{xy}$​ and the horizon $\\tau$ (with a quadratic dependence on channel width $d$). The near-linear growth in both memory and runtime shown above matches this prediction.\nRuns for 512 × 512 grids and for horizons $\\tau>10$ are queued; their results will be included in the camera-ready version.\n\n**Re: Diagnostics & interpretability (attention / saliency)**\n\nGreat suggestion. GST-UNet does not use temporal self-attention, an attention gate on the U-Net skip connections—so there are no time-axis attention weights to plot. Instead, we will be generating (i) Grad-CAM saliency maps to highlight which grid cells influence each G-head and (ii) visualisations of the learned spill-over kernel to show how neighbour treatments affect outcomes. These interpretability plots are underway and will be included in the camera-ready (or sooner if ready during the discussion).\n\n**Re: Inference-time latency on GPU and CPU**\n\nIn the wildfire simulation, the inference time is 0.0041 ms/ grid cell with the A100 GPU and 0.011 ms/ grid cell with Perlmutter’s AMD EPYC CPU.\n\n**Re: Spill-over kernel mismatch**\n\n**Clarification.** Our synthetic DGP uses a weighted mean of the neighbors, given by an asymmetric advection–diffusion convolution kernel, not a Gaussian smoother.\n\n**Does GST-UNet need the same kernel size?** No. GST-UNet learns whatever spill-over pattern exists in the data: A small 3 × 3 kernel in the first ConvLSTM, plus the U-Net’s down/upsampling, gives a large effective receptive field. \n\n**We appreciate your thoughtful feedback and hope the clarifications provided show the value of our contributions and justify a higher score.**\n\n[1] S. Ali et al. Estimating direct and indirect causal effects of spatiotemporal interventions with spatial interference. ECML PKDD, 2024.\n\n[2] T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907, 2016.\n\n[3] W. Hamilton, Z. Ying, and J. Leskovec. Inductive representation learning on large graphs. NeurIPS, 2017.\n\n[4] P. Veličković, et al. Graph attention networks. arXiv preprint arXiv:1710.10903, 2017.\n\n[5] Z. Cui et al. Traffic graph convolutional recurrent neural network: A deep learning framework for network-scale traffic learning and forecasting. IEEE Trans. Intell. Transp. Syst., 2019.\n\n[6] K. K. Bhaumik et al. STLGRU: Spatio-temporal lightweight graph GRU for traffic flow prediction. PAKDD, Springer, 2024.\n\n[7] H. Gao and S. Ji. Graph U-Nets. ICML, 2019."}, {"review_id": "3zpgvHtGw5", "reviewer": "Reviewer_zFj6", "summary": "The paper addresses a significant problem in causal inference: estimating causal effects from spatiotemporal observational data under time-varying confounding, interference, and spatial dependencies.", "questions": "What is the time complexity of GST-UNet?", "limitations": "Assumptions may be not practical in real-world applications.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\nThe proposed GST-UNet framework integrates neural representation learning (U-Net + ConvLSTM + attention) with iterative G-computation for bias-adjusted effect estimation.\n\nIt is clearly motivated by practical challenges in public health and environmental policy.\n\nThe method builds on sound causal inference principles, including the potential outcomes framework and G-computation.\n\nWeaknesses\n\nThe assumption of embedding-based time invariance might be strong in some real-world scenarios.\n\nThe comparison methods (e.g., UNet+, IPWUNet) may not be the most state-of-the-art. Why do not compare with causal transformer [1], G-Net [2] and DeepACE [3]?\n\nThe experiments conducted on one dimension time-varying confounder, so it is difficult to assess the performance of GST-UNet on multi-dimensions of time-varying confounders. \n\n\n[1] Melnychuk, Valentyn, Dennis Frauen, and Stefan Feuerriegel. \"Causal transformer for estimating counterfactual outcomes.\" In International conference on machine learning, pp. 15293-15329. PMLR, 2022.\n\n[2] 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-computation for counterfactual prediction under a dynamic treatment regime. In Machine Learning for Health, pages 282–299. PMLR, 2021.\n\n[3] 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, pages 7586–7594, 2023.", "was_revised": true, "final_justification": "Thank you for addressing my main concerns. I have no further comments on this work and will maintain my rating in support of its acceptance at NeurIPS.", "rebuttal": "**Strengths**\n\nThank 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.\n\n**Re: Embedding-based time invariance assumption**\n\nThank 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 $(\\textbf{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.\n\nThis 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.\n\nWe acknowledge that stronger empirical validation of this assumption is an important direction for future work.\n\n**Re: Comparisons with Causal Transformers and related methods**\n\nWe 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.\n\nWe perform a part-for-part comparison between our model and these methods in the table below:\n\n| Dimension                  | Prior neural G-comp methods (e.g. [1–3])                                                                                                                                                                                                 | **GST-UNet (ours)** |                                                                                                                                                                                                                                                                                      \n|---------|---------|------|\n| **Data structure & assumptions** | Independent user/patient time series; no spatial dimension; outcomes depend only on each unit’s own past.  Typical datasets contain tens of thousands of i.i.d. trajectories, so sample size is not a limiting factor.        | Single spatiotemporal chain—one set of $64 \\times 64$ lattices $(A_{s,t}, X_{s,t}, Y_{s,t})$ observed for $T$ time steps. After fixing horizon $\\tau$, the effective sample size is at most $T - \\tau$ quasi-i.i.d. prefixes. Strong space-time coupling breaks down the i.i.d. assumptions and prevents treating each grid cell as its own trajectory. |\n| **Backbone & learning algorithm** | RNN/Transformer models only temporal dependence; no neighbor treatments; no curriculum.                                                                                                                                                 | ConvLSTM-U-Net couples space and time, learns spill-over via a convolutional kernel, and uses curriculum G-computation.                                                                                                                                                                                         |\n| **Theory**                 | No identification/consistency guarantees for one lattice.                                                                                                                                                                                | Assumption 2 + Theorems 1–2 give identification & consistency for a single-trajectory, interference setting.                                                                                                                                                                                                     |\n\nWe 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.\n\n**Re: Higher-dimensional confounders**\n\nThank you for pointing this out. We extended the synthetic study to include **5 time-varying confounders**. The revised spatiotemporal DGP is:\n\nWe 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.\nSpatial convolutions are indicated with “*”. \n\n$$\n\\begin{align}\n\\mathbf{(Covariates)} \\quad X^{(c)}_{t} &= \\alpha^{(c)}_0 + \\alpha^{(c)}_1 X^{(c)}\\_{t-1}\n               + \\alpha^{(c)}_2 A\\_{t-1}\n               + \\alpha^{(c)}_3 (K_X * X^{(c)}\\_{t-1})\n               + \\varepsilon^{(c)}\\_{X,t} \\\\\\\\\n\\mathbf{(Treatment)}\\quad\nA\\_{t} &\\sim \\mathrm{Bernoulli}\\bigl(\\sigma\\bigl(\\beta_0 + \\beta_1 (K_A * \\bar X_t)\\bigr)\\bigr) \\\\\\\\\n\\mathbf{(Outcome)}\\quad Y_t &= \\gamma_0 + \\gamma_1 (K\\_{YA} * A\\_{t-1})\n        + \\gamma_2 (K\\_{YX} * \\bar X\\_{t-1})\n        + \\gamma_3 Y\\_{t-1}\n        + \\varepsilon\\_{Y,t}\n\\end{align}\n$$\n\nWe train our method and the corresponding baselines as in the paper and report the results below ($\\tau=5$):\n\n| Model     | $\\beta_1 = 0.0$   | $\\beta_1 = 0.5$   | $\\beta_1 = 1.0$   | $\\beta_1 = 1.5$   | $\\beta_1 = 2.0$   |\n|-----------|---------------------|---------------------|---------------------|---------------------|---------------------|\n| UNet+     | **0.28 ± 0.01**     | 0.39 ± 0.01         | 0.51 ± 0.01         | 0.59 ± 0.01         | 0.65 ± 0.01         |\n| STCINet   | 0.32 ± 0.01         | 0.39 ± 0.01         | 0.50 ± 0.01         | 0.57 ± 0.01         | 0.71 ± 0.01         |\n| IPWUNet   | 0.56 ± 0.01         | 0.59 ± 0.01         | 0.60 ± 0.01         | 0.58 ± 0.01         | 0.65 ± 0.01         |\n| GSTUNet   | 0.37 ± 0.01         | **0.39 ± 0.01**     | **0.44 ± 0.01**     | **0.40 ± 0.01**     | **0.41 ± 0.01**     |\n\nThese results confirm GST-UNet’s robustness to increasingly strong time-varying confounding, while plain UNet variants and existing spatial models degrade.\n\n**Re: Time complexity of the GST-UNet**\n\n*Analytical form.* A single forward/back-prop pass scales linearly in\n $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)$. \n*Empirical runtimes* (A100 GPU on NERSC’s Perlmutter cluster) for each experiment: \n\n| Task | Grid & $\\tau$ | Time |\n| ----| ---------|--------|\n| Synthetic study | 64 × 64, $\\tau$ = 10 | $\\simeq 11$ min|\n| Wildfire study | 40x44, $\\tau=10$ | $\\simeq 5$ min|\n\n**We hope these clarifications are helpful and further strengthen your confidence in the work.**"}, {"review_id": "qxL4WkTqOF", "reviewer": "Reviewer_RWWK", "summary": "The paper proposes the GST-UNet, a neural method for estimating conditional average potential outcomes (CAPOs) over time. Therein, the authors use an iterative approach to G-computation. Further, they incorporate spatial dependencies in their neural backbone, which is different to existing neural methods. The method is validated on synthetic and real world data.", "questions": "Questions:\n\n- What are the key differences to works such as [1] and [2]?\n- How can interference between units be accounted for?\n- Can the authors add more experimental results against neural baselines that are tailored for CAPO estimation over time?", "limitations": "The authors discuss limitations regarding their assumptions in Appendix E.", "rating": 2, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 1, "strengths_and_weaknesses": "Strengths:\n\n- The paper is a complete piece of work. Assumptions and proofs are provided, and the theoretical background is well explained.\n- The task of interest is highly relevant.\n\nWeaknesses:\n\n- The proposed method is very close to existing works [1,2]. Specifically, [1] uses the exact same iterative regression scheme (generating pseudo-outcomes and learning them in an iterative manner). Further, from a high-level perspective, [1] has the same neural architecture (neural backbone + G-computation heads) and proposes the exact same learning algorithm. Further, [2] has a closely related TMLE approach for estimating ATEs over time, with an additional targeting step to remove first-order bias. At the very least, the authors of the paper under review should acknowledge these existing works where appropriate (e.g., in the architecture, the learning algorithm, etc.) and point out key differences more clearly.\n- From my understanding, the GST-UNet only differs from [1] in its neural backbone, which is a U-Net consisting of ConvLSTM layers. The proposed backbone itself, however, is also not novel. Hence, the proposed paper appears to be a simple combination of existing works.\n- It remains unclear how dependencies between samples (interference) can be accounted for with the proposed backbone. The only dependencies that are additionally accounted for are spatial, but not the dependencies between units (i=1,...n), which is typically referred to as interference.\n- In the experimental section, no neural state-of-the-art methods for CAPO estimation over time are used as baseline (e.g., [1,3,4,5]). Given that one of the main claims is that these baselines cannot handle spatial dependencies, the experimental section seems insufficient.\n\n\n___\n\n[1] K. Hess, D. Frauen, V. Melnychuk, and S. Feuerriegel. G-transformer for conditional average potential outcome estimation over time. arXiv preprint arXiv:2405.21012, 2024.\n\n[2] Yi Shirakawa, Toru; Li, Yulun Wu, Sky Qiu, Yuxuan Li, Mingduo Zhao, Hiroyasu Iso, and Mark van der Laan. 2024. Longitudinal targeted minimum loss-based estimation with temporal-difference heterogeneous transformer. In ICML.\n\n[3] Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Causal transformer for estimating counterfactual outcomes. In ICML, 2022.\n\n[4] Ioana Bica, Ahmed M. Alaa, James Jordon, and Mihaela van der Schaar. Estimating counterfactual treatment outcomes over time through adversarially balanced representations. In ICLR, 2020.\n\n[5] Rui Li, Stephanie Hu, Mingyu Lu, Yuria Utsumi, Prithwish Chakraborty, Daby M. Sow, Piyush Madan, Jun Li, Mohamed Ghalwash, Zach Shahn, and Li-wei Lehman. G-Net: A recurrent network approach to G-computation for counterfactual prediction under a dynamic treatment regime. In ML4H, 2021.", "was_revised": true, "final_justification": "Thank you for your rebuttal.\n\nThe comparison to existing neural methods for estimating CAPOs made -- to some extent -- sense to me, but the novelty still appears quite limited. Further, without experiments showcasing the performance increase due to mis-specified assumptions in the aforementioned baselines, it is difficult to assess the effectiveness of the GST-UNet. For those reasons, I will maintain my score.", "rebuttal": "**Strengths:** \n\nThank 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.\n\n**Re: Novelty and backbone comparison vs. prior neural G-computation**\n\nPrior 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.\n\n| **Dimension** | **Prior neural G-comp** | **GST-UNet (ours)** | \n|-----|-----|-----|\n| Data structure & assumptions |  *Single spatiotemporal trajectory*—e.g., one 64 × 64 grid set of $(A_{s,t}, X_{s,t}, Y_{s,t})$ observed for $T$ days. Every frame is correlated with its neighbours *in space* (shared weather patterns) and with previous frames *in time* (smoke today affects pollution tomorrow). The target CAPOs themselves are *spatially structured surfaces*—we seek $ \\textbf{CAPO}_{s}(a)$ at *every* grid cell, not a single scalar per individual. Because of this tight coupling we cannot break the data into $64 × 64$ independent time-series or assume thousands of separate “videos.” To learn with only $T-\\tau$ usable prefix windows, we introduce **Assumption 2 (representation-based time-invariance)**: if a history embedding $\\phi(H\\_{1:t})$ captures the latent state, then different prefixes behave *as if* they were i.i.d. samples, enabling consistent G-computation in this data-scarce, highly correlated regime. | Independent patient/user time series; no spatial dimension. Outcomes depend only on each unit’s own history, so interference and spatial confounding are out of scope and only temporal relationships are modeled. Identification (when given) relies on i.i.d. sampling across units. |\n|Backbone & learning algorithm | ConvLSTM-UNet encoder couples space + time and leverages learned convolutional kernels that aggregates neighbour treatments/covariates, capturing both direct/spillover effects and spatial confounding. Must be sample-efficient because only $T-\\tau$ data prefixes (quasi-i.i.d samples) are available, where $T$ is the length of the spatiotemporal sequence. Curriculum training stabilises the iterative pseudo-outcome recursion, needed due to the complex data structure and scarce data. | Architectures assume arbitrary numbers of i.i.d. trajectories, allowing expressive transformers or RNNs that model *temporal* dependence only. No mechanism for neighbour treatments and covariates; no curriculum for iterative G-comp., most likely due to the simpler nature of the data structure and the larger number of i.i.d. samples |  \n| Theoretical results | Unlike previous works (e.g. the STCINet in our experiments), we state the minimally sufficient requirements (Assumption 2) for the identification results (Theorem 1) to hold. We additionally provide consistency of CAPO estimates under the single-trajectory/interference regime (Theorem 2) — new in the G-computation literature.| No prior neural G-comp. work (to our knowledge) provides identification or consistency  guarantees when only one spatiotemporal sequence is observed.| \n\nAs 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.\n\n**Re: Interference between units**\n\nIn 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$:\n\n| Kernel Size | $\\beta_1 = 0.0$ | $\\beta_1 = 0.5$ | $\\beta_1 = 1.0$ | $\\beta_1 = 1.5$ | $\\beta_1 = 2.0$ |\n|-------------|------------------|-------------------|-------------------|-------------------|-------------------|\n| \\(k = 3\\)    | 0.329 ± 0.004     | 0.348 ± 0.004      | 0.402 ± 0.005      | 0.440 ± 0.004      | 0.404 ± 0.005      | \n| \\(k = 1\\)    | 0.532 ± 0.004     | 0.547 ± 0.005      | 0.537 ± 0.005      | 0.599 ± 0.007      | 0.636 ± 0.006      | \n\nThe  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.\n\n**Re: Baselines tailored to CAPO over time**\n\n 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.\n\nWe 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.\n\n**We appreciate your thoughtful feedback and hope the clarifications provided show the value of our contributions and justify a higher score.**"}]

Author rebuttal:

Meta-review: The paper tackles the estimation of Conditional Average Potential Outcomes (CAPOs) from spatio-temporal observational data, facing the challenges of interference among units, spatial confounding, temporal carryover, and time-varying confounding. The difficulty and relevance of the goal w.r.t. epidemiology and earth sciences cannot be over-estimated.

The paper ambitions both to close a theory gap (extending the state of the art to time-varying confounding settings and offering  identification guarantees) and deliver a practical recipe. The rewriting of the paper will be more clearly focused on the theoretical contributions (addressing GdrS remarks). 

Some reviewers (XYsK, zFj6) considered the novelty to be moderate compared to prior neural G-computation work. The authors' rebuttal clarified the difference (compared to a single spatiotemporal trajectory, the proposed approach considers each spatial location as a unit, where the learned convolutional kernel aggregates neighbour treatments/covariates before the G-comp heads), hampering the benchmarking compared to the state of the art (requiring the losses and architectures to be significantly modified). AC: perhaps the naive usage of former methods would nevertheless help to appreciate the merits and novelty of the proposed GST-Unet. 

Other concerns, e.g. related to the data-scarce regime (GdrS) or the computational cost (sKpM) or the possible extension to irregular spatial graphs (sKpM) have been convincingly addressed by the authors' rebuttal.

A great suggestion from sKpM concerns the interpretability/visualisation of the influences: to be added to the revised version of the paper.
Decision: Accept (poster)