[Paper Review] Earthformer: Exploring Space-Time Transformers for Earth System Forecasting

Zhihan Gao, Xingjian Shi+5 more|arXiv (Cornell University)|Jul 12, 2022

Computational Physics and Python Applications101 citations

TL;DR

Earthformer proposes Cuboid Attention space-time transformers for Earth system forecasting, achieving state-of-the-art results on precipitation nowcasting (SEVIR) and ENSO SST forecasting (ICAR-ENSO) by combining local cuboid self-attention with a small set of global vectors.

ABSTRACT

Conventionally, Earth system (e.g., weather and climate) forecasting relies on numerical simulation with complex physical models and are hence both expensive in computation and demanding on domain expertise. With the explosive growth of the 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 Nino/Southern Oscillation (ENSO) forecasting show Earthformer achieves state-of-the-art performance. Code is available: https://github.com/amazon-science/earth-forecasting-transformer .

Motivation & Objective

• Motivate improved forecasting of Earth system variabilities using data-driven DL methods that leverage spatiotemporal observations.
• Develop a scalable space-time Transformer architecture suitable for high-dimensional Earth data.
• Design and evaluate a generic cuboid attention mechanism to enable efficient local attention with global information sharing.
• Demonstrate the effectiveness of a hierarchical encoder-decoder for coarse-to-fine forecasting.
• Investigate design choices (patterns, global vectors, hierarchy) to identify an optimal Earthformer configuration.

Proposed method

• Introduce Cuboid Attention that decomposes input tensors into non-overlapping cuboids and applies self-attention within each cuboid in parallel.
• Reduce attention complexity to approximately O(THW · bT bH bW) by operating inside cuboids of size (bT, bH, bW).
• Incorporate a collection of global vectors G that attend to all cuboids and are updated via a separate global attention step.
• Attach a hierarchical encoder-decoder with initial downsampling/upsampling to enable coarse-to-fine prediction.
• Use non-auto-regressive decoding to generate K-step forecasts directly from learned positional embeddings (with an auto-regressive variant evaluated in Appendix C).
• Explore multiple cuboid attention patterns (e.g., Axial, Divided Space-Time, Video-Swin, Spatial Local-Global) and add global vectors to assess performance benefits.

Experimental results

Research questions

• RQ1 Can a space-time Transformer with cuboid-attention match or exceed state-of-the-art performance in Earth system forecasting tasks?
• RQ2 What cuboid attention configurations (size, strategy, shift) and architectural choices (global vectors, hierarchy) yield the best forecasting accuracy and efficiency for real-world Earth data?
• RQ3 Do global vectors and hierarchical design provide consistent improvements across synthetic and real-world datasets?
• RQ4 How does Earthformer compare to established baselines on precipitation nowcasting and ENSO forecasting in terms of both accuracy and computational cost?
• RQ5 Are Axial-type cuboid patterns with global vectors a robust, efficient choice across tasks?

Key findings

• Axial pattern with global vectors delivers best overall performance in synthetic experiments and real-world benchmarks.
• Global vectors provide consistent performance gains with negligible computational overhead.
• A hierarchical encoder-decoder improves forecasting accuracy without increasing FLOPS substantially.
• Earthformer achieves state-of-the-art results on SEVIR precipitation nowcasting and ICAR-ENSO ENSO forecasting datasets.
• Earthformer outperforms baselines such as UNet, ConvLSTM, PredRNN, PhyDNet, E3D-LSTM, and Rainformer on the reported tasks when using the proposed design.
• In synthetic tests, Earthformer with Axial + global patterns demonstrates strong modeling of long-range interactions and dynamics.