NuwaDynamics: Discovering and Updating in Causal Spatio-Temporal Modeling
Authors: Kun Wang
Hao Wu
Yifan Duan
Guibin Zhang
Kai Wang
Xiaojiang Peng
Yu Zheng
Yuxuan Liang
Yang Wang
Venue: ICLR 2024 spotlight (2024)
Keywords: Spatio-temporal data mining
Causal inference
Two-stage framework
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: In this paper, we establish a two-stage spatio-temporal causal concept, NuwaDynamics, which targets to identify causal regions in observe data and endow model with causal reasoning ability.

Abstract: Spatio-temporal (ST) prediction plays a pivotal role in earth sciences, such as meteorological prediction, urban computing. Adequate high-quality data, coupled with deep models capable of inference, are both indispensable and prerequisite for achieving meaningful results. However, the sparsity of data and the high costs associated with deploying sensors lead to significant data imbalances. Models that are overly tailored and lack causal relationships further compromise the generalizabilities of inference methods. Towards this end, we first establish a causal concept for ST predictions, named  NuwaDynamics, which targets to identify causal regions in data and endow model with causal reasoning ability in a two-stage process. Concretely, we initially leverage upstream self-supervision to discern causal important patches, imbuing the model with generalized information and conducting informed interventions on complementary trivial patches to extrapolate potential test distributions. This phase is referred to as the discovery step. Advancing beyond discovery step, we transfer the data to downstream tasks for targeted ST objectives, aiding the model in recognizing a broader potential distribution and fostering its causal perceptual capabilities (refer as Update step). Our concept aligns seamlessly with the contemporary backdoor adjustment mechanism in causality theory. Extensive experiments on six real-world ST benchmarks showcase that models can gain outcomes upon the integration of the NuwaDynamics concept. NuwaDynamics also can significantly benefit a wide range of changeable ST tasks like extreme weather and long temporal step super-resolution predictions.

Published as a conference paper at ICLR 2024
NUWADYNAMICS: DISCOVERING AND UPDATING IN CAUSAL OPATIO-TEMPORT MODELING
Kun Wang Hao Wu-.i, Yifan Duan", Guibin Zhang,, Kai Wang", Xiaojiang Peng", Yu Zheng", Yuxuan Liang"", Yang Wang1.2.3,14
Key Laboratory of Precision and Intelligent Chemistry, University of Science and Technology of China (USTC) Suzhou Institute for Advanced Research, USTC
8 School of Software Engineering, USTC School of Computer Science, USTC
5 Hong Kong University of Science and Technology (Guangzhou) 6 Tongji University
7 Shenzhen Technology University" National University of Singapore JD iCity, JD Technology
(wk520529, wuhao2022, duanyi @ma e.a.l..ste...emee...mmmmmmmmmaae eustc.edu.co bin20038t.ongji. edu. en, Kai Kai.wang@ nus edu. sg
edw.
msyuzheng@out look com pengxiaojianggezze edu. cn yuxliange outlook..oooooooooooooooooooooo com", angyan@ustc. cn"
ABSTRACT
Spatio-temporal (ST) prediction plays a pivotal role in earth sciences, such as meteorological prediction, urban computing. Adequate high-quality data, cou- pled with deep models capable of inference, are both indispensable and prereq- uisite for achieving meaningful results. However, the sparsity of data and the high costs associated with deploying sensors lead to significant data imbalances. Models that are overly tailored and lack causal relationships further compromise the generalizabilities of inference methods. Towards this end, we first es- tablish causal concept for ST predictions, named NuwaDynamics, which tar- gets to identify causal regions in data and endow model with causal reasoning ability in a two-stage process. Concretely, we initially leverage upstream self- supervision discern causal important patches, imbuing the model with gen- eralized information and conducting informed interventions on complementary trivial patches to extrapolate potential test distributions. This phase is referred to as the discovery step. Advancing beyond the discovery step, we transfer the data downstream tasks for targeted ST objectives, aiding the model in recog- nizing a broader potential distribution and fostering its causal perceptual capa- bilities (denoted Update step). Our concept aligns seamlessly with the con- temporary backdoor adjustment mechanism in causality theory. Extensive experi- ments on six real- world ST benchmarks showcase that models can gain outcomes upon the integration of the NuwaDynamics concept. NuwaDynamica also can significantly benefit a wide range of changeable ST tasks like extreme weather and long temporal step super-resolution predictions. Our codes are available at https:/ / gi thub com/ wttp///aaalleaannncceeee.e www.ar/www..o..eee////aaaaaaaaaa http///aaaDDDiiiccccccc
INTRODUCTION
Modern deep learning (DL) approaches have demonstrated promising outcomes in various dynami- cal systems in natural and social science fields like weather forecasting (Schultz et al., 2021; Pathak al., 2022; Bi et al., 2022), rapid fire progression (Tam al., 2022), intelligent transportation (Kaf- fash et al., 2021; Jin et al., 2023). Such astonishing achievements primarily stem from two crucial factors. (1) With the development of computer science, vast amount of data from Earth systems is continuously being acquired (Chen et al., 2022b; Liu et al., 2023a). These ever-growing, massive datasets, with diverse sources, provide the impetus for data-hungry deep models, making learning from data possible. (2) Continual breakthroughs in deep learning algorithms and models enable us effectively adapt to diverse specific scenarios, resulting in state-of-the aff-the-art art performances.
In general, deep learning (DL) provides an efficient optimization framework for automatically and dynamically extracting intrinsic patterns from continuous observable processes. Unlike classical
Yann Wang and Yuxuan Liang are the corresponding authors. i Kun Wang and Hao Wu contribute equally
this paper.

---

Published as a conference paper at ICLR 2024
Sptalo-temporal Dynamical System
Region reconstruction
Aug. 18, 2023
processing
Observations
Tronsfer leornngg
W2
50p
Optimiie
Structure design
Extreme weather prediction
Good fit
Target
Lauee
Suplill
Spors
Deep models
(a) Distribution of the ocean system
(b) An overview of deep dynamic system and its downstream tasks
Figure 1: (a) The distribution of uneven ocean system in the Atlantic Ocean. The points represent different types of systems and we mark some red and yellow areas highlight the imbalance issue of the sensor system. (b) When constructing an ST system, the lack of interpretability often results in limited predictive capabilities for specific scenarios.
dynamic systems, which are primarily derived from first principles (Pryor, 2009; Bürkle al., 2021) and involve high computational costs, DL. approachs often sacrifice an explicit understanding of physical rules. Instead, it resorts large-scale observable data and captures implicit patterns that serve as substitutes for physical laws. Moreover, these patterns can be understood as apatio-temport (ST) correlations, and these deep models can be further regarded as ST dynamical systems.
Though promising, in the context of data-driven dynamical systems research, there still exist some clouds on the horizon. 1) Hfigh-qulityyrrsoouttnn observational data is relatively scarce, and the cost of training with such data is exceptionally high; the distribution of sensors across various regions on Earth is significantly uneven, with many areas unable to benefit from them effectively due to data scarcity (see the example in Fig (aa).. More cases are provided in Appendix A. Although some efforts have been made to address this problem, such as transfer learning ( (Wang et al., 201 8a) and active learning (Ren et al., 2021), the data-driven approach still lacks interpretability, resulting in lack of generalization ability in the transfer process and poor performance in certain extreme scenarios, e.g., cyclone tracking, and turbulence sensing. This remains a common challenge for deep models as shown in Fig 1(b). 2) Customized designs for specific tasks endow the models with specialized capabilities and high performance, however, the complicated designs make the model difficult to generalize. Consequently, unlike enormous public model 7oos in NLP and CV realms, each spatio-temporal dynamical system is primarily focused on performing a specific scene task and lacks the ability to transfer knowledge from a higher perspective. For instance, using infrared meteorological data from one region for rainfall prediction in another region.
In this paper, we propose novel research line for the first time, namely, causal spatio-temport dy- namics, aiming to provide interpretable paradigm for future large-scale ST dynamical systems. Guided by the currently prevailing technique of causal invariant learning (Arjovsky et al., 2019; Sagawa al., 2019; Rosenfeld et al., 2020; Chang al., 2020; Liu al., 2022), our primary objec- tive is to reveal the inherent correlations in available high-quality measurement data, thus providing interpretability for complex ST problems in dynamical systems such as representation learning and transfer learning. By leveraging the causal patterns inherent in the limited data, our approach bol- sters the reliability of processes such as representation learning and transfer learning. Additionally, our method subtly performs data augmentation on sparse and extreme scenarios, thereby enhancing the model's s ability to perceive and understand such cireumstances. This enables the extraction of causal features to aid in downstream tasks, leading to an improved, streamlined model performance.
Uncovering Causal Correlations. We present the first attempt to introduce the concept of causality to ST dynamical systems by establishing a novel philosophical framework termed NuwaDynami cs. Briefly put, our objective is to inject the invariant characteristics and the internal causal patterns within the data from upstream self-supervised tasks, providing a faithful and reliable framework for downstream learning. Concretely, we decompose our process into two stages Discovery and Up- date. The Discovery stage aims to answer the question of identifying the latent causal components within observed data, where we introduce supervieeed tasks to the upstream ST data reconstruc- tion. Using the popular Vision Transformer architecture (Khan al., 2022), we first patchify the
2

---

Published as a conference paper at ICLR 2024
observations at each time step, and then utilize attention maps localize crucial regions (Selvaraju et al., 2017; Wang et al., 2020a; Jiang et al., 2021a). These localizations are combined with existing pixel-space visualizations to create causal patches.
Going beyond the above process, the Update stage endeavors to evolve the downstream tasks into causal ST model. By appropriately augmenting non-causal patches (i.e., environmental patches), we are in effect generating different randomly deformed copies of the original data. As a result, the model is exposed a broader distribution of latent data and extreme scenarios, offering insights with a causal perspective for downstream tasks. This process can be further understood as backdoor ad- justment in the causal theory (Pearl, 2009; Pearl & Mackenzie, 2018). We believe that such insights will open avenues for future research on learning ST systems and their real-world applications.
Our contributions can be summarized in the following four aspects:
In this paper, we present causal and resilient philosophy (NuwaDynamics) for modeling spatio- temporal systems with the first shot. Leveraging the causal theory with good interpretability. NuwaDynamics allows the model to see a broader potential distribution of data, ensuring the model's outstanding performance across a wide range of downstream tasks.
In its elegant simplicity, NuwaDynamics cs identifies causal features in its first stage and then refines the model into a causal form. It aspires to master data invariance through upstream supf-supe training, offering more tailored and reliable foundation for specific downstream tasks.
NuwaDynamics can benefit many existing frameworks on various tasks. For some longstanding challenging issues like extreme weather perception (e.g., hurricanes, high-resolution precipita- tion), it effectively aids models in achieving perfection in detail.
We evaluate our framework using eight state-of-the-art models as backbones on six diverse bench- marks, including weather, human motion, fire evolution, pollution diffusion, etc. Empirical results show that our concept helps existing models achieve better results in ST representation learning, long-range super-resolution forecasting, and transfer learning. Even in extreme events featured by data scarcity, NuwaDynamico has showcased remarkable ability capture intricate details.
2 PRELIMINARIES
Forecasting Models mostly fall into three categories: those grounded in CNNs (Oh al., 2015; Mathieu al., 2015; Tulyakov et al., 2018). those rooted in RNNs (Srivastava al., 2015; Villegas et al., 2017; 2018; Kim et al., 2019; Wang et al., 2022b; Tan et al., 2023), and an as- sortment of other architectures which include hybrid models (Weissenborn al., 2019; Kumar et al., designs (Dosovitskiy et al., 2020; Gao et al., 2022b; Bai et al., 2022; Wu et al., 2023b). Notably, there are models that leverage graph neural networks (GNNs) primarily for graph data management (Sun al., 2020; Wang et al., 2020b; Jiang et al., 2021b; Wang al., 2022a). However, these are outside the scope of our research as we focus on the visualization ST observational data (Chen al., 2022b; Veillette al., 2020). Within our research, we formulate ST observations ST sequence {x Rg HxWxx. Based these observations, we aim to K parallelly predict the K-step-ahead future [)+++1 , X, E XCout, where H and W denote the number of spatial grids with Cin Coua-dimensional observations.
2019) and transformer-centee
Causal Inference has garnered considerable attention in the realm of deep learning (Zhang et al., 2020a; Woo al., 2022; Zheng et al., 2021; Arjovsky al., 2019) in recent years. Conceptually, causal inference (Pearl et al., 2000; Pearl, 2009) focuses on uncovering the causal relationships between variables, aiming to achieve stable and robust learning and inference. Central to the idea of discovery is the commitment to identifying spurious correlations (Geirhos et al., 2018; Sagawa al., 2019; Koh et al., 2021; Gulrajani & Lopez-Paz, 2020). Discovering spurious correlations exposes model biases that can adversely affect generalization (Wu et al., 2023c). Recently, many techniques (Selvaraju et al., 2017: 2016; Luo et al., 2020; Ying et al., 2019) have been employed for crucial causal features perception. These techniques are sufficiently versatile, exhibiting strong influence across various scenarios (Wu et al., 2023c; Fu et al., 2020; Sui et al., 2022; Zhang et al., 2020b). In this paper, we introduce novel ST causal framework; by leveraging attention techniques, we can more effectively identify causal regions in observations from the upstream task, providing a more robust foundation for downstream tasks.
Vision Transformer (ViT) Pruning. Our work closely resembles the popular ViT image token pruning techniques (Dosovitskiy al., 2020). However, NuwaDynamics emphasizes identifying
3

---

Published as a conference paper at ICLR 2024
important tokens rather than performing token pruning. These endeavors (Rao et al., 2021; Pan et al., 2021; Yuan et al., 2021: Xu et al., 2022) attempt to distinguish how informative a token is by using classification token [C1s] the guideline. However, NuwaDynamics upstream focuses on pre-training to reconstruct patches, aiming to discover causal patches, without involving [C1s] tokens. Hence, traditional ViT pruning approaches may not be suitable for our framework.
3 METHODOLOGY
In this section, systematically introduce our NuwaDynamics framework. We begin with an ex- ample of causality, which serves as the motivation behind our approach in Sec 3. Subsequently, we provide a detailed account of algorithmic process, encompassing the upstream self-supervised tasks in Sec 3.2 and the specifics of the downstream spatio-temport tasks in Sec 3.3.
3.1 MOTIVATION EXAMPLES
Let us first consider example as shown in the upper part of Fig 2. If only focus correlations between exercise duration and cholesterol levels, we may observe that longer exercise durations potentially linked higher cholesterol levels, which contradicts common sense. Merely using a framework to model this is very likely to produce incorrect conclusions. However, this issue arises because haven't taken the variable "age" into account. In reality, age affects both exercise duration and cholesterol levels, resulting in the observed data pattern. We hope to uncover more of the data's latent distribution through increased data augmentation, aiming to mitigate such issues. We also provide a more quantitative example in Appendix B.
3.2 DISCOVERY SPURIOUS CORRELATIONS
Based on the above motivations, we take a causal look at the ST data-generating process and formal- ize the principle of identifying causal and non-trivial regions in input observations, which guides dis- covery strategy (left hand in Fig 3). Naturally, we need a universal mechanism to inspect the causal and spurious regions in the input image. We resort to cur-
Ex.
Age (A)
Cholesterol (C)
Cholesterol
Exercise
fo.
rently popular ViT tools (Dosovitskiy al., 2020) which decompose images into patches of equal size and attempt locate essential patches. However, previous ViT pruning techniques (Pan et al., 2021; Yuan et al., 2021; Xu et al., 2022) have primar- ily focused classification tasks and do not trans- fer well to our spatiotemporal prediction scenarios. Hence, for the first time, we propose an upstream self-supervised reconstruction task and aim to iden- tify potential causal regions during the reconstruc- tion process. Vision transformer first splits the im- age X #x into L HW/p² non- overlapping patch tokens and embedding it into D dimension feature space. Then all tokens are added with a learnable position encoding and then fed into a stacked transformer block:
Exercise
Mon/Women
Recover
Net
Total
6/71
877263
36/25
270/80
Ignoring gender
Recover
Not Recover
Total
Trestment
273
350
Not Trestment
289
61
350
Recovery percentage with//uttout treatment (Mew
0.934400.871) Women 0.7%C/00690) 079Eyos33)))
Figure 2: The motivation of our proposal.
XMhSA
MHSA (LN (x)), XEFN XMHHA FFN (LN (2)).
where MHSA denotes multi-head self-attention (Vaswani al., 2017); FFN and LN represent feed- forward network and layer normalization, respectively. In this circumstance, the input is mapped to query, key and value matrices, i.e., Q. K, E 1.x Then, we can calculate the attention weights Att E R.. by using a softmax function, and the attention weight of the i-th patch towards the j-th patch can be represented as:
Oi, softmax (g.kz) /VDH)/H))))))) Att, where qi E Q, ky E K.
The input are sliced into A attention heads, and here D# = D/A A is the feature dimension. Here, we employ attention weight to describe patch importance without introducing any additional variables hyperparameters. However, for each row of the attention map, have ai,- Σ/ oij 1,
4

---

Published as a conference paper at ICLR 2024
DistnXplorer
Encoder
Attention
Paxls Magine
Paxds Magiis
p-
Pauch Marging
Meplae
Courai Envirosmendal parches
SwirT Book
SwirT Rock
SWiiT Nick
Minooakioo Slarkiy
Data Space
SCM
Frzz
Gerrinsqueid
Eadeddiigg
2
Graaphoo
No Exiza lasralle
Prrs
Spirious conclations
(c/as) toker
Linerr Projection of Slattened Parchex
Lokyycn
Impott observations
Panchely
Augmenied Observations
Downstream Model
Patches]
Reconstruction loss
Stage I. Causal Patches Diacovery
Freeze!
Stage 2. Causal Model Update
Figure 3: The details of NuwaDynamics, in which consists of Discovery and update two stages, For ease of understanding, we use Swin Transformer the upstream model.
in which cannot distinguish the importance of each patch. Hence, resort to the column score of the attention map and calculate the summation of column attention weights:
11 d",//H.
a+,,
(3)
Clearly, Oaj. exhibits the total attention weights of other tokens to the current token, which can be sufficient indicate the importance of the current token. o represents the average importance and the weight of j-th patch across A heads in multi-head attention. Then we move forward to calculate the normalized importance Ij of each patch and select the parts set of smallest values (Note as M) as environmental patches as:
1
Oe. EL,
M index: arg min 1)).
JElL.2)
We employ the saliency map M - 10, IR by sampling the smallest elements in 1. In this way, we can identify the causal patches za and the complementary environments 2.. In general, due to the unstable nature of spurious attributes (Wu al., 2023c), the test distribution IPe is often different from the training settings, i.e., Pt. * PPrr. NuwaDynamica ca is dedicated to enhancing model perception of the underlying essence of data. This not only facilitates representation learning and transfer learning tasks but also addresses the challenges posed by data scarcity.
DistriXplores Recall that the causal theory (Pearl, 2009; Bunge, 2017) attributes the model's weak generalization capability to the distribution shift of spurious associations, namely, the environmental part. We resort causal intervention forcibly assign values to environmental patches. Towards this end, we design DistriXplorer to modify the environmental patches, aiming to enhance sce- narios with observable environmental patches. However, intervening the patch level complex. Existing interventions primarily focus on the class level (Wu et al., 2023c) and the graph realm (Feng al., 2021). Interestingly, the characteristics of particular patch are often influenced by its spatial neighboring and temporal historical patches. Guided by this property, claborate an ST mixup DistriXplorer. Concretely, we sample spatial neighboring and temporal historical patches to mix up and generate different random deformed copies:
Aner
gh+l-1
"ner where Amer 1
Phens
(5)
J-1
Here z/S denotes the environmental patches at points. O represents the number of causal patches among the neighbors. \xer.i and 3 are the weight allocated for spatial neighboring patches and temporal decaying coefficient. Going beyond interventions, we randomly sample the surrounding patches with probability of Ω Uniform(O, 1) generate multiple ordered sequences for next training. In this way, the downstream model can learn patterns for adjusting the environment to improve the generalizability (We summarize generative methods that can be integrated into our augmentation in Future Work, here we choose Mixup for trade-off of efficiency and performance).
Causal Support of NuwaDynamics. Drawing from the causal theory, we construct a Structural Causal Model (SCM) (Pearl, 2009) by examining four variables: input observations X', ground- truth V causal patches in X' denoted as Č, and the confounder (i.e., environment) represented by S. Then we can depict the causal relationships among them by:
5

---

Published as a conference paper at ICLR 2024
Č X The input X' consists of two disjoint parts C and S. C- 6-- S. C is the only endogenous parent to determine the ground-truth y However, in practical scenarios, S is simultancously used for predicting V, leading to spurious associations.
In general, model Fø trained with Empirical Risk Minimization (ERM) often falls short of gen- cralizing to the test data Die ~ Pte- These distribution shifts are triggered by changes in the environmental patches. Therefore, environmental confounder. As shown in the right panel of Fig 3, we employ causal intervention to assist the downstream models in perceiving broader range of test distributions, i.e., Pee Peez. etc. Our framework exploits do-calculus (Pearl al., 2000) on variable C to remove the backdoor patch S by estimating P())do(())) P'm (2)():
imperative address the confounding effect exerted by the
P (21do (o) (o))=E E P. S.)
(6) where T denotes the number of environments. S; denotes the i-th environmental variable. The environmental enhancement at upstream of Nuwa aligns well with the backdoor adjustment the- ory, thereby effectively exploring the potential test environment distributions. Detailed proofs are provided in Appendix H.
3.3 UPGRADING TO CAUSAL INFERENCE: A NEW ERA IN MODELING
Typically, downstream models can be categorized into transformer and non- oonttranfforme classes. For the transformer class, ensuring consistency between the upstream and downstream models allows for rapid parameter transfer, facilitating quicker optimization of the downstream model. On the other hand, for non-transformer architectures, employ transfer-augmented data to optimize and update the downstream models, advancing their causal perception capabilities. We store the intervention data described in See 3.2 in the apatio-temporal bank ST (t). In downstream tasks, we retrieve the data from the bank, ensuring the consistency of prediction labels between the intervention data and the original data for parallel training. However, each timestamp has its corresponding environ- mental patches, denoted by g for the number of environmental patches at time i, theoretically, 2 prediction sequences can be constructed. This poses significant, or even intractable, com- putational burden on the model. In the ST scenario, we argue that historical data closer to the current moment potentially have greater influence. Therefore, we propose a temporal Gaussian decay sam- pling method to identify more influential data, aiming to enhance the model's generalization ability while reducing its computational burden:
(-(--T²) g (T, 0 exp (- -022 202
x= ==((sTTOO a((o.)))
x
(7)
We use the current moment T' as the mean value, with variance o² control the sampling ratio. In this way, we construct our training data as X. Details can be found in Appendix I.
EXPERIMENTS
In this section, present empirical results demonstrate the effectiveness of NuwaDynamics framework. The experiments aim to investigate the following research questions:
Does NuwaDynamics enhance performance prediction for existing ST backbones?
Does NuwaDynamice outperform on specific challenging tasks?
Can the acquisition of feature invariance enhance model generalization?
In the context of rare and extreme events, how effective is NuwaDynamica detection?
4.1 EXPERIMENTAL SETTINGS
We conduct extensive experiments to evaluate the effectiveness of NuwaDynamics. We implement different backbones using Pytorch and leveraging the A100-PCIE40GB support. We train all models with Adam optimizer and learning rate as 0.01. More detail can be found in Appendix C.
Datasets & Backbones We extensively evaluate our proposal on five benchmarks across diverse research domains, including TaxiBJ+ (Liang et al., 2021). )), KTH (Schuldt et al., 2004), SEVIR (Veil- lette et al., 2020), RainNet (Ayzel et al., 2020), PD, and FireSys (Chen et al., 2022a). Specifically, TaxiBJ+ tackles urban traffic, KTH focuses on human kinetics, SEVIR analyzes extreme weather, Nei forecasts precipitation, PD simulates pollutant dispersion, and FireSys monitors wildfires.

---

Published as a conference paper at ICLR 2024
Table 1: Performance comparison on different backbones, where "Ori" refers to the backbones, and "+NuWa" indicates the performance after incorporating NuwaDynamics. All experimental results are the average of five runs and the red font indicates the optimal value. Except for PD. which is 6 + 6, all others are 10 + 10.
Backbone (10 + 10)
Metric
TaxiBJ+
SEVIR (CSLMM))
RuinNet
(6 + 6)
FireSys
Ori
NaaWa Ori
NaWa
Na Ori
+NuWa Ori
+NaWa Ori
Wa
The upstream architecture is Transformer based and maintaaic consistency between upstroom and downstream structures.
MAB
3.48
2.27
59.32
34.56
37.07
46.88
0.78
0.74
83.45
24.70
3.21
3.09
VI [220
MSE
0.16
0.07
57.88
35.43
4.53
3.16
0.23
0.19
8.99
2.45
8.27
8.19
0.00
22.45
1.37
0.08
MAE
3.22
2.18
55.44
33.45
38.22
45
0.67
0.66
79.53
26.38
2.98
2.76
SwinT [2021]
MSE
0.21
0.11
52.38
33.
4.37
2.84
0.22
0.19
8.47
3.15
7.96
7.65
0.10
19:27
1.89
0.03
5.32
0.31
MAR
80.32
40.77
36.68
46.88
1,21
1.17
81,23
30.54
4.65
4.55
Rainformer [2022]
MSE
77.99
40.75
4.02
3.38
0.30
0.21
8.63
2.51
11.27
10.72
37.2
0.64
0.09
6.12
0.55
MAE
52.357
42
01:21
16'33
1.96
1.54
73.24
30.78
1.97
1.37
Earthformer [2022h]
48.65
37.21
3.88
2.96
0.20
7.32
2.44
11.44
0.92
0.01
4.88
0.23
The upstreum architectue is VET and downstrean does specify particular model architecture
MAR 5.52
3.27
128.33
53. 10
41.93
44.88
3.98
3.64
100.44
58.39
11.21
10.97
ComvLSTM [2015]
0.33
0.27
126.32
3.84
3.17
0.49
10.98
5.47
17.22
16.43
0.06
36.97
0.67
0.19
5.51
0.79
MAE
4.33
3.22
51.38
40.37
100 83'
44..9'
2.67
2.43
95.43
72.77
4.32
3.997
PredRNN-V2 [20226]
MSE
0.27
0.20
51.36
45.76
3.98
3.17
0.41
0.33
9.65
7.35
5.87
4.53
0.07
5.60
0.81
0.08
2.30
1.34
MAE
4.25
3.27
86.37
52.98
40.56
45.38
3.85
3.72
100.2
23
78.34
4.98
4.65
E3D-LSTM [2018b]
MSE
0.29
0.25
87.69
59. 49
4.37
3.89
0.38
0.29
10.31
7.35
8.76
8.12
0.04
28
0.48
0.09
2.99
0.64
MAE""
"3.07
2.56
13.39
33.98
115 198
17'09
1:27
1.00
30.93
31
1.98
1.51
SimVP [2022a]
MSE
0.14
0.07
40.93
32.89
3.44
2.92
0.28
0.20
5.48
3.24
2.65
2.42
8.04
0.52
0.23
In our predictions, we simultancously utilize the past 10 images to forecast the next 10. Additionally, due to the larger resolution of the PD dataset, we adopt a - 6 approach. Given that the upstream framework employs Transformer-ban architecture, we ensure the downstream structure both em- ulates and differentiates from the upstream one order to validate the universality of our algorithm. Concretely, we use Transformer- 2020), Swin't (Liu et al., 2021), Rainformer (Bai et al., 2022) and Earthformer (Gao et al., 2022b), as well non- Transformers such as ConvLSTM (Shi et al., 2015). PredRNN-V2 (Wang et al., 2022b), E3D-LSTM (Wang et al., 2018b) and Sim (Gao et al., 2022a). All Transformers had 12 encoder blocks, while non- Transformers used Transpose Conv2d for upsampling. This evaluation aims to clarify the efficacy of each architecture in managing NuwaDynami cs complexities, laying solid foundation for future model refinement. More detail can be found in Appendix C.
models as our backbone, such ViT (Dosovitskiy et al.,
Mcasurement metric. We delve into the metrics used by evaluation methods. Concretely, we train backbones with mean squared error (MSE), and use mean absolute error (MAE), MSE and structural similarity index measure (SSIM) as evaluation metrics. Specifically, for the SEVIR dataset, we incorporate the CSI index (Ayzel et al., 2020) replace MAE as a primary metric for comparison. For MAE and MSE, we use 1. indicates better performance, and higher value (1) denotes better results for SSIM and CSI-M. More details are placed in Appendix C.
4.2 ASSESSING THE EFFICACY OF NUWADYNAMICS (RQ1)
As a preparation, we selected both Transformer and non- Transformer architectures. For the Trans- former architecture in our upstream tasks, we ensure that the sequence lengths of the upstream reconstruction tasks and downstream prediction tasks are consistent, allowing for direct model pa- rameter transfer. For non-transformer architectures, we only transfer the data to train the downstream models. Marker and denote the decrement in MSE and the variation in CSI-M, respectively. We summarize the results in Tab 1, from the experimental results, We make the following Observations:
Obs 1. +Nuwa consistently outperforms without NuwaDynamics concept. As shown in Tab 1, we can easily observe that upon integrating the NuwaDynamico concept into the model (+Nuwa), there were consistent improvements in performance. This is evident from the reductions observed in both MSE and MAE. Specifically, for complex spatiotemporal data such PD, introducing +Nuwa can yield significant benefits: a range of 4.88 ~ 6.51 descents in Transformer scenarios and 0.08 ~ 7.35 in non-transformer scenarios across 8 different backbones based on MSE metrics.
7

---

Published as a conference paper at ICLR 2024
Imputs
Ground-rruth & Predicrions
Inputs
Ground crath & Predictions
10
11-20
10
11-20
Earthformer
Earthformer
Earthformer +Nowa
Earihformer
+Nuwa
PredRNN-V2
PredRNN-V2
PredRNN-V2 NNoa
PredRNN-V2
+Nuwa
Figure 5: Visualization on KTH & SEVIR. For simplicity, we display the results of the last frames.
Inputs
Ground-rrth Prodictions
Settings (TaxiRJ+)
Output Sequence
Backbone
Input Seq (6)
+
G
-
-> 12
32.332
SimVP
wio luura
98.67
96.43
94.32
99.12
97.77
1128
SunVP
Nina
PredRNN-V2
Busa
96.89
94.53
93.41
94.54
Earthfomner
Buwa
87.12
85.44
76.56
PrddRNN-A2
Nuva
89.92
86.43
77.12
Earthformer
Table 2: Model performances three backbones under w/o and + Nuwa con- ditions. We set the input sequence as 6 frames and the predictive length as 6, 8, 12 with SSIM performances.
Figure 4: Visualizations on backbones and +Nuwa. For layout convenience, we only display the last six frames and showcase complete results in Appendix E.
Obs 2. NuwaDynamics demonstrates remarkable adaptability across myriad of spatiotempo- scenarios. NuwaDynamice cs has been validated across a wide range of ST realm, including traffic, human motion, climate, precipitation, etc. These real-world datasets repeatedly attest to the robust generalizability of NuwaDynamics. For instance, on climate and precipitation datasets like SEVIR and RainNet, there was an average MSE reduction of approximately 0.13 and 0.91, respectively. Notably, when integrating causal perturbations on the SEVIR dataset under the CSI-M metric, im- provements of approximately 7.40 and 3.28 were observed for Transformer and non-transformer architectures, respectively. Further analysis of visualization results are presented in Fig 12.
Obs 3. NuwaDynamics excels in capturing predictive details. When attached NuwaDynamicas Earthformer and PredRNN-V V2 present the predictions good sharpness. In KTT, NuwaDynami.c could help the model predict the sharpest sequence compared with original backbones and largely enrich the details for each part of the body, especially for the arms and legs. In SEVIR, the model achieves more reliable predictions on details, with texture information becoming more pronounced.
4.3 EVALUATING THE PERFORMANCES ON CHALLENGING TASK (RQ2)
Although the aforementioned experiments have demonstrated the efficacy of NuwaDynamics, val- idations have been limited. For instance, we have only verified the model under conditions where the temporal lengths and image sizes between upstream and downstream tasks are consistent. To further clucidate the robust adaptability, we choose a more challenging ST task, i.e., the long tempo- ral step super-resolution prediction. Specifically, we selected TaxiBJ+ as the validation benchmark, since it has intricate temporal dynamics. The input images were downsampled to 32 x 32, and we utilize the past frames to predict the next frames at a resolution 128 x 128. To accommodate various spatial resolutions and diverse temporal lengths, we employ spatial sampling module for downstream framework (Here we choose Si,VP, PredRNN-V2 and Earthformer). For a fairer comparison. only perform data transfer without transferring model parameters. We have placed the setting details in Appendix F. From the Fig and Tab 2, we make observations:
Obs 4. NuwaDynamics shows great prominence in challenging task. We find that for long- range super-resolution prediction tasks, all models benefit from Nuwa. As depicted in Fig 4, Earthformer exhibits the most pronounced advantage. In long-distance forecasting scenarios, especially for 12-
8

---

Published as a conference paper at ICLR 2024
time-step predictions, Sim PredRNN- V2, and Earthformer all achieve an improvement ranging from 0.56 to 1.44 on SSIM. This further attests to the efficacy of our model.
4.4 TRANSFERABILIT OF NUWADYNAMICS (RQ3)
ST transfer learning has long been considered challenging problem, given its intricate ST cor- relations (Yao et al., 2020; Wang et al., 2018a). Few works have managed to effectively transfer certain ST patterns to assist another scene. RegionTrans (Wang et al., 2018a) takes the first step to achieve inter-city knowledge transfer, and (Yao et al., 2020) designs a differentiable frame for unsupervised transfer learning across multiple ST tasks. In this study, we chose more complex sense of meteorological prediction as backdrop to explore whether ST transfer tasks can bene- fit from NuwaDynamics, Concretely, we select SEVIR (100.6GB), J, as source data and target data (Note as RainNet = SEVIR), respectively, Fur- ther, we choose SOTA frameworks based on CNN (SimVP), RNN (PredRNN- V2), and Transformer (Earthformer) to systematically validate the feasibility of our framework.
meteorological datasets, RainNet (50GB) and
Obs 5. NuwaDynami cs greatly improves transferability of general models. From holistic in- sight, the model's transfer capability has improved by approximately 1.34 ~ 7.75 on SSIM. This demonstrates Nuwa's assistance transfer learning. An intriguing observation was made by us: Earthformer exhibits greater capability than SimVP and PredRNN-V2, further supporting the po- tential superiority transformers in transfer learning tasks and effectiveness of our upstream task. We showcase the visualizations of the transfer learning results in Appendix G.
Table 3: Performances under w/o and + Nuwa. Marker and denote the RainNet (R) SEVIR (S), and RainNet (R) + SEVIR (S) performances. denotes improvements in SSIM metrics.
Sim VP, (10 -> 10)
PredRNN-V2, (10
10)
Earthformer, (10
10)
w/o Nuwa
+Nuwa
w/o Nuwa
+Nuwa
A
w/o Nuwa
+Nuwa
A
R R
72.12 65.49
76.98
4.86
65.43
66.97
1.54
82.18 75.32
85.67
3.49
S R
66.97
1.48
64.37
72.12
7.75
76.66
1.34
4.5 EXTREME WEATHER FORECASTING OF NUWA (RQ4)
Extreme weather forecasting has always been considered highly challenging task with significant real-world implications (Bi et al., 2022). Due to the rarity of extreme events, current research often struggles to achieve high fidelity in detailing (Scher & Messori, 2019; Schultz et al., 2021; Keisler, 2022). SEVIR and RainNet contain a vast collection of high-quality extreme events like Storm, Hurricane Florence and Squall. We showcase the visualizations in Fig to illustrate that the en- hancements in Nuwa's upstream can aid in better discerning the distribution of potential extreme events, thereby improving perceptual capability. As expected, Nuwa helps the model in capturing the intricate details of extreme weather, achieving exceptional local fidelity.
Storm
Hurricane Florence
Squall
Gound- trit
Reinfomee
Raisieene (Noso
Figure 6: Visualization of storm, hurricane florence, and squall.
CONCLUSION
In this paper, we present the first attempt to introduce the causality philosophy in ST forecasting tasks. We propose a two-stage causal framework, NuwaDynamics, to discover non-trivial regions in data and update the model into causal frameworks. It performs self-supervised learning in up- stream reconstruction tasks for intervening in environmental regions and augmentation on trivial part (aligned with the backdoor adjustment). Extensive experiments across six real-world spatiotem- poral (ST) benchmarks demonstrate that models enhanced with the NuwaDynamics concept yield improved results. In the future, we plan explore causal learning spatio-temporal graphs.

---

Published as a conference paper at ICLR 2024
6 ACKNOWLEDGEMENT
This paper is partially supported by the National Natural Science Foundation of China N2.66072427, No. 12222227,,) the Project of Stable Support for Youth Team in Basic Research Field, CAS (No. YSBR-005), Academic Leaders Cultivation Program, USTC. The authors also thank Guangzhhu-HKKSSTGGG Joint Funding Program (No. 2024A03J0620), National Natural Science Foundation of China (62176165). This work is supported by the National Research Foundation, Singapore under its AI Singapore Programme (AISG Award No: AISG2-PhD-2021-00 008).
REFERENCES
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Pate Invariant risk minimization. arXiv preprint arXiv: 1XX.vv997700293 2019.
Georgy Ayzel, Tobias Scheffer, and Maik Heistermann. Rainnet v1. 0: convolutional neural net- work for radar-based precipitation nowcasting. Geoscientific Model Development, 13(6):2631- 2644, 2020.
Cong Bai, Feng Sun, Jinglin Zhang, Yi Song, and Shengyong Chen. Rainformer: Features extrac- tion balanced network for radar-based precipitation nowcasting. IEEE Geoscience and Remote Sensing Letters, 19:1-5, 2022.
Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Pangu-weather: A 3d high-resolution model for fast and accurate global weather forecast. arXiv preprint arXiv:2211. 22221.10556. 2022.
Mario Bunge. Causality and modern science. Routledge, 2017.
Marius Bürkle, Umesha Perera, Florian Gimbert, Hisao Nakamura, Masaaki Kawata, and Yoshihiro Asai. Deep-learning approach to first-principles transport simulations. Physical Review Letters, 126(17):17700,, 2021.
Shiyu Chang, Yang Zhang, Mo Yu, and Tommi Jaakkola. Invariant rationalization. In International Conference on Machine Learning, pp. 1448-14 1458. PMLR, 2020.
Boyuan Chen, Kuang Huang, Sunand Raghupathi, Ishaan Chandratreya, Qiang Du, and Hod Lip- son. Automated discovery of fundamental variables hidden in experimental data. Nature Compu- tational Science, 2(7))433-444 2022a.
Xuanhong Chen, Kairui Feng, Naiyuan Liu, Bingbing Ni, Yifan Lu, Zhengyan Tong, and Ziang Liu. Rainnet: A large-scale imagery dataset and benchmark for spatial precipitation downscaling. Advances in Neural Information Processing Systems, 35:9797-9812, 2022b.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 1929, 2020.
Junfeng Fang, Wei Liu, An Zhang, Xiang Wang, Xiangnan He, Kun Wang, and Tat-Seng Chua. On regularization for explaining graph neural networks: An information theory perspective. 2022.
Junfeng Fang, Wei Liu, Yuan Gao, Zemin Liu, An Zhang, Xiang Wang, and Xiangnan He. Evalu- ating post-hoc explanations for graph neural networks via robustness analysis. In Thirth-seventh Conference on Neural Information Processing Systems, 2023a. URL htRs:t/s:://// /ttttps://operreviee net/forrmm forum? (d=eD534mPP.
Junfeng Fang, Wei Liu, Yuan Gao, Zemin Liu, An Zhang, Xiang Wang, and Xiangnan He. Evalu- ating post-hoc explanations for graph neural networks via robustness analysis. In Thirty-seventh Conference on Neural Information Processing Systems, 2023b.
Junfeng Fang, Xiang Wang, An Zhang, Zemin Liu, Xiangnan He, and Tat-Seng Chua. Cooperative explanations of graph neural networks. In WSDM, pp. 616-624. ACM, 2023c.
10

---

Published as a conference paper at ICLR 2024
Junfeng Fang, Xinglin Li. Yongduo Sui, Yuan Gao, Guibin Zhang, Kun Wang, Xiang Wang, and Xiangnan He. Exge: Bridging efficiency and explainability in graph condensation, 2024a.
Junfeng Fang, Shuai Zhang, Chang Wu, Zhiyuan Liu, Sihang Li, Kun Wang, Wenjie Du, Xiang Wang, and Xiangnan He. Moltc: Towards molecular relational modeling in language models, 2024b.
Fuli Feng, Weiran Huang, Xiangnan He, Xin Xin, Qifan Wang, and Tat-Seng Chua. Should graph convolution trust neighbors? simple causal inference method. In Proceedings of the 44th In- ternational ACM SIGIR Conference on Research and Development Information Retrieval, pp. 1208-1218, 2021.
Ruigang Fu, Qingyong Hu. Xiaohu Dong, Yulan Guo. Yinghui Gao, and Biao Li. Axiom based grad-cam Towards accurate visualization and explanation of cnns. arXiv preprint 0rXiv:2008.02312. 2020.
Zhangyang Gao, Cheng Tan, Lirong Wu, and Stan Z Li. Simvp: Simpler yet better video prediction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 3170-3180, 2022a.
Zhihan Gao, Xingjian Shi, Hao Wang, Yi Zhu, Yuyang Bernie Wang. Mu Li, and Dit- Yan Yeung. Earthformer: Exploring space-time transformers for earth system forecasting. Advances in Neural Information Processing Systems, 35:25390- -24403,, 2022b.
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Imagenet-trained cnns are biased towards texture; increasing shape bias im- proves accuracy and robustness. arXiv preprint arXiv:/8/1.12233, 2018.
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization, arXiv preprint 2xXv:20077.01444, 2020.
Catalin Ionescu, Dragos Papava, Vlad Olaru, and Cristian Sminchisescu. Human3. 6m: Large scale datasets and predictive methods for 3d human sensing in natural environments. JEEE transactions on pattern analysis and machine intelligence, 36(7): 1215::3323333333 1339, 2013.
Peng-Tao Tao Jiang, Chang-Bing Bin Zhang, Qibin Hou, Ming-Ming Cheng, and Yunchao Wei. Layercam: Exploring hierarchical class activation maps for localization. IEEE Transactions on Image Pro- 2021a.
cessing,
Renhe Jiang, Du Yin, Zhaonan Wang, Yizhuo Wang, Jiewen Deng, Hangchen Liu, Zekun Cai, Jin- liang Deng, Xuan Song, and Ryosuke Shibasaki. DI-traff: Survey and benchmark of deep learning models for urban traffic prediction. Proceedings of the 30th ACM international conference on information & knowledge management, pp. 4515-4525, 2021b.
Guangyin Jin, Yuxuan Liang, Yuchen Fang, Jincai Huang, Junbo Zhang, and Yu Zheng. Spatio- temporal graph neural networks for predictive learning in urban computing: A survey. arXiv preprint arXiv:2303. /4483, 2023.
Sepideh Kaffash, An Truong Nguyen, and Joe Zhu. Big data algorithms and applications in in- telligent transportation system: review and bibliometric analysis. International Journal of Production Economics, 231:107868, 2021.
Ryan Keisler. Forecasting global weather with graph neural networks. arXiv preprint arXiv:2202.07555, 2022.
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah. Transformers in vision: A survey. ACM computing surveys (CSUR), 54(10s): 1-41, 2022.
Ali Rahimi Khojasteh, Sylvain Laizet, Dominique Heitz, and Yin Yang. Lagrangian and eulerian dataset of the wake downstream smooth cylinder at a reynolds number equal to 3900. Data in brief, 40: 107725, 2022.
11

---

Published as a conference paper at ICLR 2024
Taesup Kim, Sungjin Ahn, and Yoshua Bengio. Variational temporal abstraction. Advances in Neural Information Processing Systems, 32, 2019.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang. Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning, pp- 5637-5664. PMLR, 2021.
Manoj Kumar, Mohammad Babaeizadeh, Dumitru Erhan, Chelsea Finn, Sergey Levine, Laurent Dinh, and Durk Kingma. Videoflow: A flow-based generative model for video. arXiv preprint arXiv:1903.01434 2(5):3, 2019.
Xinglin Li, Kun Wang, Hanhui Deng, Yuxuan Liang, and Di Wu. Attend who is weak: Enhancing graph condensation via cross-free adversarial training. arXiv preprint arXiv:2311.15772 2023.
Yuxuan Liang, Kun Ouyang, Junkai Sun, Yiwei Wang, Junbo Zhang, Yu Zheng, David Rosenblum, and Roger Zimmermann. Fine-graines urban flow prediction. In Proceedings of the Web Confer- ence 2021, pp. 1833-1845 2021.
Xu Liu, Yutong Xia, Yuxuan Liang, Junfeng Hu, Yiwei Wang, Lei Bai, Chao Huang, Zhenguang Liu, Bryan Hooi, and Roger Zimmermann. Largest: A benchmark dataset for large-scale traffic forecasting, arXiv preprini arXiv:2306.08255: 2023a.
Yiheng Liu, Tianle Han, Siyuan Ma, Jiayue Zhang, Yuanyuan Yang, Jiaming Tian, Hao He, Antong Li, Mengshen He, Zhengliang Liu, et al. Summary of chatgpt-related research and perspective towards the future of large language models. Meta-Radiology, pp. 100017, 2023b.
Yuejiang Liu, Riccardo Cadei, Jonas Schweizer, Sherwin Bahmani, and Alexandre Alahi. Towards robust and adaptive motion forecasting: A causal representation perspective. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 17081-17099 2022.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Zemin Liu, Xingtong Yu, Yuan Fang, and Xinming Zhang. Graphprompt: Unifying pre-training and downstream tasks for graph neural networks. In Proceedings of the ACM Web Conference 2023, pp. 417-428, 2023c.
Dongsheng Luo, Wei Cheng, Dongkuan Xu, Wenchao Yu, Bo Zong, Haifeng Chen, and Xiang Zhang. Parameterized explainer for graph neural network. Advances neural information pro- cessing systems, 18:19220-19333 2020.
Michael Mathieu, Camille Couprie, and Yann LeCun. Deep multi-scale video prediction beyond mean square error. arXiv preprint arXiv: 1.1004440 2015.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard I. Lewis, and Satinder Singh. Action-condiiin video prediction using deep networks in atari games. Advances neural information processing systems, 28, 2015.
Bowen Pan, Rameswar Panda, Yifan Jiang, Zhangyang Wang, Rogerio Feris, and Aude Oliva. la- red 2: Interpretabilit-awwar redundancy reduction for vision transformers. Advances in Neural Information Processing Systems, 014-2888-244111 2021.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Four- castnet: A global data-driven high-resolution weather model using adaptive fourier neural opera- tors. arXiv preprint arXiv:2202. /2 2022.
Judea Pearl. Causality. Cambridge university press, 2009.
Judea Pearl and Dana Mackenzie. The book of why: the new science of cause and effect. Basic
books, 2018
12

---

Published as a conference paper at ICLR 2024
Judea Pearl al. Models, reasoning and inference. Cambridge, UK: CombridgeUnnvverttyyresss 19(2):3, 2000.
Roger W Pryor. Mulriphysics modeling using COMSOL@: first principles approach. Jones & Bartlett Publishers, 2009.
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. pp. 13937-13949, 2021.
Pengzhen Ren, Yun Xiao, Xiaojun Chang, Po- Yao Huang, Zhihui Li, Brij B Gupta, Xiaojiang Chen, and Xin Wang. A survey deep active learning. ACM computing surveys (CSUR), 54(9):1-40) 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. arXiv preprint arXiv:2010.05761, 2020.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. arXiv preprint arXiv:191.008731 2019.
Sebastian Scher and Gabriele Messori. Weather and climate forecasting with neural networks: using general circulation models (gems) with different complexity as study ground. Geoscientific Model Development, 1277)2797-2809, 2019.
Christian Schuldt, Ivan Laptev, and Barbara Caputo. Recognizing human actions: local svm approach. In Proceedings of the 17th International Conference on Pattern Recognition, 2004. ICPR 2004., volume 3, pp. 32-36 IEEE, 2004.
Martin G Schultz, Clara Betancourt, Bing Gong, Felix Kleinert, Michael Langguth, Lukas Hubert Leufen, Amirpasha Mozaffari, and Scarlet Stadtler. Can deep learning beat numerical weather prediction? Philosophical Transactions of the Royal Society A, 379(2194):202009977 2021.
Ramprasaath R Selvaraju, Abhishek Das, Ramakrishna Vedantam, Michael Cogswell, Devi Parikh, and Dhruv Batra. Grad-cam: Why did you say that? arXiv preprint 16 2016.
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based local- ization. In Proceedings of the IEEE international conference on computer vision, pp. 618-626, 2017.
Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai kin Wong, and Wang chun Woo. Convolutional Istm network: A machine learning approach for precipitation nowcasting, 2015.
Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov. Unsupervised learning of video rep- resentations using Istms. In International conference machine learning, pp. 843-852. PMLR, 2015.
Yongduo Sui, Xiang Wang, Jiancan Wu, Min Lin, Xiangnan He, and Tat-Seng Chua. Causal at- tention for interpretable and generalizable graph classification. Proceedings of the 28th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 1696-1705... 2022.
Junkai Sun, Junbo Zhang, Qiaofei Li, Xiuwen Yi, Yuxuan Liang, and Yu Zheng. Predicting citywide crowd flows in irregular regions using multi-view graph convolutional networks. IEEE Transac- tions on Knowledge and Data Engineering. 34(5):2344-233) 2020.
Wai Cheong Tam, Eugene Yujun Fu, Jiajia Li, Xinyan Huang, Jian Chen, and Michael Xuelin Huang. A spatial temporal graph neural network model for predicting flashover in arbitrary building floor- plans. Engineering Applications of Artificial Intelligence, 105:50 2022.
13

---

Published as a conference paper at ICLR 2024
Cheng Tan, Zhangyang Gao, Lirong Wu, Yongjie Xu, Jun Xia, Siyuan Li, and Stan Z Li. Tem- poral attention unit: Towards efficient spatiotemporal predictive learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 18770-18782, 2023.
Sergey Tulyakov, Ming- Yu Liu, Xiaodong Yang, and Jan Kautz. Mocogan: Decomposing motion and content for video generation. In Proceedings of the IEEE conference on computer vision and paltern recognition, pp. 1526-1535, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Mark Veillette, Siddharth Samsi, and Chris Mattioli. Sevir: A storm event imagery dataset for deep learning applications in radar and satellite meteorology. Advances in Neural Information Processing Systems, 33:22009-22019, 2020.
Ruben Villegas, Jimei Yang. Yuliang Zou, Sungryull Sohn, Xunyu Lin, and Honglak Lee. Learning to generate long-tern tergttermmnn future via hierarchical prediction. In international conference on machine learning, pp. 3560-3569. PMLR, 2017.
Ruben Villegas, Dumitru Erhan, Honglak Lee, et. Hierarchical long-term video prediction without supervision. In International Conference on Machine Learning, pp. 6038 6046. PMLR, 2018.
Haofan Wang, Zifan Wang, Mengnan Du, Fan Yang, Zijian Zhang, Sirui Ding, Piotr Mardziel, and Xia Hu. Score-cam: Score-weighted visual explanations for convolutional neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pp. 24-25, 2020a.
Kun Wang, Zhengyang Zhou, Xu Wang, Pengkun Wang. Qi Fang, and Yang Wang. A2djp: A two graph-based component fused learning framework for urban anomaly distribution and duration joint-prediction. IEEE Transactions Knowledge and Data Engineering, 2022a.
Kun Wang, Yuxuan Liang, Xinglin Li, Guohao Li, Bernard Ghanem, Roger Zimmermann, Huahui Yi, Yudong Zhang, Yang Wang, et al. Brave the wind and the waves; Discovering robust and gen- eralizable graph lottery tickets. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023a.
Kun Wang, Yuxuan Liang, Pengkun Wang, Xu Wang, Pengfei Gu, Junfeng Fang. and Yang Wang. Searching lottery tickets in graph neural networks: dual perspective. In The Eleventh Interna- tional Conference on Learning Representations, 2023b. URL htttps://openreviiee net/ Forum?id=Dvs-aBee a3aymee.
Kun Wang, Zhengyang Zhou, Xu Wang, Pengkun Wang, Qi Fang, and Yang Wang. A2DJP: A two graph-based component fused learning framework for urban anomaly distribution and duration joint-prediction. IEEE Trans. Knowl. Data Eng., 35(12): 1984 11 11998, 2023c.
Kun Wang, Hao Wu, Guibin Zhang, Junfeng Fang. Yuxuan Liang, Yuankai Wu, Roger Zimmer- mann, and Yang Wang. Modeling spatio-temporal dynamical systems with neural discrete learn- ing and lfvels-of-experts. IEEE Transactions on Knowledge and Data Engineering. 2024.
Leye Wang, Xu Geng, Xiaojuan Ma, Feng Liu, and Qiang Yang. Cross-city transfer learning for deep spatio-temporal prediction. arXiv preprint arXiv: arXiv:/802.00s 1802.00386, 201 8a.
Senzhang Wang, Jiannong Cao, and S Yu Philip. Deep learning for spatio-temporal data mining: A survey. IEEE transactions on knowledge and data engineering, 34(8):3681 -3700, 2020b.
Yunbo Wang, Lu Jiang, Ming-Hsuan Yang, Li-Jia Li, Mingsheng Long, and Li Fei-Fei. Eidetic 3d Istm: A model for video prediction and beyond. In International conference on learning representations, 2018b.
Yunbo Wang, Haixu Wu, Jianjin Zhang, Zhifeng Gao, Jianmin Wang, S Yu Philip, and Mingsheng Long. Predmn: A recurrent neural network for spatiotemporal predictive learning. IEEE Trans- actions Pattern Analysis and Machine Intelligence, 45(2):2208-2222 2022b.
14

---

Published as a conference paper at ICLR 2024
Dirk Weissenborn, Oscar Täckström, and Jakob Uszkoreit. Scaling autoregressive video models. arXiv preprint 1.XXvv:906.0266. 2019.
Gerald Woo, Chenghao Liu, Doyen Sahoo, Akshat Kumar, and Steven Hoi. Cost: Contrastive
learning of disentangled seasonal-trent representations for time series forecasting. arXiv preprint arXiv:2202.01575, 2022.
Hao Wu, Shilong Wang, Yuxuan Liang, Zhengyang Zhou, Wei Huang, Wei Xiong, and Kun Wang. Earthfarseer: Versatile spatio-temporal dynamical systems modeling in one model. arXiv preprint arXiv:2312.08400, 2023a.
Hao Wu, Wei Xiong, Fan Xu, Xiao Luo, Chong Chen, Xian-Sheng Hua, and Haixin Wang. Past- net: Introducing physical inductive biases for spatio-temporal video prediction. arXiv preprint arXiv:2305.1142/ 2023b.
Shirley Wu, Mert Yuksekgonul, Linjun Zhang, and James Zou. Discover and cure: Concept-aware mitigation of spurious correlation. arXiv preprint arXiv:2305.00M 2023c.
Yutong Xia, Yuxuan Liang, Haomin Wen, Xu Liu, Kun Wang, Zhengyang Zhou, and Roger Zim- mermann. Deciphering spatio-temporal graph forecasting: A causal lens and treatment. CoRR, abs/2309.13378. 2023.
Yifan Xu, Zhijie Zhang, Mengdan Zhang. Kekai Sheng, Ke Li, Weiming Dong, Liqing Zhang, Changsheng Xu, and Xing Sun. Evo-vit: Slow-fast token evolution for dynamic vision trans- former. In Proceedingy of the AAAI Conference on Artificial Intelligence, volume 36, pp. 2964- 2972, 2022.
Zhiyu Yao, Yunbo Wang, Mingsheng Long, and Jianmin Wang. Unsupervised transfer learning for spatiotemporal predictive networks. In International Conference on Machine Learning, pp. PMLR, 2020.
Rex Ying, Dylan Bourgeois, Jiaxuan You, Marinka Zitnik, and Jure Leskovec. Gnnexplainer: Gen- erating explanations for graph neural networks, 2019.
Xingtong Yu, Zemin Liu, Yuan Fang, and Xinming Zhang. Hgprompt: Bridging homogeneous and heterogeneous graphs for few-shot prompt learning. arXiv preprint arXiv:23/2.01878, 2023a.
Xingtong Yu, Zhenghao Liu, Yuan Fang, Zemin Liu, Sihong Chen, and Xinming Zhang. General- ized graph prompt: Toward a unification of pre-training and downstream tasks on graphs. arXiv preprint arXiv:2311.15317, 2023b.
Xingtong Yu, Chang Zhou, Yuan Fang, and Xinming Zhang. Multigprompt for multi-task pre- training and prompting graphs. arXiv preprint ar arXiv:2312.03731, 2023c.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi-Hang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 558-567, 2021.
Dong Zhang, Hanwang Zhang, Jinhui Tang, Xian-Sheng Hua, and Qianru Sun. Causal interven- tion for weakly-supervised semantic segmentation. Advances in Neural Information Processing Systems, 33:655 1665,666,222 2020a.
Guibin Zhang, Yanwei Yue, Kun Wang, Junfeng Fang, Yongduo Sui, Kai Wang. Yuxuan Liang, Dawei Cheng, Shirui Pan, and Tianlong Chen. Two heads are better than one: Boosting graph sparse training via semantic and topological awareness. arXiv preprint 2rX.v:2402.01222. 2024.
Jingzhao Zhang, Aditya Menon, Andreas Veit, Srinadh Bhojanapalli, Sanjiv Kumar, and Suvrit Sra. Coping with label shift via distributionally robust optimisation. arXiv preprint arXiv:2010. 12230, 2020b.
Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision pp. 3836-3847, 2023.
15

---

Published as a conference paper at ICLR 2024
Yu Zheng, Chen Gao, Xiang Li, Xiangnan He, Yong Li, and Depeng Jin. Disentangling user in- terest and conformity for recommendation with causal embedding. In Proceedings of the Web Conference 2021, pp. 2021.
16

---

Published as a conference paper at ICLR 2024
A EXAMPLES
We present the global distribution of PH in 2022, sourced from the average of monthly observations distributed at the Global Disaster Alert and Coordination System (GDACS) over the calendar year (https: / / www. ocean-ops org/). From the data (Fig 7). we easily observe a significant imbalance in the PH distribution within global ocean currents. Within each grid area, the number of grids ranging from 0.01 to 0.50 more than five times greater than the number of grids ranging from 1.01 to 1.50. During data analysis, it becomes challenging to predict extreme regions (c.g., regions with PH 3.0) using models.
0.00 (2662)
(05-0.50(554)
953-1.00:219
1:00-.90 (3)
(.000300
Yearly Coverage gH
Average monthl dhservato dtorbbteed 5oNci over calendar yer
Figure 7: Global distribution PH 2022, sourced from the average of monthly observations distributed at the Global Disaster Alert and Coordination System (GDACS) over the calendar year.
Another case of uneven sensor distribution (Fig 8) reveals that the deployment of observation equip- ment in the Indian Ocean is significantly sparser compared to the Pacific and Atlantic Oceans. This data collection. Leveraging data from regions with higher density of sensor deployment to guide regions with relatively fewer deployments can provide significant assistance addressing these challenges.
imbalance in observation deployment can lead to challenges
Pacific Ocean
Atlantie Ocean
Indian Ocean
Dense
Sparse
Dense
Argo
Yearly Deployments
2022
Launch profiing deployed during the year
Figure 8: Launch locations for all profile buoy deployments in 2022.
17

---

Published as a conference paper at ICLR 2024
B AN QUANTITATIVE EXAMPLE SUPPORT MOTIVATION
We next contemplate a more quantitative example in the lower part of Fig 2. Suppose we are in- vestigating whether drug aids recovery from an illness. For males, out of 270/87 who did not take or took the medicine, and the recovery rate among those who took the medicine was 0.93, which is higher than those who did not. A similar phenomenon can be observed hort (0.73>0.66 0.69). However, when we disregard gender, we reach the opposite conclusion, with the recovery rates being 0.79 for those who took the medicine and 0.83 for those who did not. This phenomenon is known Simpson's 's Paradox (Pearl et al., 2000), caused by the unobserved gender variable the aggregate data. By thoroughly traversing the confounder variables, we can effectively mitigate the above issue. This approach forms the crux of our framework and is also referred as "backdoor adjustment" (Pearl & Mackenzie, 2018; Pearl et al., 2000).
the female co-
C EXPERIMENTAL SETTINGS
Evaluation metrics. The Critical Success Index (CSI) (Ayzel et al., 2020; Gao et al., 2022b) serves as prevalent metric in precipitation forecasting to gauge prediction accuracy. CSI is given by: CSI + Misses + F.A larms). Here, Hits, Misses, and F. F.Alarms are the quantities of true positives, false negatives, and false positives. To compute these quantities, we rescale the predicted and true values to lie between O and 255 and determine binary classifications using the thresholds [16, 74, 133, 160, 181, 219]
By computing CSI values across various thresholds, assess the predictive performance of the model, employing the mean CSI-M as a comprehensive evaluation metric. A superior CSI value signifies precise precipitation prediction by the model, whereas inferior CSI suggests room for enhancement in predictive capability. Thus, the CSI stands as a pivotal metric precipitation forecasting, offering insight into the model's efficacy and directing refinements in the prediction algorithm.
Hits, Misses, and F.Alarms important indicators for evaluating the performance of the prediction. Specifically;
True positive (Hits): Denotes the model's accurate prediction precipitation occurrence.
False negative (Misses): Indicates the model's oversight in predicting an actual precipitation event.
False positive (F (l.alams)): Reflects the model's erroneous prediction of precipitation, specifically when precipitation does not materialize.
In predictive modeling, greater count of Hits, Misses, and F.Alarms implies diminished perfor- mance. In precipitation forecasting, our objective maximize the number of Hits, concurrently minimizing Misses and F. F.Alarms, thereby enhancing the accuracy and trustworthiness of the pre- dictions.
Details of benchmarks. Here we provide a systematic introduction to the benchmark we used. For a clearer and better understanding, we have placed the statistical characteristics in Tab 4.
Table 4:
Dataset statistics. Nutr and Ne denote the number of instances in the training and test sets. The lengths of the input and prediction sequences Ij and O₁, respectively.
Dataset
N fr
N_te
(C, H, W)
I
O
Interval
TaxiBJ+
3555
445
(2, 128, 128)
10
10
30 mins
KTH
108717
4086
(1, 128, 128)
10
10
SEVIR
4158
500
384, 384)
10
10
mins
RainNet
6000
1500
(1, 208, 333)
10
10
1 hour
PD
2000
500
(3, 1400, 1400)
6
6
5 seconds
FireSys
2000
500
(3, 128, 128)
10
10
This dataset encompasses trajectory information sourced from Beijing taxis' GPS, de- lineated into two distinct channels: inflow and outflow. Furthermore, we've augmented the origi- nal dataset by gathering recent trajectory details from Beijing and enhancing the resolution from 32x32 to 128x 128, designating TaxiBJ+
TaxiBB+:
18

---

Published as a conference paper at ICLR 2024
This dataset comprises 25 individuals executing six distinct actions: walking, jogging, running, boxing, waving, and clapping. The intricacy of human movements stems from the unique variations individuals exhibit when performing these actions. By analyzing preceding frames, the model can grasp the nuances of human dynamics and anticipate future prolonged postural shifts.
KTH:
The SEVIR dataset features radar-based readings of vertical accumulation liquid (VIL., captured at 5-minute intervals with 1 km resolution. This dataset serves the foundational source for rain and hail detection.
SEVIR:
RainNet: This benchmark boasts over 62,400 pairs of top-notch low/high-resolution precipitation maps spanning more than 17 years, primed to facilitate the advancement of deep learning models precipitation downscaling.
Dollutant-Diffusiie (PD): This data is derived from the computational fluid dynamics (CFD) simulation outcomes related to pollutant dispersion within a designated area. We selected wind speed of 15m/s, with the wind direction set to due north, and utilized the centering point as the dynamic data for the pollutant release point.
FireSys: The FireSys dataset encompasses data related to fire observations, where both temporal and spatial trends of fire evolution accurately reflect the progression status in nature.
D VISUALIZATIONS TO ANSWER RQI
Input
Ground- Truth
SimVP
-Nawa
Input
Ground-
Truth
SimVP +Nuwa
Figure 9: Visualization results of SEVIR under VinWaa we can observe that incorporating NuwalDyammic significantly enhances the model's ability to capture fine-grained details.
1,1
1.1
4
&
Input
Ground- Truth
SimVP +Nowa
Figure 10: Visualization results of FireSys under SimVP+Nuwa, We find that upon integrating NuwaDynamics, the model's predictive outcomes adeptly capture the edge information of flames, offering commendable prediction in terms of fine details.
19

---

Published as a conference paper at ICLR 2024
E VISUALIZATIONS TO ANSWER RQ2.
In this section, we present the complete visualization results on TaxiBJ+. It's evident that SimVP achieves the best visual details in prediction, while Earthformer's visualization performance is in- consistent. However, when enhanced with Nuwa, all models achieve improved visualization out- comes. This improvement is most notably observed in Earthformer's results.
Inputs
Ground-truth Predictions
t 7~18
SimVP
PredRNN-V2
Earthformer
SimVP
+Nuwa
PredRNN-V2
+Nuwa
Earthformer
+Nuwa
Figure 11: Visualization results of TaxiBJ+ under VimVP+Nuwa, PredRNN-V2+Nuwa and Earth- former+Nuwa. Here we showcase the last eight frames for ease of understanding.
F EXPERIMENTAL SETTINGS ON LONG TEMPORAL STEP SUPER RESOLUTION PREDICTION
In this section, have delved into scenarios where the input-output dimensions of pre-trained models differ from those of downstream tasks. We addressed this in two parts: (1) handling low- resolution data inputs, and (2) managing varying prediction lengths. Specifically, the details are as follows:
Spatial Upsampling in Spatio-temporal Data The spatio-tempore upsampler contains a spe- cially designed upsampling module aimed at enhancing the spatial resolution of time series data. Given five-dimensional input tensor z with the shape B x T x H x W, where B represents the batch size, T denotes the number of time steps, C' stands for the channel count, and H and W respec- tively describe the height and width of the feature map. This tensor is initially reshaped into four- dimensional form as shown by Treshape reshape(r, (B C, H, W)). Subsequently, through transpose convolution operation with a stride of 4, kernel size of 4, and zero padding, the spatial di- mensions of the feature map are expanded, yielding Fupsompled =t4,4.0 (Zreshape). ) with the result- ing dimensions being 4H x 4W. Ultimately, the output feature map final is reshaped back into its original five-dimensional shape, expressed as final = reshape (Fapsampled; (B, T, C, 4H, 4W)). In summary, the entire upsampling process can be represented as I.final = reshape (14,4.0 (reshape( (ar, (B x T, C, H, W))). (B, T,, CH, 4W)), thereby facilitating a precise transformation from low to high resolution.
Adaptive Temporal Forecasting with Autoregressive Process In the context time series fore- casting using CNN-based method with input tensor of dimensions B x T C H x W. there
20

---

Published as a conference paper at ICLR 2024
exists challenge in flexibly extending the temporal dimension. While expanding temporal channels offers a means to alter the length of output predicted frames, a more computationally efficient strat- egy is sought. Mimicking RNNs provides solution: RNNS inherently generate long-term forecasts by recycling prior predictions as present inputs. When the desired prediction length K' is shorter than the input sequence length T, the most recent K timesteps are sliced from the input, adjusting to B x K x C x H x W. This adjustment ensures that the model's autoregressive predictions align with the intended temporal horizon.
G TRANSFERABILITY OF NUWADYNAMICS
Ground-truth Predictions
Ground-truth Predictions
NNnwo
Earthform
Earthformer
or
Niny
Si
Sim VP
PredRNN
Ninw
ProdRNN-V2
+
Figure 12: Visualizations of the transfer learning. We only display the last six frames for conve- nience.
H PROOFS OF BACKDOOR ADJUSTMENT IN NUWADYNAMICS
Backdoor adjustment refers to a method used in causal inference eliminate or control for con- founding variables that may affect the relationship between the treatment and the outcome. This is achieved by conditioning on a set of variables (the backdoor criterion) that blocks backdoor paths from the treatment the outcome through confounders. By doing fect of the treatment on the outcome from the biases introduced by confounders. this work, employ the backdoor adjustment mechanism to better assist downstream models in perceiving potential test distributions (Yu et al., 2023c;b;a; Liu et al., 2023c).
can isolate the causal ef-
The do-caleulus is of three rules introduced by (Pearl al., 2000) part of the causal inference framework. It's a mathematical formalism for reasoning about interventions and causal effects. The do-calculus is utilized for deriving expressions for causal effects in terms of observed distributions, which can be evaluated from data. The rules of do-calculus allow for the manipulation of expressions involving "do" operators, which correspond interventions shown in Fig 13, based on the above descriptions, we can apply the following three rules:
causal model. As
(Yldo (c), P(YIC)) since the environment
Rule Insertion/dellttion of observations.
variable does not affect the prediction of C to V. L.e. VLSIĆ.
(y/do (c) do (S)) P(Yido (c).s) if Ō not
Rule Action/observation exchange.
descendant of S.
Rule 3: Rule of Reversal. P (CIdo(21) P(ČI)) if Ö' is not descendant of y
Backdoor Adjustment. Based on the above three rules, we showcase the relevance of our algorithm and backdoor adjustment. Our algorithm can be well-understood to enhance the potential data and remove backdoor paths. Uniquely, due to the complexity of envi- ronmental variables in the backdoor paths, it significantly increases the training burden. In Section 3.3, we introduced the concept of spatio-temporal (ST) bank to better select influential patches, thereby achieving trade-off between performance and computational resources.
form of backdoor adjustment
21

---

Published as a conference paper at ICLR 2024
SCM
X
SCM
SCM
C
S
C
S
(a)
(b)
(c)
Figure 13: Fig (a) represents the general deep models prediction processes, which consider the envi- ronmental parts. Fig (b) illustrates that within the input, there exists an environmental portion does not contribute to the model's prediction, which may consequently lead to spurious associations. (c) denotes the model prediction after backdoor adjustment, we can remove spurious correlations by traversing the potential test distributions.
T P (yido (c))= Σ >/ (V/do (J) , S.) p ps+= St do (c)
I
E ( s- s.) P(S= S.)
Rule
= si)
5.) p(s=
Rule
I DESCRIPTIONS OF ST BANK AND GAUSSIAN SAMPLING
We employ a discretized Gaussian formula to extract data from historical time steps to construct sequences, aiming to better aid the model in creating backups of apatio-temporal prediction data:
(- (x (-6--
T) (x)
(T,
exp
0/22%
2x
The aforementioned formula can be further illustrated in Fig 15, we use the current moment as the mean, and with variance of o², apply Gaussian distribution probability sample the intervention data in the ST bank. Evidently, in the ST bank, the closer the data to the current moment, the higher the sampling ratio, and vice versa.
Not Filtered
Sampling Probability
Po
P1
T-2
P2
7-3
T-4
T
7-27-1 Ty
Figure 14: An illustration of discretized Gaussian formula.
22

---

Published as a conference paper at ICLR 2024
J ABLATION EXPERIMENTS ON SPATIAL AND TEMPORAL AUGMENTATION
In this section, we meticulously execute three distinct ablations to elucidate the influence of key components within our proposed model. These methodologies are articulated as follows:
This method primarily involves the manipulation of only the causal patches at the current time point, deliberately omitting historical patches. This approach is intended isolate and evaluate the impact of immediate causal effects, devoid of historical influences, thereby providing insight into the temporal immediacy of the model's performance.
1. Strategy A1:
2. Strategy A2: In this variant, our attention pivots to historical data, but with a significant alteration: all historical patches are accorded equal significance, effectively disregarding the time decay factor. Each time point is uniformly weighted, with value of 1. This modification is designed to probe the model's s sensitivity to temporal variations and to ascertain the importance of differentially weighting historical data based on their temporal proximity.
3. Strategy A3: This strategy concentrates on the examination of historical patches specific to the region of interest, while excluding the broader spatial context and other causal patches. These patches are incorporated with a decay factor, with the objective of exploring the localized temporal dynamics and their isolated influence the model's 's predictive accuracy.
For clarity and consistency in our analysis, these strategies are systematically designated as A1, A2, and A3. The outcomes of these ablation tests, particularly in terms of MAE, tabulated. This structured presentation is chosen to enable a lucid comparison and an in-depth understanding of the distinct and collective impacts of these strategies sive ablation study, we aim to unravel the complex inter-dependencies among causal, temporal, and spatial elements in our analytical construct.
model's efficacy. Through this comprehen-
TaxiBI+
KTH
RainNet
PD
A1
A2
A3
A1
A2
A3
A1
A2
A3
A1
A2
A3
VIT SimVP
2.89
2.72
3.02
37.64
35.25
43.26
1.16
1.08
1.25
40.86
35.64
50.62
SWin SimVP
2.72
2.61
2.94
35.56
34.12
42.64
1.05
0.99
1.16
39.34
34.16
49.13
ViT PredRNN
3.81
3.54
4.28
45.42
42.06
51.25
2.55
2.49
2.62
84.36
79.92
92.86
SWin + PredRNN
3.72
3.41
4.15
44.32
41.17
50.14
2.45
2.48
2.53
82.32
77.43
90.95
Table 5: Ablation study results showing MAE across different datasets and backbones.
The comparison across strategies indicates that Strategy A2, which treats all historical data equally and does not consider time decay, typically offers superior performance. This suggests the pre- dominance of historical data in enhancing predictive accuracy. Conversely, Strategy A1, which concentrates solely on the current causal patches, advantageous for datasets where present data is more indicative of future outcomes. However, Strategy A3 leads to increased MAE for the KTH and PD datasets, underscoring the significance of spatial context in these scenarios. Collectively, these insights reveal that while historical data is paramount, spatial context is indispensable for datasets with intricate spatial dependencies.
K ADDITIONAL RESULTS ON HIGH-RESOLUTION DATASETS
In order to deeply study Nuwa's capability in handling high-resolution datasets (Khojasteh et al., 2022), we chose a cylinder dataset with 768x768 resolution to systematically validate the perfor- mance and effectiveness of our algorithm. In order maintain the consistency of the main paper, we chose ViT, Swin Transformer, Rainformer, Earthformer, as well as ConvLSTM, PredRNN-V2. E3D-LSTM, and VP the backbone networks for our experiments. We ensure that the experi- mental setup remains consistent with the main part.
The implementation of Nuwa across various deep learning models consistently enhances model accuracy, with notable reductions in both MSE and MAE metrics. The Earthformer model exhibits the most dramatic improvement, dropping from 0.49 to 0.33 in MSE and from 0.43 to 0.32 in MAE.
23

---

Published as a conference paper at ICLR 2024
Table 6: Experimental results on top of the high-resolution cylinder dataset
Backhone
MSE
MAE
Oni
+NuWa
Ori
+NuWa
ViT
0.67
0.37
0.56
0.37
SwinT
0.61
0.34
0.55
0.33
Rainformer
0.55
0.49
0.36
0.40
Earthformer
0.49
0.33
0.43
0.32
ConvLSTM
0.61
0.48
0.53
0.38
PredRNN-V2
0.70
0.61
0.49
0.39
E3D-LSTM
0.54
0.31
0.43
0.31
SimVP
0.40
0.31
0.37
0.30
The experimental results also show the robustness of Sim which maintains the lowest MSE and MAE both before and after the enhancement using Nuwa. While the Rainformer model displays an unusual increase in MAE, suggesting trade-off introduced NuWa Improvements vary among models, with Earthformer, SimVP. and E3D-LSTM benefiting substantially, illustrating NuWa's variable impact.
Impux
Grons-touth
StaVP
PadRRNN-V2
PrddRNN-V2
Nooa
Figure 15: Visualization of spatiotemporal prediction with cylinder dataset. The figure demonstrates that upon incorporating Nuwa, the model achieves greater precision in details, with predictions displaying more accurate textural information.
L EVALUATIONS OF CAUSAL DISCOVERY
In light of the lack of well-defined causal region in traditional datasets, this section elucidates the accuracy our causal discovery through a real pedestrian movement dataset. We select the Human3.6m 6m for upstream self-supervised tasks and employ attention maps to visualize the areas of importance loneseu et (2013). The results, as depicted the following figure, demonstrate that our method proficiently correlates pedestrians with causal regions.
As illustrated in the accompanying Figure 16. It is readily apparent that our upstream model profi- ciently identifies significant areas. These results serve to validate the causal discovery capabilities of our algorithm.
M FUTURE WORK
In this paper, we have pioneered the investigation into the integration of causal reasoning with spatio- temporal observable data, systematically formulating a solution that addresses out-of-distribution generalization issues in spatio-temporal contexts. By implementing Mixup data augmentation technique, we have laid the groundwork for substantial improvements in model performance. Rec- ognizing the potential for further advancements. Our future work will delve into extending the appli- cation our work to a broader spectrum downstream graph learning tasks, such as graph pruning (Wang et al., 2023b; Xia et al., 2023; Wang al., 2023c;a), graph sparification (Li et al., 2023; Wu
24

---

Published as a conference paper at ICLR 2024
Figure 16: Left. Attention map regions on the Human3.mm a.mam dataset. Right. Augmented samples by Nuwa.
et al., 2023a; Wang et al., 2024; Zhang et al., 2024; Fang et al., 2024a) and graph explainability (Fang et al., 2023b;e; 2022). Moreover, we plan to explore the utility of stable diffusion (Rombach et al., 2022; Zhang et al., 2023) for environmental patch enhancement and leverage Language Model Learning (L.LM) (Liu al., 2023b; Fang al., 2024b) for more sophisticated data description and send these textual information for guide agumentation. These endeavors aim to refine our model's predictive capabilities and generalizability, thereby contributing to the evolution of robust analytical tools in spatio-temporal data analysis. Interestingly, method's ability to identify causal compo- nents can be further refined and optimized using advanced quantitative metrics. A notable example the OAR metric (Fang al., 2023a), which takes the first step to tackle the inherent OOD issues of traditional metrics in deep learning explainability domains. On the other hand, we acknowledge that there is still room for improvement in the efficiency of Nuwa, which can be optimized using the latest pruning paradigms. For instance, RGLT (Wang et al., 2023a) achieves joint pruning of data and networks while maintaining generalization and robustness through causal pruning theory.
TaxiBJ+
KTH
RainNet
PD
A1
A2
A3
A1
A2
A3
Al
A2
A3
A1
A2
A3
VIT SimVP
2.89
2.72
3.02
37.64
35.25
43.26
1.16
1.08
1.25
40.86
35.64
50.62
SWin SimVP
2.72
2.61
2.94
35.56
34.12
42.64
1.05
0.99
1.16
39.34
34.16
49.13
ViT PredRNN
3.81
3.54
4.28
45.42
42.06
51.25
2.55
2.49
2.62
84.36
79.92
92.86
SWin PredRNN
3.72
3.41
4.15
44.32
41.17
50.14
2.45
2.48
2.53
82.32
77.43
90.95
Table 7: Ablation study results showing MAE across different datasets and backbones.
25


Reviews:
[{"review_id": "m18hu7ZkCK", "reviewer": "Reviewer_PkBe", "summary": "This paper proposes a novel causal framework called NuwaDynamics for spatio-temporal (ST) prediction tasks. The key idea is to leverage self-supervised reconstruction tasks to identify causal regions in the input data, and perform interventions on non-causal regions to expose the model to a broader distribution of data. The framework consists of two main stages: 1) Causal patch discovery, where a Vision Transformer is trained on a reconstruction task and attention maps are used to localize causal patches. 2) Causal model update, where environmental patches are modified through interventions and used to train the downstream model. Experiments on 6 benchmarks demonstrate improved performance over strong baselines.", "questions": "1. What are the limitations of current attention map methods for identifying causal regions?\n2. How sensitive is the approach to the choice of interventions?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper is well-organized and clearly written.\n2. Introducing causality concepts into spatio-temporal modeling is novel.\n3. The motivation is clear and reasonable.\n4. Visualizations demonstrate improved detail and extreme weather perception.\n5. The experiments across multiple datasets are thorough.", "weaknesses": "1. It seems that the computational complexity is high due to generating many intervened sequences.\n2. Lack of ablation studies to validate design choices.\n3. Except for the KTH dataset, the performance gains are relatively low.\n4. Lacks of the comparison on computational complexity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B9qMivpHXs", "reviewer": "Reviewer_zp2y", "summary": "This paper established a causal concept for spatio-temporal predictions. In particular, it firstly identifies causal regions at the Discovery step. It then augments non-causal patches at the Update stage. By doing so, the model is able to see a broader potential distribution of data, achieving improved downstream performance. Extensive experiments show that the proposed model can benefit many existing frameworks on various tasks.", "questions": "Please see the weakness section for my major concerns. In addition,\n\n1.\tHow many augmented samples are required to achieve reasonable performance? Does the downstream performance sensitive to the number of augmented samples?\n2.\tWhat if we directly perform data enhancement using the attention map? \n3.\tCan you show the visualization of the identified causal regions and the augmented samples?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "**Originality:** This paper established a causal concept for spatio-temporal predictions and introduced a novel philosophical framework. Its causal perspective is interesting and original. It is also addressing an important challenge faced by AI community which is the interpretability and generalizability. \n\n**Presentation:** This paper is well written. The flow of the paper is smooth and easy to follow. Figure 1 and Figure 2 are very helpful to illustrate the central idea of the paper. \n\n**Experiments:** the experiments are extensive, covering a wide range of backbones and tasks. The experimental analysis is well organized with key questions and itemized observations.", "weaknesses": "**The experiments are insufficient.** The performed experiments are focused on the comparison between the performance of an original backbone and the performance of adding NuWa. There is a lack of comparison to SOTA performance on the evaluated tasks. Besides, it is also important to compare to other data-augmentation methods (even though they are not using causality), which is also missing in the paper.  \n\nBesides, there is a lack of evaluation on the causal discovery accuracy. Causal discovery always requires a large amount of data to ensure accurate discovery of causal relationships, and it can suffer from highly imbalanced data. How does your model address this issue as you are particularly interested in such regime?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gg0WpY1Ype", "reviewer": "Reviewer_oBVm", "summary": "This paper explores the application of deep learning to dynamical systems and proposes a new philosophical framework called \"NuwaDynamics\", which aims to mine data for causal patterns to enhance the explanatory and generalization capabilities of models.  The approach is divided into two phases: Discovery, which identifies the causal components of the data, and Update, which applies the causal model to downstream tasks.  In addition, NuwaDynamics enhances the model's performance in sparse and extreme situations.  Experimental results in the article show that the framework achieves excellent results in several benchmark tests.", "questions": "(1) Despite the extensive experiments presented in this paper, I'm still curious about the effectiveness of this method on certain spatio-temporal datasets, such as [1-2]. These real pedestrian movement datasets could offer a more convincing validation of the model's efficacy.  However, an intriguing question arises: given that Nuwa's environment is populated with causal patches, is its performance reliable on pedestrian movement datasets that have clear divisions between causal and non-causal segments?\n\n\n[1] Pedestrian detection: A benchmark  **Conference on Computer Vision and Pattern Recognition**\n[2] Human3. 6m: Large scale datasets and predictive methods for 3d human sensing in natural environments **IEEE transactions on pattern analysis and machine intelligence**", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "(1) A conceptually and technically highly innovative paper which seamlessly integrates causal theory with spatio-temporal data mining. The integration of causal theory's environment discovery and augmentation into the popular vision transformer's attention map is ingenious. Coupled with a straightforward mixup augmentation, it achieves a very refined alignment. From my perspective, NuwaDynamics has the potential to revolutionize how models recognize causality in ST data mining realm. This may provide a new approach to solving such problems (data scarcity and the high cost of deploying sensors). \n\n(2) The paper reads very well, and hardly has any errors or inconsistencies. Information is provided at the right point, is complete and accurate. The split between main paper and supplementary material is good. The narrative and exposition flow well. Meanwhile, the paper is well-structured, I appreciate its visual aids.\n\n(3) Providing a causal intervention perspective for data augmentation seems interesting. In terms of practical implementation, utilizing attention scores without introducing additional complex designs seems both simple and effective. Additionally, I believe the backdoor adjustment mechanism proposed by the authors serves as a robust theoretical foundation. This contributes significantly to enhancing generalization, especially in out-of-distribution scenarios.\n\n(4) Baselines are strong, relevant, discussed well and evaluated fairly. A large number of analytic experiments complete the main findings. Meanwhile, experimental setups are presented accurately and consistently, at an adequate level of detail.", "weaknesses": "(1) The paper does not seem to make it clear exactly how the process of obtaining the importance rankings in the patches is done, and it is hoped that the authors will add that detail.\n\n(2) In referring to the fact that DL methods often forgo an explicit understanding of the rules of physics, the practical implications and potential risks of this sacrifice are not explained in detail. More discussion or explanation might have given the reader a better understanding of the consequences of this choice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vxtsdSfENl", "reviewer": "Reviewer_TYFb", "summary": "The paper delves into the spatio-temporal challenges inherent in earth sciences, encompassing areas like meteorological forecasting and urban computing. Given the hurdles of data sparsity and imbalance, the authors put forth a novel causal prediction framework dubbed \"NuwaDynamics\". Operating through two phases - \"Discovery\" and \"Update\" - this paradigm pinpoints causal regions within datasets and endows the model with causal reasoning prowess. The authors also substantiate their methodology with proofs drawn from causal theory and employ a myriad of experimental results to demonstrate its efficacy across various spatio-temporal benchmark tests, especially in tasks like extreme weather forecasting and long-term temporal predictions. In a nod to fostering collaboration and transparency, the code from this research has been shared with the public, marking a valuable contribution to the open-source arena.", "questions": "See weaknesses.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well-written and exhibits a commendable layout; the motivation is intuitive and clear.\n\n2. It is interesting to employ ViT and Mixup in spatiotemporal data for causal environmental discovery and augmentation, and the work can potentially expand the model's latent observational space in many scarce scenarios.\n\n3. Extensive experiments demonstrate the effectiveness of the framework. The authors introduced architectures tailored for both ViT and non-ViT, while also including control experiments across frameworks such as RNN and CNN. Comparisons across multiple mainstream tasks, including transfer learning and super-resolution prediction, further demonstrates its capabilities in tasks like extreme event prediction.\n\n4. Detailed theoretical validation based on causal manipulation is provided to demonstrate the feasibility of the framework. Interestingly, in real-world scenarios, the author doesn't explore all possible environments but instead introduces partial perturbations, which conveniently addresses the computational constraints and limited data availability in practice.", "weaknesses": "1. I am concerned about whether there are clear guiding principles for environmental augmentation. Especially when it comes to the mixup technique, is it possible to automatically adjust the fusion hyperparameters of each environmental patch based on a specific dataset? If this cannot be achieved at the moment, I hope the authors will offer potential solutions in their section of future work .\n\n2. I also have a keen interest in environmental perturbation methods. In fact, environmental perturbation can be considered a form of data augmentation. I wonder, are there any generative methods in computer vision (CV) that can be integrated into this? If the answer is affirmative, I hope the authors delve deeper into this in the \"Related Work\" and \"Model\" sections.\n\n3. The experimental section of the paper is rich in content, and I am particularly interested in spatio-temporal system modeling. Current research literature mainly focuses on datasets with lower resolution, for example, roughly 256x256 in size. I hope the authors can consider more experiments to explore whether these methods remain competitive on high-resolution datasets. If the authors can achieve superior results on higher-resolution datasets, I might consider further raising my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)