Earthformer: Exploring Space-Time Transformers for Earth System Forecasting
Authors: Zhihan Gao
Xingjian Shi
Hao Wang
Yi Zhu
Bernie Wang
Mu Li
Dit-Yan Yeung
Venue: NeurIPS 2022 Accept (2022)
Keywords: Machine Learning for Earth Science
Spatiotemporal Forecasting
Transformers

Abstract: Conventionally, Earth system (e.g., weather and climate) forecasting relies on numerical simulation with complex physical models and hence is both expensive in computation and demanding on domain expertise. With the explosive growth of spatiotemporal Earth observation data in the past decade, data-driven models that apply Deep Learning (DL) are demonstrating impressive potential for various Earth system forecasting tasks. The Transformer as an emerging DL architecture, despite its broad success in other domains, has limited adoption in this area. In this paper, we propose Earthformer, a space-time Transformer for Earth system forecasting. Earthformer is based on a generic, flexible and efficient space-time attention block, named Cuboid Attention. The idea is to decompose the data into cuboids and apply cuboid-level self-attention in parallel. These cuboids are further connected with a collection of global vectors. We conduct experiments on the MovingMNIST dataset and a newly proposed chaotic $N$-body MNIST dataset to verify the effectiveness of cuboid attention and figure out the best design of Earthformer. Experiments on two real-world benchmarks about precipitation nowcasting and El Niño/Southern Oscillation (ENSO) forecasting show that Earthformer achieves state-of-the-art performance.

Earthformer: Exploring Space-Time Transformers for Earth System Forecasting
Zhihan Gao"
Xingjian Shif Amazon Web Services x shi @amazon. com com
Hong Kong University of Science and Technology zhihan gww.connect......oooooom........ ust hk
Hao Wang
Yi Zhu
Yuyang Wang Amazon Web Services
Rutgers University
Amazon Web Services
hw4886cs. rrw.ers.uurreer.... edu
yzaws@amazon.com com
yuyawanggamazon com
Mu Li
Dan Yeung
Amazon Web Services ml @amamzzoo.... com
Hong Kong University of Science and Technology dyyeungêcse. .utee.......... ust.hk hk
Abstract
Conventionally, Earth system (e.g., weather and climate) forecasting relies on numerical simulation with complex physical models and hence is both expensive in computation and demanding on domain expertise. With the explosive growth spatiotemporal Earth observation data in the past decade, data-driven models that apply Deep Learning (DL) are demonstrating impressive potential for various Earth system forecasting tasks. The Transformer its broad success other domains, has limited adoption in this area. In this paper, we propose Earthformer, space-time Transformer for Earth system forecasting. Earthformer is based on a generic, flexible and efficient space-time attention block, named Cuboid Attention. The idea decompose the data into cuboids and apply cuboid-level self-attention in parallel. These cuboids are further connected with a collection global vectors. We conduct experiments the MovingMNIST dataset and newly proposed chaotic N N-body MNIST dataset to verify the effectiveness of cuboid attention and figure out the best design of Earthformer. Experiments on two real-world benchmarks about precipitation nowcasting and El Niño/Southern Oscillation (ENSO) forecasting show that Farthformer achieves state-of-the-art performance.
emerging DL. architecture, despite
Introduction
The Earth is a complex system. Variabilities of the Earth system, ranging from regular events like temperature fluctuation to extreme events like drought, hail storm, and Niño/Southern Oscillation (ENSO), impact our daily life. Among all the consequences, Earth system variabilities can influence crop yields, delay airlines, cause floods and forest fires. Precise and timely forecasting these vari- abilities can help people take necessary precautions to avoid crisis, better utilize natural resources such as wind and solar energy. Thus, improving forecasting models for Earth variabilities (e.g., weather and climate) has huge socioeconomic impact. Despite its importance, the operational weather and climate forecasting systems have not fundamentally changed for almost 50 years [4 These operational models, including the state-of-the-art High Resolution Ensemble Forecast (HREF)
Work done while being an intern Amazon Web Services. Contact person.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Min.
Mts.
Mit.
Figure 1: Example Vertically Integrated Liquid (VIL) observation sequence from the Storm EVent ImageRy (SEVIR) dataset. The observation intensity is mapped to pixel value of the range 0-255. The larger value indicates the higher precipitation intensity.
rainfall nowcasting model used in National Oceanic and Atmospheric Administration (NOAA) 32)
rely on meticulous numerical simulation of physical models. Such simulation-based systems in- evitably fall short in the ability to incorporate signals from newly emerging geophysical observation systems [2]. . or take advantage of the Petabytes-scale Earth observation data 43
As an appealing alternative, deep learning (DL) is offering a new approach for Earth system forecast- ing [4]. Instead of explicitly incorporating physical rules, DL--based forecasting models are trained on the Earth observation data [36). By learning from a large amount of observations, DL models able to figure out the system's intrinsic physical rules and generate predictions that outperform simulation-bar based models 1.. Such technique has demonstrated success in several applications, including precipitation nowcasting 2 d and ENSO forecasting [13]. Because the Earth system is chaotic 21 high-dingg digh-diimessonal, and spatiotemporal, designing appropriate DL architecture for model- ing the system is particularly challenging. Previous works relied on the combination of Recurrent Neural Networks (RNN) and Convolutional Neural Networks (CNN) 336 (37 43 13 45]. These two architectures impose temporal and spatial inductive biases that help capturing spatiotemporal patterns. However, as chaotic system, variabilities of the Earth system, such as rainfall and ENSO, are highly sensitive the system's initial conditions and can respond abruptly to internal changes. is unclear whether the inductive biases in RNN and CNN models still hold for such complex systems.
On the other hand, recent years have witnessed major breakthroughs in DL brought by the wide adoption of Transformer. The model was originally proposed for natural language processing 42 7 and later has been extended to computer vision S 22 multimodal text-image generation graph learning [52], etc. Transformer relies on the attention mechanism to capture data correlations and is powerful at modeling complex and long-range dependencies, both of which appear in Earth systems (See Fig. ffor an example of Earth observation data). Despite being suitable for the problem, Transformer sees Timited adoption for Earth system forecasting. Naively applying the Transformer architecture is infeasible because the O(N2) attention mechanism is too computationally expensive for the digensiinnni Earth observation data. How to design space-time Transformer that is good at predicting the future of the Earth systems is largely an open problem to the community.
In this paper, we propose Earthformer, a space-time Transformer for Earth system forecasting. To better explore the design of space-time attention, we propose Cuboid Attention, which is a generic building block for efficient space-time attention. The idea is to decompose the input tensor to non-overlapping cuboids and apply cuboid-level self-attention in parallel. Since we limit the ((N2) self-attention inside the local cuboids, the overall complexity is greatly reduced. Different types of correlations can be captured via different cuboid decompositions. By stacking multiple cuboid attention layers with different hyperparameters, we are able to subsume several previously proposed video Transformers [19 233 as special cases, and also come up with new attention patterns that were not studied before. A limitation of this design the lack of a mechanism for the local cuboids to communicate with each other. Thus, we introduce a collection of global vectors that attend to all the local cuboids, thereby gathering the overall status of the system. By attending to the global vectors, the local cuboids can grasp the general dynamics of the system and share information with each other.
To verify the effectiieneess of cuboid attention and figure out the best design under the Earth system forecasting scenario, we conducted extensive experiments on two synthetic datasets: the MovingM- NIST (36] dataset and newly proposed N N-body MNIST dataset. Digits in the NN-body MNIST follow the chaotic 3 3-body motion pattern (2.1 which makes the dataset not only more challenging than MovingMNIIT but also more relevant to Earth system forecasting. The synthetic experiments reveal the following findings: 1) stacking cuboid attention layers with the Axial attention pattern is both efficient and effective, achieving the best overall performance, 2) adding global vectors provides consistent performance gain without increasing the computational cost, 3) adding hierarchy in the encoder-decoder architecture can improve performance. Based on these findings, we figured out the optimal design for Earthformer and made comparisons with other baselines the SEVIR 43]
2

---

benchmark for precipitation noweasting and the ICAR-ENSO dataset [5 for ENSO forecasting. Experiments show that Earthformer achieves slate-of-the-ard (SOTA) performance on both tasks.
2 Related Work
Deep learning architectures for Earth system forecasting, Conventional DL models for Earth system forecasting are based on CNN and RNN. U-Net with either 2D CNN or 3D CNN have been used for precipitation nowcasting 4
Seasonal Arctic Sea ice prediction M. and ENSO forecasting [S]. Shi al. 36 proposed the ConVLSTM network that combines CNN and L.STM
proposed PredRNN which adds the spatiotemporal memory flow structure to ConvLSTMM. To better learn long-term high-level relations, Wang al. ++ proposed E3D-LSTM that integrates 3D CNN to LSTM. To disentangle PDE dynamics from unknown complementary information, PhyDNet 13 incorporates a new recurrent physical cell to perform PDE-constrained prediction in latent space. Espeholt et al. [9] proposed MetNet-2 that outperforms HREF for forecasting precipitation. The architecture based on ConvLSTM and dilated CNN. Very recently, there are works that tried to apply Transformer for solving Earth system forecasting problems. Pathak et al. [28] proposed the FourCastNet for global weather forecasting, which is based Adaptive Fourier Neural Operators (AFNO) [4] Bai al. 3 proposed Rainformer for precipitation nowcasting, which is based on an architecture that combines CNN and Swin- Transformer 22 In our experiments, we can see that Earthformer outperforms Rainformer.
for precipitation nowcasting. Wang et al. 45
Space-time Transformers for video modeling. Inspired by the success of VIT S for image classification, space-time Transformer is adopted for improved video understanding. In order to bypass the huge memory consumption brought by joint spatiotemporal attention, several pioneering work proposed efficient alternatives, such as divided attention [4], axial attention [19 4 factorized encoder [27 [ and separable attention [54]] Beyond minimal adaptation from ViT. some recent work introduced more prior to the design of space-tima transformers, including trajectory [29]. multi-scale [23 [ and multi-view
4 However, no prior work focuses on exploring the design of space-time Transformers for Earth system forecasting.
Global and local attention in vision Transformers. To make self-attention more efficient in terms of both memory consumption and speed, recent works have adapted the essence of CNN to perform local attention in transformers [6 52 HaloNets E1 develops a new self-attention model family consisting of simple local self-attention and convolutional hybrids, which outperform both CNN and vanilla ViT on a range of downstream vision tasks. GLiT [5] introduces a locality module and uses neural architecture search to find efficient backbone. Focal transformer [5]] proposes focal self-attention that can incorporate both fine- mme-grained local and coarse- coarse-grained global interactions. However, these architectures are not directly applicable spatiotemporal forecasting. Besides, they are also different from our design because we keep K global vectors to summarize the statistics of the dynamic system and connect the local cuboids; experiments show that such global vector design is crucial to successful spatiotemporal forecasting.
3 Model
formulate Earth system forecasting a spatiotempo- ral sequence forecasting problem. The Earth observation data, such as radar echo maps from NEXRAD [17 and climate data from CIMP6 [10). are represented
Similar previous works 36
43
BJ.
spatiotemporal sequence [M]] X; R HxWxC, Based on these observations, the model predicts the K-step-ahead future Dr+++ E XHxXWXCa Here, H, W denote the spatial resolution, and Cins Cout denote the number of measurements available at each space-time coordinate from the input and target sequences, respectively. As illustrated in Fig. our proposed Earthformer is a hierarchical Transformer encoder- decoder based Cuboid Attention. The input observations are encoded as a hierarchy of hidden states and then decoded to the prediction target. In what follows, we will present the detailed design of cuboid attention and the hierarchical encoder-decoder architecture adopted in Earthformer.
[D++,K K
3.1 Cuboid Attention
Compared with images and text, spatiotemporal data in Earth systems usually have higher dimension- ality. As consequence, applying Transformers to this task is challenging. For example, for 3D
3

---

Inbal Postiiona Embeoding
Guboiii Attention Back X.D
Ciboid Atlention Block x.D.
Mx
M
Downsamplc
Upsample
Cubold
Downcmmple
Uppampple
Attend
Allend
[oot]]
[X
[9/ +:l
Figure 2: Illustration of the Earthformer architecture. It is hierarchical Transformer encoder-decodss based cuboid attention. The input sequence has length T and the target sequence has length K. "x D" means stack D cuboid attention blocks with residual connection. " "Mx" means to have M layers of hierarchies.
Murgo
Xout
Sout
Figure Illustration of the cuboid atten- tion layer with global vectors.
(a) strategy= "local" shift=(0,0,00
(b) strategy="dilaste "dlattdd" shift=(0,0, 0,
(c) strategy= strategy="local" "local" shift = (0, 1,
Figure Illustration of cuboid decomposition strategies when the input shape (T, H, W) (6, 4, 4), and cuboid size (br. by, bw) = (3, 2). Cells that have the same color belong to the same cuboid and will attend each other. shift (0, 1, 1) shifts the cuboid decomposition by 1 pixel along height and width dimensions. strategy = "local" means to aggregate contiguous (br, bu, bw) pixels as a cuboid. strategy = "dilated" means to aggregate pixels every br by
W 1) steps along time (height, width) dimension. (Best viewed in color).
bw
tensor with shape (T. H, W), the complexity vanilla self-attention is O and can be computationally infeasible. Previous literature proposed various structure-aware space-time attention mechanisms to reduce the complexity
These space-time attention mechanisms share the common design of stacking multiple elementary attention layers that focus on different types of data correlations (c.g., temporal correlation and spatial correlation). Steming from this observation, propose the generic cuboid attention layer that involves three steps: "decompose", "attend", and "merge".
Decompose. We first decompose the input spatiotemporal tensor RTxHxWXC into a sequence of cuboids (n)
(x (n)) Decompose(X, cuboid size, strategy, shift).
(1)
where cuboid size (br, bH bw) is the size of the local cuboid, strategy
["local", "dilated") controls whether to adopt the local decomposition strategy or the dilated decomposition strategy 4. shift (ST, sH, sw) is the window shift offset [22. Fig. provides three examples showing how an input tensor will be decomposed following different hyperparameters of Decompose(-). There are total number of 1 bw cuboids in (x) . To simplify the
notation, we assume that T, H, W are divisible by br. bir bw. In the implementation, we pad the input tensor if it is not divisible.

---

Assume (²)) is the (nr, nH, nw) )- cuboid in [-("). The (i. k)-th element of x(1) can be mapped the (t', k")th element of X via Eqn. 2 if the strategy is "local" or Eqn. Bi the strategy is "dilated".
1' ++ ST bT(nT 1) + mod T
+ ST + by 1) nt mod T
j' sn + bu (nH 1) + j mod H
(2) j' ++ sn bu 1) + nM mod H (3)
k' sw (nw - 1) + A mod W
k ++ sw + bw (k 1) + nw mod W
Since the mapping is bijective, one can then map the elements from X' to {x!")} via the inverse operation.
Attend. After decomposing the input tensor into sequence of non-overlapping cuboids (x(1) }. we apply self-attention within each cuboid in parallel.
T out = Aatentiong(x"") x(1), x 1
(4)
The query, key, and value matrices Q. K. and V of Attentions(Q K, V) Softmax (WQQ))W>X))> the resulting matrix back to a 3D tensor. Wo. WK and Wv are linear projection weights and are abbreviated together as θ. The self-attention parameter θ are shared across all cuboids. The compu- tational complexity of the "attend" step is which scales linearly with the cuboid size. Since the cuboid size can be much smaller than the size of the input tensor, the layer is more efficient than full attention.
(WvV) are all flattened versions of (n)), and we unravel
W by bw (brbubw) )²) R O(THW-brbrbw )
Merge. Merge(.) is the inverse operation of Decompose((). The sequence of cuboids obtained after
the attention step ( } are merged back to the original input shape to produce the final output of cuboid attention, as shown in Equ. The mapping follows the same bijections in Equ. and Eqn.
Xaut Merge( (f=00))))))))
( Jn. cuboid_ size, strategy, shift).
(5)
We combine the "decompose", "attend" and "merge" steps described in Eqn. 1 construct the generic cuboid attention as in Eqn. Ø
Xout CubAttne (X, cuboid_ size, strategy, shift).
(6)
Explore cuboid attention patterns. By stacking multiple cuboid attention layers with different choices of "cuboid_siiz" "strategy" and "shift", potentially more effective space-time attention. In this paper, we explore the cuboid attention patterns as listed in Table From the table, space-time attention methods like axial attention, video wwin-Transformer, and divided space-time attention. Also, manually picked the patterns that are reasonable and not computationally expensive as our search space. The flexibility of cuboid attention allows us to conduct Neural Architecture Search (NAS) to automatically search for a pattern but we will leave it as future work.
able efficiently explore existing and
see that cuboid attention subsumes previously proposed
3.2 Global Vectors
One limitation of the previous formulation is that the cuboids do not communicate with each other. This is undesirable because each cuboid is not capable of understanding the global dynamics of the
system. Thus, inspired by the [CLS] token adopted in BERT 7 53) we propose to introduce a collection of P global vectors G Px help cuboids scatter and gather crucial global information. When each cuboid is performing the self-attention, the elements will not only attend to the other elements within the same cuboid but also attend the global vectors G. We revise Eqn. to Eqn. to enable local-global information exchange. We also use Eqn. to update the global vectors G by aggregating the information from all elements of the input tensor X.
(n) out Attentiong (x))), (") 9), (at( .)) in N.
(7)
Gout Attentiong (G, Cat(G, Cat (G.)))
(8)
5

---

Table I: Configurations of the cuboid attention patterns explored in the paper. The input tensor has shape (T, H, W). If "shift" or "strategy" is not given, we use shift (0, 0, 0) and strategy "local" by default. When stacking multiple cuboid attention layers, each layer will be coupled with layer normalization layers and feed-forward network as in the Pre-LN Transformer 48 The first
row shows the configuration of the generic cuboid attention.
Name
Configurations
Values
cuboid_sizz
(T,.Hr, N:]
(Ts, Hb)
(T.. Mr.
Generic Cuboid Attention
shift
(Ps, M,) + (P2. Ma)
- .... (PL, ML, M,]
stratogy
"1oc. di1."" "Ioc. /dil."
lloc./dil
Axial
cuboid size
(T. + (1, (1,1,M)
Divided Space- Time
cuhoid size
[T,1, + (1,
cuboid size shift
(P. -+ (P...
K)
Video-Swii P M
(0,0.0) (P/2.M/2,//2)
cuboid_siae size
(I. (1. - (1,M,
Local-Dilate--N Spatial Local- Dilate-M
strategy
"local" -> "local" - "dilated"
cuboid_sizi
(T, - (1, -+ (1, 1) (1.1,1/M) + (,/1, "local" "dilated" "local" "dilated" "local"
Axial Space Dilate-M
strategy
Here, Cat(:) flattens and concatenates input tensors. By combining Eqn. the overall computation of the cuboid attention layer with global vectors in Eqn.
1
we abbreviate
Xout CubAttna (X, , cuboid_ strategy, shift),
(9)
Gour Attn globul ())
The additional complexity caused by the global vectors is approximately O (THW P+P2). Given
that P usually small (in our experiments, P is most 8), the computational overhead induced by the global structure is negligible. The architecture of the cuboid attention layer is illustrated in Fig.
3.3 Hierarchical Encoder-Decoder Architecture
Earthformer adopts hierarchical encoder-decoder architecture illustrated in Fig. 2 The hierarchical architecture gradually encodes the input sequence to multiple levels of representations and generates the prediction via a coarse-to-fine procedure. Each hierarchy stacks D cuboid attention blocks. The cuboid attention block in the encoder uses one of the patterns described in Table [ and each cuboid block in the decoder adopts the "Axial" pattern. To reduce the spatial resolution of the input to cuboid attention layers, we include pair of initial downsampling and upsampling modules that consist of stacked 2D-CNN and Nearest Neighbor Interpolation (NNI) layers. Different from other papers that adopt Transformer for video prediction 19 30], we generate the predictions in a non-autt-reeressiv fashion rather than auto-regressive patch-by-patch fashion. This means that our decoder directly generates the predictions from the initial learned positional embeddings. We also conducted experiments with augo-regressive decoder based on visual codebook [33]. However, the ruto-regressive decoder underperforms the non-auto-rege ron-auto-regressvve decoder in terms of forecasting skill scores. The comparison between non-auto-regreessvve decoder and auto-regressive decoder is shown in Appendix C
Experiments
We first conducted experiments two synthetic datasets, MovingMNIST and newly proposed N-body MNIST, to verify the effectiveness of Earthformer and conduct ablation study on our design choices. Results these two datasets lead the following findings: 1) Among all patterns listed in Table "Axial" achieves the best overall performance; 2) Global vectors bring consistent performance gain with negligible increase in computational cost; 3) Using hierarchical coarse- to-fine structure can boost the performance. Based on these findings, we figured out the optimal design of Earthformer and compared it with other state-of-the-ar models on two real-world datasets: SEVIR 4 and ICAR-ENSQ On both datasets, Earthformer achieved the best overall performance. The statisties of all the datasets used in the experiments shown in Table 2 We normalized the data to the range [0, 1] and trained all the models with the Mean-Squared Error (MSE) loss. More implementation details are shown in Appendix
Dataset available at https////aaaaaaacaaccc//ccostttt

---

Table 2: Statistics of the datasets used in the experiments.
Size
Seq. Len.
Spatial Resolution
Datasel
train
Test
Out
x
MovingMNIST
8,100
900
1,000
10
10
61 x 61
N-body MNIST
20,000
1,000
1,000
10
64 64
SEVIR
35,718
9,060
12,159
13
12
381 x 381
ICAR-ENSO
5,205
1,667
24
Table 3: Ablation study the importance of adopting a hierarchical encoder-decoder. We conducted
experiments on MoVingMNIST. "Depth D means the model stacks cuboid attention blocks and there is no hierarchical structure. "Depth D1, D2 means the model stacks D1 cuboid attention blocks, applies the pooling layer, and stacks another D2 cuboid attention blocks.
Metrics MAE
Model
#Param, (M)
(PPS
MSE
SSIM 1
Depth
1.4
17.9
63.80
140.6
0.8324
Depth
3.1
36.3
52.46
114.8
0.8685
Depth
4.9
54.6
50.49
110.0
0.8738
Depth
6.6
73.0
48.04
104.6
0.8797
Depth 1.1
1.4
11.5
60.99
135.7
0.8388
Depth 2.2
3.1
18.9
50.41
106.9
0.8805
Depth 3.3
4.9
26.3
47.69
100.1
0.8873
Depth 4,4
6.6
33.7
46.91
101.5
0.8825
4.1 Experiments on Synthetic Datasets
MovingMNIIT. We follow 38 to use the public MovingMNIST datase The dataset contains 10,000 sequences. Each sequence shows 2 digits moving inside a 64 x 64 frame. We split the dataset to use 8, 8,100 samples for training. 900 samples for validation and 1,000 samples for testing. The task to predict the future 10 frames for each sequence conditioned on the first 10 frames.
N-body MNIST. The Earth is complex system in which an extremely large number of variables interact with each other. Compared with the Earth system, the dynamics of the synthetic MovingM- NIST dataset, in which the digits move independently with constant speed, over-simplified. Thus, achieving good performance on MovingMNIST does not imply that the model is capable of modeling complex interactions the Earth system. On the other hand, the real-worke Earth observation data, though experiencing rapid development, are still noisy and may not provide useful insights for model development. Therefore, we extend MovingMNIST to N-body MNIST, where N digits are moving with the N N-body motion pattern inside a 64 x 64 frame. Each digit has its mass and is subjected to the gravity from other digits. We choose N = in the experiments so that the digits will follow the chaotic 3-body motion 25.. The highly non-linear interactions in N N-body MNIST make it much more challenging than the original MovingMNIST. We generate 20,000 sequences for training, 1,000 for validation and 1, 1,000 for testing. Perceptual examples of the dataset can be found at the first two rows of Fig. 5 In Appendix D we demonstrate the chaotic behavior of N-body MNIST.
Hierarchical v.s. non-hierarchical. We choose "Axial" without global vectors our cuboid attention pattern and compare the performance of non-hierarchical and hierarchical architectures on MovingMNIIST. The ablation study on the importance of adopting a hierarchical encoder-decoder shown in Table kJ We can see that the hierarchical architecture has similar FLOPS with the non- hierarchical architectures while being better in MSE. This observation is consistent as we increase the depth until the performance saturates.
Cuboid pattern search. The design of cuboid attention greatly facilitates the search for optimal space-time attention. We compare the patterns listed in Table both MovingMNIST and N -odody MNIST to investigate the effectiveness and efficiency of different space-time attention methods on spatiotemporal forecasting tasks. Besides the previously proposed space-time attention methods, we also include new configurations that are reasonable and not computationally expensive in search space. For each pattern, we also compare the variant that uses global vectors. Results are summarized in Table We find that the "Axial" pattern is both effective and efficient and adding global vectors improves performance for all patterns while having similar FLOPS. We thus pick "Axial + global" as the pattern in Earthformer when conducting experiments real-worke datasets.
MMvinngNNIST http:////thaaatteooeaaassa////as

---

Table 4: Ablation study of different cuboid attention patterns and the effect of global vectors on MovingMNIST and N-body MNIST. The variant that achieved the best performance is in boldface while the second best is underscored. We also compared the performance of the cuboid attention patterns with and without global vectors and highlight the better one with grey background.
MovingMNIST MAE
N-hody MNIST
Model
#Param. (M)
GFLOPS
MSB 4
SSIM 1
MSB.
MAE L
SSIM
Axial
6.61
33.7
46.91
101.5
0.8825
15.89
41.38
0.9510
global *
7.61
34.0
41.79
92.78
0.8961
14.82
39,93
0.9538
DST
5.70
35.2
57.43
118.6
0.8623
18.24
45.88
0.9435
global
6.37
35.5
52,92
108.3
0.8760
17.77
45.84
0.9433
Video Swin 2x8
5.66
31.1
54.45
111.7
0.8715
19.89
49.02
0.9374
+ global
6.33
31.4
52.70
108.5
0.8766
19.53
48.43
0.9389
Video Swin 10x8
5.89
39.2
63.34
125.3
0.8525
23.35
53.17
0.9274
globall
6.56
39.4
62.15
123,4
0.8541
22.81
52.94
0.9293
Spatial Local-Gllbal 2 global
6.61
33.3
59.88
122.1
0.8572
23.24
54.63
0.9263
7.61
33.7
59:42
122.9
0.8565
21.88
52.49
0.9305
Spatial Gaauall 4 global
6.61
33.5
58.72
118.5
0.8600
21.02
49.82
0.9344
7.61
33.9
54.84
115.5
0.8585
19.82
48.12
0.9371
Axial Space Dilate 2
8.59
41.8
50.11
104.4
0.8814
15.97
42.19
0.9494
+ global
10.30
42.4
46.86
98.95
0.8884
15.73
41.85
0.9510
Axial Space Dilate 4
8.59
41.6
47.40
99.31
0.8865
19.49
51.04
0.9352
+ global
10.30
42.2
45.11
95.98
0.8928
17.91
46.35
0.9440
Table 5: Comparison of Earthformer with baselines on MovingMNIST and N --body MNIST.
Moving MNIST
N N-hody MNIST
Model
#Param. (M)
GFLOPS
MSE
MAE
SSIM
MSE
MAE
SSIM
UNet 3k
16.6
0.9
110.4
249.4
0.6170
38.90
94.29
0.8260
ConvI_STM
16
14.0
30.1
126.9
0.8477
32.15
72.64
PredRNN
35
23.8
232.0
52.07
108.9
0,8831
21.76
54.32
0,9288
PhyDNet
13
3.1
15.3
58. 70
124.1
0.8350
28.97
78.66
0.8206
E3D-LSTM
44
12.9
302.0
55.31
101.6
0.8821
22.98
62.52
0,9131
Rainformer
19.2
1.2
85.83
189.2
0,7301
38.89
96.47
0.8036
Earthformer w/o global
6.6
33.7
46.91
101.5
0.8825
15.89
41.38
0.9510
Fanthformer
34.0
41.79
92.78
0.8961
14.82
39.93
0.9538
Comparison to the state of the art. We evaluate six spatiotemporal forecasting algorithms: UNet [33] ConvLSTM (6). PredRNN [5). PhyDNet 13 E3D-LSTM 44]] and Rainformer []]
The results are in Table 5 Note that the MovingMNIST performance on several papers [13] is obtained by training the model with on-the-fly generated digits while we pre-generate the digits and train all models on a fixed dataset. Comparing the numbers in the table with numbers shown in these papers not fair. We train all baselines from scratch on both MovingMNIST and N-body MNIST using the default hyperparameters and configurations in their officially released code Fl
Qualitative results on N-body MNIST. Fig- 5 shows the generation results of different methods on a sample sequence from the N-body MNIST test set. The qualitative example demonstrates that our Earthformer is capable of learning rangrange interactions among digits and correctly predicting their future motion trajectories. Also, we can see that Earthformer able to more accurately predict the position of the digits with the help of global vectors. On the contrary, none of the baseline algorithms that achieved solid performance on MovingMNIST gives the correct and precise position of the digit "0" in the last frame. They either predict incorrect motion trajectories (PredRNN and E3D-LSTM), generate highly blurry predictions (Rainformer, UNet and PhyDNet) to accommodate the uncertainty about the future.
4.2 SEVIR Precipitation Nowcasting
Storm EVent ImageRy (SEVIR) 4 is spatiotemporally aligned dataset containing over 10,000 weather events. Each event consists of 384 km x 384 km image sequences spanning over 4 hours. Images in SEVIR were sampled and aligned across five different data types: three channels (C02,
'Except for Rainformer which originally has 212M parameters and thus suffers from overfitting severely.
8

---

Input
Target
46.34
CTRVLSTM
MSE=46.06
PhyDNet
MSE=38.93
E3D-LSTM
MSE=32.38
2
PredRNN
MSE-28.38
2
Earthformer w/o global
MSE-22.75
2
Earthformer
MSE-16.73
step
Figure 5: A set of examples showing the perceptual quality the predictions the N-body MNIST test set. From top to bottom: input frames, target frames, predictions by Rainformer (., UNet [3
ConvLSTM 36). PhyDNet [3]. E3D-LSTM 44], PredRNN 45 Earthformer without using global vectors, Earthformer. The results are sorted according to the MSE.
Table 6: Performance comparison on SEVIR. We include Cri tritical Success Index (CSI) besides MSE as evaluation metrics. The CSI, k.a intersection over union (IOU), is calculated different precipitation thresholds and denoted as CSI-thresh.
Model
GPLOPS
Menics
CS1 74.1
Prosideanc
UNet
15
0.3720
LSTM
0.1239
0.2675
0.7533
4,1702
Rainfore:
0.0831
0.2167
Farhifome plobal
13.1
253
0.4155
0.1572
0.2016
0.31
0.4214
0.6300
0.7637
3,300
Eathfiemer
15.1
257
0.4419
0.1791
0.2845
0,3333
0.4271
0.6860
0.7513
3,6457
CO9, C13) from the G0ES-16 advanced baseline imager, NEXRAD Vertically Integrated Liquid (VIL) mosaics, and GOES-16 Geostationary Lightning Mapper (GLM) flashes. The SEVIR bench- mark supports scientific research on multiple metcorological applications including precipitation nowcasting, synthetic radar generation, front detection, etc. We adopt SEVIR for benchmarking precipitation nowcasting, i.c., to predict the future VIL up to minutes (12 frames) given 65 minutes context VIL (13 frames). Fig. I shows example of VIL observation sequences in SEVIR.
Besides MSE, we also include the Critical Success Index (CSI), which is commonly used #Bits in precipitation nowcasting and is defined as CSI #Bita+ #F Alamms
count the #Hits (truth=1, pred=1), #Misses (truth=1, pred=0) and #F Alarms (truth=0, pred=1), the prediction and the ground-truth are rescaled back to the range 0-255 and binarized at thresholds [16, 74, 133, 160, 181, 219]. We report CSI at different thresholds and also their mean CSI-M.
SEVIR is much larger than MovIngMNIST and N-body MNIST and has higher resolution. We thus fair comparison.
slightly adjust the configurations of baselines based on those for MovingMNIST Detailed configurations are shown in Appendix A The experiment results are listed in Table Earth- former consistently outperforms baselines on almost all metrics and brings significant performance gain especially at high thresholds like 219, which are more valued by the communities.

---

Table 7: Performance comparison on ICAR-ENSO. C-Nino3.44 4-M and C-Nino3.4 WMM.. are the mean and the weighted mean of the correlation skill C over K forecasting steps. C-Nino3 4M assigns more weights to longer-tern prediction scores. MSE is calculated between the spatiotemporal SST anomalies prediction and the corresponding ground-truth.
Model
#Pamm. (M)
GILLOPS
C-Wino3
C-Mine 4-MX
Metrics
MSE 1)
Pensistence
0.3221
0.447
4.581
UNet
[k]
12.1
0.6926
2.102
2,868
ConLLSTM
14.0
ProdRNN
23.8
85.8
0.6492
1,910
3.444
PNy
[
3.1
5.7
0.6646
1,965
2,708
E3D-LSTM
44
12.9
99.8
0,7040
2. 125
3,095
Rainformer
19.2
0.7106
2.153
3,043
Earthformer w/o global
Earthformer
7.6
23,9
0.7329
2.259
2,546
4.3 ICAR-ENSO Sea Surface Temperature Anomalies Forecasting
El Niñio/Southern Oscillation (ENSO) has wide range of associations with regional climate extremes and ecosystem impacts. ENSO sea surface temperature (SST) anomalies forecasting for lead times up to one year (12 steps) valuable and challenging problem. Nino3.4 44 index, which the area- averaged SST anomalies across certain area (1709-120%, N))-5"N) of the Pacific, serves crucial indicator of this climate event. The forecast quality is evaluated by the correlation skill [T5]
of the siree-month-mene moving-averagede Nino3.4 index Ciio33.4
IR
(X=
calculated on the whole test set of size N, where Y E RRNK is the ground-truth of K-step Nino3.4 index, X R- is the corresponding prediction of Nino3.4 index.
ICAR-ENSO consists historical climate observation and stimulation data provided by Institute for Climate and Application Research (ICAR). We forecast the SST anomalies up to 14 steps (2 steps more than one year for calculating
given a context of 12 steps of SST anomalies observations. Table 7 compares the performance of our Earthformer with baselines
the ICAR-ENSO dataset. We report the mean correlation skill C-Nino3.4: M-MM K Σn C CMiin3.4 and the weighted mean correlation skill C-Nino3.4-NM as well as the MSE between the spatiotemporal SST anomalies prediction and the corresponding ground-truth truth. We can find that Earthformer consistently outperforms the baselines in all concerned evaluation metries and that using global vectors further improves the performance.
Σx Ok -C-Ciinn... over K = 12 forecasting steps
5 Conclusions and Broader Impact
In this paper, propose Earthformer, space-time Transformer for Earth system forecasting. Earthformer is based on a generic and efficient building block called Cuboid Attention. achieves SOTA on MovingMNIST, our newly proposed N-body MNIST, SEVIR, and
Our work has certain limitations. The first one that Earthformer a deterministic model that does not model uncertainty. This may result in predicting the average of all plausible futures, causing the model to generate blurry predictions of low perceptual quality and be lack of valuable small-scale details. In fact, the community lacks appropriate metrics that measure the uncertainty component in Earth system forecasting models. Extending Earthformer to a probabilistic forecasting model can be an exciting future direction. We include more detailed discussions and preliminary experiments about handling uncertainty Appendix C The second one that the model is purely data-driven and does not take advantage of the physical knowledge of the Earth system. Recent studies on adding physical constraints 220 26 and ensembling the predictions from a data-driven model and physics- based model 3 imply that it is active and promising research direction pursue. We plan to study how to incorporate physical knowledge into Earthformer in the future.
Acknowledgments and Disclosure of Funding
This work has been made possible by Research Impact Fund project (R6003-21) and Innovation and Technology Fund project (ITS/004/2IFP) funded by the Hong Kong Government.
aa be In k. where by L.5, for < 4; by 2. 4 l S 11: by 3, for & 11.
10

---

References
[1] Tom R Andersson, Scott Hosking, María Pérez-Ortiz, Brooks Paige, Andrew Elliott, Chris Russell, Stephen Law, Daniel C Jones, Jeremy Wilkinson, Tony Phillips, et al. Seasonal Arctic sea ice forecasting with probabilistic deep learning. Nature communications, 12(1):1-12, 2021.
[2] Anurag Arna,, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luvcić, and Cordelia Schmid. VIVIT: A video vision transformer. In International Conference Computer Vision (ICCV), 2021.
[3] Cong Bai, Feng Sun, Jinglin Zhang, Yi Song, and Shengyong Chen. Rainformer: Features extraction balanced network for radar-based precipitation nowcasting. IEEE Geoscience and Remote Sensing Letters, 19:1-5, 2022.
[4] Gedas Bertasius, Heng Wang. and Lorenzo Torresani. Is space-time attention all you need for video understanding. arXiv preprint arXiv:2 050000 2(3):4, 2021.
[5] Boyu Chen, Peixia Li, Chuming Li, Baopu Li. Lei Bai, Chen Lin, Ming Sun, Junjie Yan, and Wanli Ouyang. GLiT: Neural architecture search for global and local image transformer. In Proceedings of the IEEF/CVF International Conference Computer Vision pages 12-21, 2021.
[6] Christian Schroeder de Witt, Catherine Tong, Valentina Zantedeschi, Daniele De Martini, Freddie Kalaitzis, Matthew Chantry, Duncan Watson-Parris, and Piotr Bilinski. RainBench: towards global precipitation forecasting from satellite imagery. In AAAI, 2021.
[7] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 2018.
[8] 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 1xXi=:2000.11929, 2020.
[9] Lasse Espeholt, Shreya Agrawal, Casper Sønderby, Manoj Kumar, Jonathan Heek, Carla Bromberg, Cenk Gazen, Jason Hickey, Aaron Bell, and Nal Kalchbrenner. Skillful twelve hour precipitation forecasts using large context neural networks. arXiv preprint arXiv:2111.0777 2021.
[10] Veronika Eyring, Sandrine Bony, Gerald Meehl, Catherine A Senior, Bjorn Stevens, Ronald Stouffer, and Karl E Taylor. Overview of the coupled model intercomparison project phase 6 (CMIP6) experimental design and organization. Geoscientific Model Development, 9(5): 1958, 2016.
[11] Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. In ICCV, 2021.
[12] Steven Goodman, Timothy J Schmit, Jaime Daniels, and Robert J Redmon. The GOES-R series: new generation of geostationary environmental satellites. Elsevier, 2019.
[13] Vincent Le Guen and Nicolas Thome. Disentangling physical dynamics from unknown factors for unsupervised video prediction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1 11474-1 11484, 2020.
[14] John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro, Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:21 135877 2021.
[15] Yoo-Geur Ham, Jeong- Hwan Kim, and Jing-Jia Luo. Deep learning for multi-year ENSO forecasts. Nature, 573(7775):568-5777 572, 2019.
[16] Kai Han, An Xiao, Enhua Wu, Jianuaan Guo, Chunjing Xu, and Yunhe Wang. Transformer in transformer. NeurIPS, 2021.
[17] William H Heiss, David L. McGrew, and Dale Sirmans. NEXRAD: next generation weather radar (wsr-88d). Microwave Journal, 33(1):79-89, 1990.
[18] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
axXiv:1606.08415,200 2016.

---

[19] Jonathan Ho. Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv preprint arXiv:1912.12180, 2019.
[20] Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Characterizing possible failure modes in physics-informan neural networks. Advances in Neural Information Processing Systems, 34:26548-266 265660 2021.
[21] Christophe Letellier. Chaos in nature, volume 94. World Scientific, 2019.
[22] 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 the IEEE/CVF International Conference on Computer Vision, pages 10012-10022, 2021.
[23] Ze Liu, Jía Ning, Yue Cao, Yixuan Wei, Zheng Zhang, Stephen Lin, and Han Hu. Video swin transformer. arXiv preprint arXiv:2 13230,, 2021.
[24] Mario Lucic, Karol Kurach, Marcin Michalski, Sylvain Gelly, and Olivier Bousquet. Are GANs created equal? large-scale study. In NeurlPS, 2018.
[25] Valtonen MJ, Mauri Valtonen, and Hannu Karttunen. The three-body problem. Cambridge University Press, 2006.
[26] Geoffrey Négiar, Michael W Mahoney, and Aditi S Krishnapriyan. Learning differentiable solvers for systems with hard constraints. arXiv preprint urXiv:2207.0866 2022.
[27] Daniel Neimark, Omri Bar, Maya Zohar, and Dotan Asselmann. Video transformer network. 2021.
arXiv preprint arXiv:2102.000
[28] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadeneshelli et al. FourCastNet: global dataen high-resotution weather model using adaptive fourier neural operators. arXiv preprint arXiv: 1214, 2022.
[29] Mandela Patrick, Dylan Campbell, Yuki M. Asano, Ishan Misra Florian Metze, Christoph Feichtenhofer, Andrea Vedaldi, and João Henriques. Keeping your eye on the ball: Trajec- tory attention in video transformers. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[30] Ruslan Rakhimov, Denis Volkhonskiy, Alexey Artemov, Denis Zorin, and Evgeny Burnaev. Latent video transformer. arXiv preprint arXiv:2006. 10704, 2020.
[31] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pages 8821-8831. PMLR, 2021.
[32] Suman Ravuri, Karel Lenc, Matthew Willson, Dmitry Kangin, Remi Lam, Piotr Mirowski, Megan Fitzsimons, Maria Athanassiadou, Sheleem Kashem, Sam Madge, et al. Skilful pre- cipitation nowcasting using deep generative models of radar. Nature, 597(7878):722-677, 2021.
[33] Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with VQ-VAE-2. Advances in neural information processing systems, 32, 2019.
[34] Markus Reichstein, Gustau Camps- Valls, Bjorn Stevens, Martin Jung, Joachim Denzler, Nuno Carvalhais, al. Deep learning and process understanding for data-driven earth system science.
116677743):155-244 2044, 2019.
[35] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma. PixeICNN++: A PixeICNN implementation with discretized logistic mixture likelihood and other modifications. In ICLR, 2017.
[36] Xingjian Shi, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional LSTM network: A machine learning approach for precipitation nowcasting. In NeurlPS, volume 28, 2015.
[37] Xingjian Shi, Zhihan Gao, Leonard Lausen, Hao Wang, Dit-Yan Yeung, Wai-kin Wong, and Wang-chun Woo. Deep learning for precipitation nowcasting: A benchmark and model. In NeurIPS. volume 30, 2017.
[38] Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov. Unsupervised learning of video representations using LSTMs. In ICML, pages 843-852. PMLR, 2015.
12

---

[39] Byron D Tapley, Srinivas Bettadpur. John C Ries, Paul F TThmmson, and Michael M Watkins. GRACE measurements of mass variability in the earth system. science, 305(5683):503-500) 2004.
[40] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
[41] Ashish Vaswani, Prajit Ramachandran, Aravind Srinivas, Niki Parmar, Blake Hechtman, and Jonathon Shlens. Scaling local self-attention for parameter efficient visual backbones. In CVPR, 2021.
[42] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz. Kaiser, and Illia Polosukhin. Attention all you need. NeurlPS, volume 30, 2017.
[43] Mark Veillette, Siddharth Samsi, and Chris Mattioli. SEVIR: A storm event imagery dataset for deep learning applications in radar and satellite metcorology. Advances in Neural Information Processing Systems, 33:22009-22019, 2020.
[44] Yunbo Wang, Lu Jiang, Ming-Hsuan Yang. Li-Jia Li, Mingsheng Long, and Li Fei-Fei. Eidetic
3D LSTM: model for video prediction and beyond. In International conference on learning representations, 2018.
[45] Yunbo Wang, Haixu Wu, Jianjin Zhang, Zhifeng Gao, Jianmin Wang, Philip Yu, and Mingsheng Long. PredRNN: A recurrent neural network for spatiotemporal predictive learning. IEEE Transactions Pattern Analysis and Machine Intelligence, 2022.
[46] Dirk Weissenborn, Oscar Täckström, and Jakob Uszkoreit. Scaling autoregressive video models. In International Conference on Learning Representations, 2019.
[47] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3-19, 2018.
[48] Ruibin Xiong, Yunchang Yang, Di Hc, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533, PMLR, 2020.
[49] Shen Yan, Xuehan Xiong, Anurag Arnab, Zhichao Lu, Mi Zhang, Chen Sun, and Cordelia Schmid. Multiview transformers for video recognition. In CVPR, 2022.
[50] Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. VideoGPT: Video generation using vq-vae and transformers. arXiv preprint arXiv:2104. 10157, 2021.
[51] Jianwei Yang, Chunyuan Li, Pengchuan Zhang, Xiyang Dai, Bin Xiao, Lu Yuan, and Jianfeng Gao. Focal self-attention for local-global interactions in vision transformers. In NewrIPS, 2021.
[52] Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie Yan Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34, 2021.
[53] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33:1 11117288872277 2020.
[54] Yanyi Zhang, Xinyu Li, Chunhui Liu, Bing Shuai, Yi Zhu, Biagio Brattoli, Hao Chen, Ivan Marsic, and Joseph Tighe. Vidtr: Video transformer without convolutions. In Proceedings of the IEEE/CVF International Conference on Computer Vision pages 13577-13587, 2021.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section S
(c) Did you discuss any potential negative societal impacts of your work? [No] Exploring
Earth system forecasting has no negative societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [No]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any assets either in the supplemental material URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you's 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "wRi8Y6z2ykb", "reviewer": "Reviewer_fTQV", "summary": "The authors propose a spacetime transformer and apply it to forecasting events in Earth systems. In detail:\n- the authors propose Cuboid attention mechanism which splits the input tensor to non-overlapping parts\n  - to reduce computational cost by local attention\n  - to compute self-attention in parallel\n- the authors generalize previous studied video transformer by stacking cuboid attention layers with various hyperparameters\n- the authors fulfill the communication gaps of local Cuboid by introducing shared global vectors\n- the authors propose an N-body MNIST dataset developed from MINIST\n- the proposed Earthformer achieves SOTA performance on both SEVIR and ICAR-ENSO", "questions": "- I am not sure why N-body MNIST is related to the chaotic 3-body motion problem and thus performing good on this dataset means the model is capable to handle complex system. It would be better if the authors illustrate the reasons and intuitions behind this?\n- It would be better if you can give some training details. For example, how you encode the input data and what is the dimension of the embeddings.\n", "limitations": "- The authors do not address the limitation adequately\n- This framework is not physical dynamic-aware such as the PhyDNet paper which is also mentioned by the authors. To handle a complex system, I think it is important to introduce knowledge from the dynamics.\n- There are few analyses showing why this framework shall only be specialized on this task. They might extend this framework to video frame prediction tasks and comparing to the SOTA in these datasets (e.g. UCF101).\n- It would be nicer if they can post some failure cases and discuss the possible issues.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n - originality: From my knowledge, it is one of the first paper working on predicting ICAR-ENSO using transformer.\n - quality: The authors use both synthetic dataset--N-body MNIST and SEVIR, ICAR-ENSO to proof the significance improvement of their work.\n - clarity: The Cuboid Attention is introduced very clearly with formula and detailed figure illustration\n - significance: The improvement of the performance on the model is significant\n\nWeaknesses:\n - originality: There are a lot of work of transformers with local attention and shared global connection between these local attentions. The model architecture is not novel from a machine learning model aspects.\n - quality: It would be nicer if the author can use other transformer based space-time model to train on these benchmarks and compare the performance with them. After all, ConvLSTM is from 2015, PredRNN and PhyDNetis not based on transformer, so does E3D-LSTM which is from 2017. Rainformer is the only one using transformer, but not specialized in this task.\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "sXF42U7z4S5", "reviewer": "Reviewer_DF9s", "summary": "This paper proposed cuboid attention as an alternative to regular attention for forecasting Earth system-like data. The paper presents results on a toy N-body MNIST dataset and two real datasets, one on precipitation nowcasting and one one El-Nino Southern Oscillation\n\n\n\n", "questions": "1. The motivation behind not using CNN/RNN is that the inductive bias might not capture the internal variability of the climate system. Still, the proposed architecture also has not been studied in terms of why it should.\n\n2. It's important to understand that for problems in physics, a single number like MSE and MAE may not be a good representative of actually how well the algorithm performs in terms of forecasting. For example, in the motivation, the authors talk about how internal climate variability can affect the prediction performance-- that's a solid motivation but is completely absent in the analysis. Does their model capture large scales better or do they capture small scales for a longer period of time? All of these can be very well explored with a Lorenz-type system instead of this MNIST data set. Weather is more than just chaotic--it is multi-scale without any scale separation, unlike the MNIST data\n\n2. The El Nino and nowcasting predictions results need further elaboration. It's hard to say whether it's good or bad. Also, there are good baselines for both these problems. To my knowledge, this paper is SOTA for DL-based ENSO:https://www.nature.com/articles/s41586-019-1559-7. Can the authors compare their model with the proposed model? \n\n3. For precipitation nowcasting, I believe this paper is SOTA:https://www.nature.com/articles/s41586-021-03854-z. Is it possible to compare the results directly against this paper?\n\n4. I would finally also like to point out, that even if the proposed model is not outperforming some of the baselines it does not necessarily take anything away from the architecture, but the authors need to do some analysis and show the usefulness or show the connection between how this architecture is really useful for physical/earth-system data. In the current form, it's not really there", "limitations": "N/A", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\nThe novelty in this paper lies in the fact that they use cuboid attention in their transformer architecture. The motivation is that internal variability in climate may not be captured by the inductive biases of CNN/RNN type architectures. However, there is no reason to think that transformer architectures would also capture internal variability just as much as CNN/RNN or for that matter any DL architecture\n\nWeakness:\n\nThis paper has several weaknesses, starting from the motivation behind using cuboid attention and baselines. Furthermore, while I found the N-body system's example as a nice toy example to prove that the proposed algorithm has superiority over other baselines, it is actually not that relevant to weather. I detail it further in my comments. ", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "mkdWTAiE3aU", "reviewer": "Reviewer_LEpa", "summary": "This paper proposes the framework of \"cuboid attention\" for thinking about different spatio-temporal attention mechanisms. This framework includes vanilla self-attention and various lower-complexity alternatives such as axial attention or divided space-time attention as special cases.\n\nFurthermore, this paper proposes Earthformer, a spatio-temporal transformer model with a particular set of cuboid attention blocks. The Earthformer architecture further features a set of global vectors that all cuboids attend to in order to more-readily share information between cuboids.\n\nThis paper also proposes a dataset called N-body MNIST (with N=3), where each dataset example is a sequence 64x64 frames consisting of 3 MNIST digits moving according to chaotic 3-body motion. This dataset contains 22,000 examples.\n\nFinally, the authors evaluate the Earthformer architecture against several existing algorithms (UNet, ConvLSTM, PredRNN, PhyDNet, E3D-LSTM, and Rainformer) on several datasets (MovingMNIST, N-body MNIST, SEVIR, and ICAR-ENSO). The results show that Earthformer consistently outperforms the other models, sometimes by a large margin.", "questions": "**Q1**: To me, it was unclear whether the results of baseline models reported in Tables 4-6 were computed by the Earthformer authors, or the original authors of those baseline models. Could the Earthformer authors please clarify?\n\n**Q2**: Will the N-body MNIST dataset be released publicly?\n\n**Q3**: Why call the dataset \"N-body MNIST\" instead of \"3-body MNIST\" if you only tried N=3? Given that you only tried N=3, I think it's a bit misleading to write N-body.\n\n**Q4**: You write for cuboid attention on its own, \"the cuboids do not communicate with each other.\" This sentence is a bit vague, and perhaps misleading. If divided space-time attention is applied sequentially (first space, then time), then there is information passing between all cuboids, right?", "limitations": "To me, there are no clear limitations / negative impact. The authors did not discuss any negative impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n**S1: New N-body MNIST dataset**\\\nCompared to MovingMNIST, the N-body MNIST dataset intuitively seems more challenging and representative of realistic dynamical systems. I believe that this dataset could well serve as a standard benchmark for future models.\n\n**S2: Comprehensive ablation study on the benefit of global attention vectors**\\\nI appreciated the large amount of ablation studies showing the significant performance gain with using global attention vectors on the MovingMNIST and N-body MNIST datasets, regardless of the attention variant.\n\n**S3: Generalization of attention via Cuboid Attention framework**\\\nThe cuboid attention framework for thinking about attention mechanisms serves to provide a common language for discussing various forms of attention. In particular, the authors show how axial attention and divided space-time attention are special cases of cuboid attention.\n\n**S4: Overall impressive results**\\\nThe comparison tables and figure really demonstrate the superiority of Earthformer against the baseline models.\n\n## Weaknesses\n\n**W1: Missing some experimental datasets**\\\nWhile the authors show good results of Earthformer on MovingMNIST, N-body MNIST, SEVIR, and ICAR-ENSO, I would be interested in seeing experiments on other types of spatio-temporal datasets, such as action-recognition (Kinetics-400/600, SSV2, Diving-48, or YouTube-8M).\n\n**W2: Missing some comparison models**\\\nWhile the set of model comparisons is large, to me it seems that some more generic models are missing. For example, how would a standard TimeSformer model (with divided space-time attention) do on the experiments? A comparison against TimeSformer, for example, would really help cement the importance of (1) cuboid attention, and (2) global attention vectors. An appropriate comparison would also be against TimeSformer, but with global attention added to the TimeSformer model.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "dAyrYaRFfo3", "reviewer": "Reviewer_YCyZ", "summary": "In this paper, the authors propose an earth system forecast model based on transformers. This is achieved by first dividing the entire space-time field as non-overlapping cuboid blocks, applying attention in each cuboid block together with several \"global vectors\" that allow for interactions between the individual blocks, and then merging the blocks together to yield the predictions. They find state-of-the-art results on several synthetic and real datasets including moving MNIST, N-body MNIST (an extension of moving MNIST), precipitation nowcasting and SST anomaly forecasting.", "questions": "There are several lines throughout the text that I found difficult to comprehend. Please could you clarify this?\n- Lines 238-240:\n> Note that the MovingMNIST performance on several papers [13] are obtained by training the model with on-the-fly generated digits while we pre-generate the digits and train all models on a fixed dataset.\n\nWhat are \"on-the-fly generated digits\"? What was the reasoning behind training the model differently in this paper vs the other papers?\n\n- Lines 184-185:\n> generates the prediction via a coarse-to-fine procedure\n\nWhat exactly is the coarse-to-fine procedure?\n\n- What do the subscripts $(a,b,c)$ in the latent visual code $Z_{(a,b,c)}^{\\text{code}}$ mean in the VQ-VAE? This notation is not introduced. Also, shouldn't $Z^{\\text{code}}$ in the LHS of equation (9), Appendix D be $Z_i^{\\text{code}}$ (since the RHS has index $i$)?\n\n- Could you please explain in more details the idea behind auto-regressive Earthformer? This is to incorporate temporal correlations? But if so, shouldn't the global vectors already take that into account?\n\n\n", "limitations": "One limitation is that the model may have good performance with respect to standard machine learning metrics, however in practice, this may not necessarily translate to their usability, as demonstrated in the recent Deepmind paper on precipitation nowcasting (Ravuri et al. 2021). This is already visible in Figures 7-10, where the blurry predictions have better MSE score than the more realistic, high fidelity predictions that are clearly better. This indicates a search for better metrics for earth system forecasting or approaches that don't rely on loss optimisation.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Thorough benchmarking to demonstrate the performance on real and synthetic data in addition to nice ablation studies.\n- Results on the N-body MNIST look great. Predictions are much less blurry than the baselines, which are a common problem for these data-driven nowcasting/forecasting models.\n\nWeaknesses:\n- In terms of novelty, with so many recent papers applying transformers to vision tasks, the idea itself is not very original. There are many papers that applies transformers to video prediction tasks, e.g. \"Anticipative Video Transformer (Girdhar and Grauman 2021)\", that the authors also acknowledge. However, I have not seen any paper that applies cuboid attention, so this paper may be novel in this respect.\n- Comparisons with physics-based models would have been nice to see as for many problems in weather forecasting, physics-based model still perform better than data-driven models. It would be interesting to know how the transformer approach fare against these traditional modelling approaches.\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept