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 .
翻译:传统上,地球系统(如天气和气候)预测依赖于基于复杂物理模型的数值模拟,因此既在计算上昂贵,又需要大量领域专业知识。随着过去十年时空地球观测数据的爆炸式增长,应用深度学习的数据驱动模型正展现出解决各类地球系统预测任务的巨大潜力。尽管Transformer作为一种新兴的深度学习架构在其他领域取得了广泛成功,但在该领域的应用仍然有限。本文提出Earthformer,一种用于地球系统预测的时空Transformer。Earthformer基于一个通用、灵活且高效的时空注意力模块,称为立方体注意力(Cuboid Attention)。其核心思想是将数据分解为立方体,并并行地应用立方体级别的自注意力机制。这些立方体进一步通过一组全局向量进行连接。我们在MovingMNIST数据集和一个新提出的混沌N体MNIST数据集上开展实验,以验证立方体注意力的有效性并确定Earthformer的最佳设计。针对降水临近预报和厄尔尼诺-南方涛动(ENSO)预测这两个真实世界基准的实验表明,Earthformer达到了最先进的性能。代码地址:https://github.com/amazon-science/earth-forecasting-transformer。