Earthformer: Exploring Space-Time Transformers for Earth System Forecasting

TL;DR

Earthformer introduces a novel space-time Transformer architecture with cuboid attention for improved Earth system forecasting, demonstrating state-of-the-art results on real-world benchmarks like precipitation nowcasting and ENSO prediction.

Contribution

The paper presents Earthformer, a flexible and efficient space-time Transformer with cuboid attention, specifically designed for Earth system forecasting tasks.

Findings

Earthformer outperforms existing models on precipitation forecasting.

It achieves state-of-the-art results on ENSO forecasting.

The cuboid attention mechanism effectively captures spatiotemporal dependencies.

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 .