2 个月前

Earthformer:探索用于地球系统预报的空间-时间变换器

Zhihan Gao; Xingjian Shi; Hao Wang; Yi Zhu; Yuyang Wang; Mu Li; Dit-Yan Yeung
Earthformer:探索用于地球系统预报的空间-时间变换器
摘要

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