[论文解读] 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 Applications被引用 101

一句话总结

Earthformer 提出 Cuboid Attention 空间-时间变换器用于地球系统预测，通过将局部立方体自注意力与少量全局向量相结合，在对降水即时预测（SEVIR）和 ENSO SST 预测（ICAR-ENSO）方面达到最先进的成果。

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 .

研究动机与目标

• 通过利用时空观测数据的基于数据驱动的深度学习方法，推动对地球系统变率的改进预测。
• 开发适用于高维地球数据的可扩展时空 Transformer 架构。
• 设计并评估一种通用的立方体注意力机制，以实现高效的局部注意力并共享全局信息。
• 展示分层编码器-解码器在粗到细预测中的有效性。
• 研究设计选择（模式、全局向量、层次结构）以确定最优的 Earthformer 配置。

提出的方法

• 引入 Cuboid Attention，将输入张量分解为非重叠的立方体，在每个立方体内部并行应用自注意力。
• 通过在大小为 (bT, bH, bW) 的立方体内进行操作，将注意力复杂度降低到大约 O(THW · bT bH bW)。
• 引入一个全局向量集合 G，使其对所有立方体进行关注，并通过单独的全局注意力步骤进行更新。
• 附加带有初始下采样/上采样的分层编码器-解码器，以实现粗到细的预测。
• 使用非自回归解码直接从学习到的位置嵌入生成 K 步预测（在附录 C 中评估了自回归变体）。
• 探索多种立方体注意力模式（如 Axial、Divided Space-Time、Video-Swin、Spatial Local-Global），并加入全局向量以评估性能提升。

实验结果

研究问题

• RQ1 具备立方体注意力的时空 Transformer 能否在地球系统预测任务中达到或超过最先进的性能？
• RQ2 哪些立方体注意力配置（大小、策略、偏移）和架构选择（全局向量、层次结构）在现实地球数据上能实现最佳预测精度和效率？
• RQ3 全局向量和分层设计是否在合成和真实数据集上都提供一致的改进？
• RQ4 在降水即时预测和 ENSO 预测方面，Earthformer 在准确性和计算成本上与既有基线相比如何？
• RQ5 Axial 型立方体模式结合全局向量在各种任务中是否是稳健且高效的选择？

主要发现

• Axial 模式搭配全局向量在合成实验和真实世界基准测试中提供最佳整体性能。
• 全局向量提供稳定的性能提升，计算开销几乎可以忽略。
• 分层编码器-解码器在不显著增加 FLOPS 的情况下提高了预测准确性。
• Earthformer 在 SEVIR 降水即时预测和 ICAR-ENSO ENSO 预测数据集上达到最新的状态。
• 在所报告的任务上，Earthformer 使用所提设计优于基线如 UNet、ConvLSTM、PredRNN、PhyDNet、E3D-LSTM 和 Rainformer。
• 在合成测试中，使用 Axial + 全局模式的 Earthformer 展示了对长程相互作用和动力学的强大建模能力。