PaperlayerView exact revision on arXiv

Source-linked AI summary

Earthformer: Exploring Space-Time Transformers for Earth System Forecasting

Zhihan Gao, Xingjian Shi, Hao Wang, Yi Zhu, Yuyang Wang, Mu Li, Dit-Yan Yeung

arXiv:2207.05833v2cs.LGcs.AIcs.CV

TL;DR

Earth system forecasting is important but conventional physical simulation is costly and limited in using emerging observations and massive Earth-observation data. Earthformer introduces cuboid-based space-time attention with global vectors in a hierarchical encoder-decoder architecture, achieving state-of-the-art performance across synthetic and real-world forecasting benchmarks. Its main limitations are deterministic uncertainty handling and the absence of explicit physical knowledge.

  • Problem

    Conventional Earth system forecasting relies on expensive, expertise-demanding physical simulation, while Transformer adoption remains limited despite the growth of spatiotemporal Earth-observation data.

  • Method

    Earthformer uses efficient cuboid-level space-time self-attention, global vectors for communication across local cuboids, and a hierarchical encoder-decoder architecture.

  • Results

    Earthformer achieves state-of-the-art performance on MovingMNIST, N-body MNIST, SEVIR precipitation nowcasting, and ICAR-ENSO forecasting.

  • Takeaways & Limitations

    Cuboid Attention provides a generic, efficient building block for space-time Transformers applied to Earth system forecasting.

  • Takeaways & Limitations

    Earthformer is deterministic and purely data-driven, so it does not model uncertainty or explicitly use physical knowledge of the Earth system.

Abstract

from arXiv · show

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 .

1 Introduction

Earth system forecasting has substantial socioeconomic importance, but conventional simulation systems are costly and limited in using emerging observations and massive Earth-observation data. Earthformer addresses limited Transformer adoption with efficient cuboid-based space-time attention and achieves state-of-the-art results on real-world precipitation and ENSO forecasting benchmarks.

  • Motivation: Earth system variability affects crop yields, transportation, floods, and fires, making precise and timely forecasting economically and socially valuable.Forecasting can help people take precautions and better use wind and solar resources.
  • Motivation: Simulation-based operational forecasting relies on meticulous physical modeling and struggles to incorporate emerging geophysical observations and Petabytes-scale Earth-observation data.The passage identifies these limitations for operational weather and climate systems, including HREF and NOAA precipitation nowcasting.
  • Motivation: Transformers suit Earth systems because attention captures complex, long-range dependencies, yet their adoption for Earth system forecasting remains limited.Deep learning models instead train on Earth observation data and have shown success in precipitation nowcasting and ENSO forecasting.
  • Earthformer: Earthformer decomposes inputs into non-overlapping cuboids, applies parallel local self-attention, and uses global vectors to connect local regions.Limiting O(N^2) attention to local cuboids reduces complexity, while global vectors gather system status and share information among cuboids.
  • Evaluation: Synthetic experiments use MovingMNIST and chaotic N-body MNIST to assess cuboid attention and identify an effective Earthformer design.N-body MNIST introduces interacting three-body dynamics to make synthetic forecasting more challenging and more relevant to Earth system forecasting.
  • Results: Earthformer achieves state-of-the-art performance on precipitation nowcasting and ENSO forecasting benchmarks.The real-world benchmarks are the SEVIR precipitation-nowcasting benchmark and the ICAR-ENSO dataset.

2 Related Work

Related work has developed convolutional and recurrent architectures for Earth system forecasting alongside efficient space-time and local-attention Transformers for video and vision tasks. These approaches motivate Earthformer's space-time attention design.

  • Deep learning architectures for Earth system forecasting: CNN-, RNN-, and hybrid CNN-LSTM architectures have been applied to precipitation nowcasting, sea-ice prediction, and ENSO forecasting.Examples include U-Net, ConvLSTM, PredRNN, and E3D-LSTM.
  • Space-time Transformers for video modeling: Video Transformers use divided, axial, factorized, and separable attention to reduce the memory demands of joint spatiotemporal attention.These methods adapt Transformer architectures to space-time video modeling.
  • Global and local attention in vision Transformers: Vision Transformers improve attention efficiency through local-attention designs and CNN-inspired locality mechanisms.Examples include HaloNets, GLiT, and Focal Transformer.

3 Model

Earthformer formulates Earth system forecasting as spatiotemporal sequence prediction and uses a hierarchical encoder-decoder built around efficient cuboid attention. Cuboid attention decomposes inputs, performs local self-attention, reconnects regions through global vectors, and supports multiple attention patterns.

  • Forecasting formulation: Earthformer represents observations as a spatiotemporal sequence and predicts a K-step-ahead target sequence with spatial resolution H × W.Input and output channels may differ between observations and targets.
  • Motivation: Vanilla self-attention on a tensor shaped (T, H, W) has complexity O(T^2H^2W^2), making direct application computationally infeasible.The high dimensionality of Earth system data creates a central challenge for Transformers.
  • Cuboid Attention: Cuboid attention decomposes the input tensor into non-overlapping cuboids, applies self-attention within each cuboid in parallel, and merges the outputs back to the original shape.The generic layer consists of “decompose”, “attend”, and “merge” steps.
  • Cuboid Attention: Cuboid decomposition varies through cuboid size, local or dilated strategy, and shift offsets, enabling different space-time correlations to be modeled.The implementation pads inputs when dimensions are not divisible by the cuboid size.
  • Efficiency: Restricting self-attention to local cuboids reduces complexity to approximately O(THW · bT bH bW), which scales linearly with cuboid size.Because cuboids can be much smaller than the input tensor, the layer is more efficient than full attention.
  • Attention patterns: Cuboid attention can reproduce axial, video Swin-Transformer, and divided space-time attention patterns by stacking layers with different configurations.The paper manually selects reasonable, computationally inexpensive patterns and leaves automatic neural architecture search for future work.
  • Global communication and architecture: Global vectors connect otherwise isolated cuboids by allowing local elements to attend to shared vectors, which aggregate information from the full input tensor.Earthformer uses a hierarchical encoder-decoder with coarse-to-fine prediction; each hierarchy stacks D cuboid attention blocks.

4 Experiments

Experiments on synthetic datasets support Earthformer’s hierarchical and attention design, while benchmarks show strong performance for precipitation nowcasting and ENSO forecasting.

  • Synthetic datasets: Synthetic experiments identify Axial cuboid attention, global vectors, and hierarchical coarse-to-fine encoding as effective design choices.Axial performs best overall; global vectors consistently improve performance with negligible computational cost, and hierarchy boosts performance.
  • Synthetic datasets: N-body MNIST uses chaotic three-body interactions to provide a more challenging and Earth-system-relevant test than MovingMNIST.Three digits interact through gravity, producing highly nonlinear dynamics unlike independently moving digits.
  • Synthetic datasets: Hierarchical architectures achieve better MovingMNIST MSE than non-hierarchical architectures with similar FLOPS.The advantage remains as depth increases until performance saturates.
  • Synthetic datasets: Earthformer learns long-range interactions and more accurate digit trajectories on N-body MNIST, whereas several baselines produce incorrect trajectories or blurry predictions.Qualitative comparisons include UNet, ConvLSTM, PhyDNet, E3D-LSTM, PredRNN, Rainformer, and Earthformer variants.
  • Real-world benchmarks: On SEVIR, Earthformer consistently outperforms baselines on almost all metrics, with especially large gains at the high CSI-219 threshold.The benchmark predicts 12 future VIL frames from 13 context frames and evaluates MSE and thresholded CSI.

5 Conclusions and Broader Impact

Earthformer combines a generic Cuboid Attention building block with space-time Transformer forecasting and achieves state-of-the-art performance across four datasets. The authors identify uncertainty modeling and physical knowledge integration as important limitations and future directions.

  • Earthformer is a space-time Transformer built around the generic and efficient Cuboid Attention block.The paper’s conclusion identifies Cuboid Attention as Earthformer’s core building block.
  • Earthformer achieves state-of-the-art performance on MovingMNIST, N-body MNIST, SEVIR, and ICAR-ENSO.
  • As a deterministic model, Earthformer does not represent uncertainty and may produce blurry predictions lacking valuable small-scale details.The authors describe probabilistic forecasting as a future direction.
  • Earthformer is purely data-driven and does not use physical knowledge of the Earth system.The authors identify incorporating physical knowledge as an active future research direction.

Checklist

The checklist reports that the paper addresses its claimed contributions and limitations, provides reproducibility materials and training details, and documents asset usage. It reports no error bars and no negative societal impacts for this work.

  • The authors affirm that the paper describes its contributions, scope, and limitations, with limitations discussed in Section 5.
  • The authors report no potential negative societal impacts for Earth system forecasting.
  • The paper provides code, data, reproduction instructions, training details, asset citations, and asset licenses.
  • The experiments report no error bars, while compute resources are documented.
  • The checklist marks consent, personally identifiable information, offensive content, human-subject instructions, participant risks, and compensation as not applicable.

A Implementation Details

All experiments use NVIDIA V100 GPUs, and Earthformer and baseline models fit on a single GPU with gradient checkpointing and no model parallelization.

  • All experiments are conducted on machines with NVIDIA V100 GPUs.
  • Earthformer and the baselines can fit in a single GPU with gradient checkpointing.
  • The models are trained without model parallelization.

A.1 Earthformer

Earthformer’s cuboid attention blocks sequentially apply configurable local attention patterns, normalization, feed-forward layers, and global-vector attention. The architecture and optimization settings are adapted slightly for SEVIR because of its resolution and dataset size.

  • Cuboid attention block: A cuboid attention block contains L cuboid attention layers with potentially different hyperparameters and attention patterns.Axial uses three consecutive patterns, while Divided Space-Time and Video Swin P × M use two.
  • Cuboid attention block: Local cuboid attention combines decomposition, attention, and merging, while global attention follows the paper’s global-attention formulation.
  • Cuboid attention block: Each block processes tensor and global-vector inputs through layer normalization, cuboid attention, and feed-forward layers with residual connections.The final pattern outputs X_d and G_d for the whole block.
  • Architecture configuration: The same architecture configurations are used for MovingMNIST, N-body MNIST, and ICAR-ENSO, with slight SEVIR adjustments for resolution and dataset size.
  • Optimization: All Earthformer variants use AdamW, train for up to 100 epochs with validation-based early stopping, and use warm-up plus cosine learning-rate decay.The total batch size is 64 through data parallelism and gradient accumulation.
  • Baseline training: Baseline models follow their official configurations, with dataset-specific tuning of learning rates, schedulers, resolutions, and related settings.

B More Ablation Analysis on Synthetic Datasets

Additional synthetic-data analyses examine cuboid-pattern selection, model depth, implementation settings, and the hierarchical auto-regressive variant.

  • Cuboid pattern search: The Axial cuboid pattern remains optimal when model depth is reduced to 2.The shallower-model results confirm the pattern-search finding.
  • Implementation details: Table 8 specifies Earthformer configurations for MovingMNIST, N-body MNIST, and ICAR-ENSO datasets.The listed components include convolution, normalization, activation, feed-forward, and patch-merging modules.
  • Implementation details: Table 9 provides the corresponding Earthformer implementation details for the SEVIR dataset.These settings include convolution, normalization, activation, feed-forward, and patch-merging components.
  • Optimization: Table 10 summarizes AdamW optimizer hyperparameters used across the four datasets.The datasets are MovingMNIST, N-body MNIST, SEVIR, and ICAR-ENSO.
  • Baselines: Table 11 lists dataset-specific modifications to officially released baseline implementations.The table also defines reverse encoder-decoder and CNN downsampler/upsampler wrapper modifications.
  • Auto-regressive variant: Figure 6 illustrates Earthformer AR as a hierarchical encoder-decoder using stacked causal cuboid-attention blocks.Its input has length 3 and its target sequence has length 2; D denotes stacked blocks and M denotes hierarchy layers.

C Non-Auto-Regressive v.s. Auto-Regressive

The paper compares non-auto-regressive and auto-regressive Earthformer decoders on SEVIR, finding a trade-off between perceptual realism and skill scores.

  • Auto-regressive decoder: Earthformer AR generates discrete visual codes one by one in raster-scan order and decodes them into pixel-space outputs.A VQ-VAE encoder supplies the input and target codes, while the decoder reconstructs the generated codes as images.
  • Auto-regressive decoder: Earthformer AR conditions each current code on the input context and previously generated codes.This defines its sequential generation procedure.
  • SEVIR comparison: Earthformer AR produces more perceptually satisfying SEVIR predictions than Earthformer but performs worse on skill scores.Its outputs look more like real VIL images and avoid blurry predictions, yet its performance is worse than Earthformer and some simple baselines.
  • Sampling ablation: Argmax sampling improves skill scores but reduces perceptual similarity to real VIL images.The comparison is between step-by-step argmax generation and random sampling from the predicted distribution.
  • Model selection: The main experiments use the non-auto-regressive decoder because Earthformer AR has worse skill scores despite more perceptually satisfying predictions.The paper notes that no well-established metric currently evaluates perceptual quality for Earth-system forecasting outputs.
  • Synthetic comparison: The synthetic N-body MNIST sequences illustrate stronger sensitivity to small initial-velocity disturbances than MovingMNIST sequences.This qualitative comparison motivates N-body MNIST as a more chaotic forecasting setting.

D Chaos in N-Body MNIST Dataset

N-body MNIST is designed to represent chaotic dynamics in which small changes to initial velocities produce substantial differences in later digit positions.

  • Chaotic dynamics: After 20 steps, slightly changing the digits’ initial velocities causes their positions to change significantly.The effect is much larger in N-body MNIST than in MovingMNIST.
  • Forecasting challenge: N-body MNIST is therefore more challenging than MovingMNIST for forecasting.Its difficulty follows from the greater sensitivity of future states to initial conditions.
  • Dataset design: The dataset uses a chaotic 3-body motion pattern for the digits.This makes N-body MNIST more relevant to Earth-system forecasting than MovingMNIST, according to the paper’s dataset description.

E Evaluation Metrics in SEVIR

The SEVIR evaluation uses Critical Success Index variants computed across precipitation thresholds and forecasting horizons, with Earthformer consistently outperforming baselines on the challenge metrics.

  • CSI definition: The Critical Success Index is computed as #Hits(τ) divided by #Hits(τ) + #Misses(τ) + #F.Alarms(τ).Hits, misses, and false alarms are counted after predictions and ground truth are rescaled and thresholded.
  • Threshold aggregation: CSI-M averages CSI across the six thresholds [219, 181, 160, 133, 74, 16].The thresholds correspond to pixel values after rescaling to the range 0-255.
  • Evaluation variants: The SEVIR Dataset Challenge uses a slightly different CSI calculation from the metric used in the main SEVIR evaluation.The challenge evaluation is specified separately from the implementation followed in the main experiments.
  • Threshold aggregation: CSI-M3 averages CSI over thresholds [133, 74, 16], whereas CSI-M6 averages it over [219, 181, 160, 133, 74, 16].These two averages correspond to the three- and six-threshold evaluation settings.
  • Temporal aggregation: CSI-τ is averaged over the forecasting horizon T before threshold-level aggregation.The horizon average is based on CSI-τ(t) across forecast times.
  • Results: Earthformer consistently outperforms baselines in almost all metrics under the SEVIR Dataset Challenge evaluation.These results are reported in Table 13.
Loading 2207.05833v2…