摘要

传统上，地球系统（如天气和气候）预报依赖于复杂的物理模型进行数值模拟，因此在计算成本和领域专业知识方面都具有较高的要求。近十年来，随着时空地球观测数据的迅猛增长，应用深度学习（DL）的数据驱动模型在各种地球系统预报任务中展现出巨大的潜力。尽管Transformer作为一种新兴的深度学习架构在其他领域取得了广泛的成功，但在这一领域的应用却相对有限。本文提出了一种用于地球系统预报的空间时间Transformer——Earthformer。Earthformer基于一种通用、灵活且高效的空间时间注意力模块，称为立方体注意力（Cuboid Attention）。其核心思想是将数据分解为立方体，并并行地应用立方体级别的自注意力机制。这些立方体进一步通过一组全局向量连接起来。我们使用MovingMNIST数据集以及一个新提出的混沌N体MNIST数据集进行了实验，以验证立方体注意力的有效性并确定Earthformer的最佳设计。在两个关于降水临近预报和厄尔尼诺/南方涛动（ENSO）预报的真实世界基准测试中的实验表明，Earthformer达到了最先进的性能。代码已公开：https://github.com/amazon-science/earth-forecasting-transformer 。