速度升45000倍，英伟达用傅里叶模型实现前所未有天气预报准确率

　　机器之心报道
　　机器之心编辑部
　　英伟达在天气预报上放出了「大招」。
　　现代数值天气预报（numerical weather prediction, NWP）可以追溯到 1920 年代。当今，数值天气预报已经无处不在，并且为交通、物流、农业和能源生产等关键部门的经济规划做出了贡献。准确的天气预报通过提前通知极端事件挽救了无数人的生命。
　　过去几十年，天气预报的质量一直在稳步提高。最早的单点动态建模数值天气预报是由英国科学家 Lewis Fry Richardson 在 1922 年使用计算尺和对数表计算得出的，并用六周时间计算出 6 小时（6-hour）的大气预报。到了 1950 年代，早期的电子计算机极大地提高了预测速度，使得运行预测的计算速度足够快，从而有助于未来天气预测。除了更好的计算能力之外，通过更深入地了解小尺度过程的物理特性和更高质量的大气观测，来更好地参数化小尺度过程，从而改进天气预报效果。
　　与 SOTA NWP 模型相比，数据驱动的深度学习模型在成本上低了数个量级，因此现在研究人员越来越有兴趣开发这类方法来预测天气。许多研究都试图通过建立数据驱动的模型来预测大气的规模环流，并且要么在气候模型输出、大气环流模型（GCM）、再分析产品，或者结合气候模型输出和再分析产品进行训练。数据驱动的模型通过克服 NWP 模型中存在的模型偏差，以及通过以较低的计算成本生成用于概率预测和数据同化的大型集合，因而在改进天气预报方面具有巨大潜力。
　　大多数据驱动的天气模型使用低分辨率数据来训练，通常是 Rasp and Thuerey [2021b]中的 5.625  分辨率或者 Weyn et al. [2020] 中的 2  分辨率。这些先验尝试在预测一些粗粒度、低分辨率大气变量方面取得了很好的效果。然而，粗化过程会导致关键的细粒度物理信息的丢失。要使数据驱动模型真正具有影响力，它们生成的预报的分辨率必须与当前 SOTA 数值天气模型相同或更高，分辨率大约在 0.1 。
　　在近日的一篇论文中，英伟达、劳伦斯伯克利国家实验室、密歇根大学安娜堡分校、莱斯大学等机构的研究者开发了一种基于傅里叶的神经网络预测模型 FourCastNet，它能以 0.25  的分辨率生成关键天气变量的全球数据驱动预测，相当于赤道附近大约 30 30 km 的空间分辨率和 720 1440 像素的全球网格大小。这使得我们首次能够与欧洲中期天气预报中心（ECMWF）的高分辨率综合预测系统（IFS）模型进行直接比较。
　　论文地址：https://arxiv.org/pdf/2202.11214.pdf
　　FourCastNet 在节点小时（node-hour）基础上比传统 NWP 模型快约 45,000 倍。FourCastNet 这种数量级的加速以及在高分辨率下前所未有的准确性，使得它能够以很低的成本生成超大规模集合预测。FourCastNet 极大地改善了概率天气预报的效果，使用它可以在几秒钟内生成对飓风、大气层河流和极端降水等事件的大规模集合预报，从而可以实现更及时、更明智的灾难响应。
　　此外，FourCastNet 对近地表风速的可靠、快速和低廉预测可以改善陆海风电场的风能资源规划。训练 FourCastNet 所需的能量大约等于使用 IFS 模型生成 10 天预测所需的能量（50 个成员）。然而，一旦经过训练，FourCastNet 生成预测所需的能量比 IFS 模型少 12,000 倍。研究者希望 FourCastNet 只训练一次，并且后续微调的能耗可以忽略不计。
　　在实现技术上，FourCastNet 使用基于傅里叶变换的token混合方法 [Guibas et al., 2022] 和ViT骨干 [Dosovitskiy et al., 2021]。这一方法基于最近的的傅里叶神经算子，该算子以分辨率不变的方式学习，并在建模流体动力学等具有挑战性的偏微分方程中取得了成功。此外，他们选择 ViT 骨干的原因是它能够很好地建模长程依赖。ViT 和基于傅里叶的 token 方法混合生成了 SOTA 高分辨率模型，它可以解析细粒度的特征，并能够很好地随分辨率和数据集大小扩展。研究者表示，这一方法能够以真正前所未有的高分辨率训练高保真数据驱动的模型。
　　训练模型
　　欧洲中期天气预报中心（ECMWF）提供了一个公开可用的综合数据集 ERA5，该研究使用 ERA5 来训练 FourCastNet。他们专注于两个大气变量，即（1）距离地球表面 10m 处的风速和（2）6 小时总降水量，除此以外，该研究还预测了其他几个变量，包括几个不同垂直高度的位势高度、温度、风速和相对湿度，一些近地表变量，如地面气压和平均海平面气压以等。
　　FourCastNet模型
　　为了生成高分辨率预测，研究采用自适应傅里叶神经算子 (Adaptive Fourier Neural Operator，AFNO) 模型。这种神经网络架构专为高分辨率输入而设计，并将深度学习（DL）最新关键进展综合到一个模型中。也就是说，FourCastNet 结合了 傅里叶神经算子 (FNO) ，该方法已被证明在建模具有挑战性的 PDE 系统方面表现良好，并且还具有强大的 ViT 主干。
　　目前，研究者已经提出了几种降低计算复杂度的 ViT 变体，然而，AFNO 模型的独特之处在于它将混合操作构建为连续的全局卷积，在傅里叶域中通过 FFT 有效实现，这允许灵活且可扩展地建模跨空间和通道维度的依赖关系。通过这样的设计，空间混合复杂度降低到 O(N log N)，其中 N 是图像 patch 或 token 的数量。这种缩放使 AFNO 模型非常适合 0.25  分辨率的高分辨率数据，以及进行未来更高分辨率的研究。
　　模型的计算流程：首先将 720   1440 lat-lon 网格上的输入变量投影到 patch 的 2D 网格 (h   w)（patch大小为 p   p，例如 p = 8），每个 patch 表示为一个 d 维 token。然后，将 patch 序列连同位置编码一起馈送到一系列 AFNO 层。
　　训练
　　该研究主要关注预测地表风速和降水，但复杂的大气系统包含多个变量之间的强非线性相互作用，例如温度、地面气压、湿度、从地球表面到平流层的水分含量等。为了模拟这些相互作用，该研究选择了一些变量（表 1）来表示大气的瞬时状态。
　　该研究将 ERA5 数据集分为三组，即训练、验证和测试数据集。训练数据集由 1979 年至 2015 年的数据组成，验证数据集包含 2016 年和 2017 年的数据，测试数据集包含 2018 年及以后的数据。
　　该研究使用张量 X(k t) 表示变量，其中k为时间指数， t 为训练数据集中连续快照之间的时间间隔，研究将 ERA5 数据集视为真值，用 X_true(k t) 表示真值变量， t 固定在 6 小时。训练包括预训练和微调两个阶段：  预训练阶段，采用训练数据集以有监督的方式训练 AFNO 模型，使其学习从 X(k) 到 X(k + 1) 的映射；  微调阶段，该研究从之前训练好的模型开始对模型进行优化，即模型首先从输入 X(k) 生成输出 X(k + 1)，然后模型使用输出 X(k + 1) 作为输入，并生成输出 X(k + 2)。通过比较 X(k + 1) 和 X(k + 2) 与训练数据中各自的真值来计算训练损失，并使用这两个训练损失的总和来优化模型。
　　整个训练过程是在 64 个 Nvidia A100 GPU 的集群上完成，端到端训练大约需要 16 小时。
　　降水模型
　　ERA5 re-analysis 数据集中的总降水量（total precipitation，TP）是一个变量，表示通过降雨和降雪落到地球表面的累积液态和冰冻水。该研究将总降水量作为诊断变量，并用 p(kΔt) 表示。用于训练主干模型的数据集（20 个变量）不包括总降水量这个变量。相反，研究者训练一个单独的 AFNO 模型 ，使用主干模型输出来诊断 T P，如图 2(c) 所示。这种方法将降水建模的困难与预测大气状态的一般任务脱离。此外，一旦经过训练，诊断 TP 模型可能会与其他预测模型（传统的 NWP 或数据驱动的预测）结合使用。
　　用于从主干输出诊断降水的模型具有相同的基本 AFNO 架构，此外还具有额外的 2D 卷积层和 ReLU 激活，用于强制执行非负降水输出。由于主干模型以 6 小时为增量进行预测，因此该研究训练诊断降水模型来预测 6 小时累积的总降水量。
　　实验结果
　　下图1定性地展示了 FourCastNet 模型以 0.25 -lat-long 的分辨率预测全球地表风速的能力。其中，使用风速的纬向和经向分量计算地表风速的大小，即
　　。
　　下图1显示了使用FourCastNet生成的提前96小时的说明性全球近地表风速预测。研究者重点介绍了预测的已解析和准确跟踪的关键高分辨率细节，包括超级台风山竹（Mangkhut，2018年）和向美国东海岸进发的三个命名气旋（Florence、Issac和Helene）。
　　其中，图 1(a) 显示了模型初始化时的风速。图 1(b) 显示了提前 96 小时的模型预测（上）和当时相应的真实风速（下）。可以看到，FourCastNet 模型能够提前 96 个小时预测风速，并且具有出色的保真度和正确的精细尺度（fine-scale）特征。
　　值得注意的是，图1说明了一个名为山竹的超级台风的形成和轨迹的预测情况，该台风在大约 10  N、210  W 的初始位置开始形成。该模型以惊人的保真度定性地预测了该台风的增强及其在 4 天内的轨迹。这表明 FourCastNet 模型能够很好地预测台风等天气现象的形成和移动轨迹。
　　该研究使用来自样本外测试数据集的初始条件初始化 FourCastNet 模型。基于此初始条件，模型可以在推理模式下自由运行 16 个时间步长，如下图 2(d) 所示。
　　相比于台风，降水存在间歇性和随机性，因此预测降水被认为是一项极其困难的任务。下图 3 展示了 FourCastNet 模型预测全球总降水量的能力。图 3(a) 为初始化的降水条件，图 3(b) 是提前 36 个小时的模型预测和实际情况的对比。该图以美国西海岸局部区域的降水为例，突出了 FourCastNet 模型以极高的精度解析和预测局部降水的能力。研究者观察到 FourCastNet 在捕获短期高分辨率降水特征方面性能出色，这对预测极端天气事件具有重大意义。值得注意的是，这是第一次成功利用深度学习模型进行大规模降水预测。
　　接下来，研究者进一步探讨了开发深度学习模型来预测飓风的潜在效用。飓风作为一种具有巨大破坏潜力的极端天气事件。构建快速可用、计算成本低廉的大气预警模型对于减轻生命和财产损失损失非常重要，同时这类模型也需要更严格的要求以避免错误预测。下图 4 展示了 FourCastNet 在预测迅速多变的飓风方面的强大能力。
　　除了台风和降水，FourCastNet 模型还能够预测大气层河流（ atmospheric river）的形成和演变。下图5显示了一个于 2018 年 4 月最终登陆北加州海岸线的大气层河流的预测情况。
　　此外，FourCastNet 在量化技巧方面也有很大的优势，比 IFS 更佳。对于 ACC 和 RMSE 两项指标，FourCastNet 在更短的时间（约 48 小时或更短）内获得了比 IFS 更好的性能。并且 FourCastNet 仅使用了 IFS 模型完整变量集的一部分，并且计算成本也低了很多。
　　集成是数值天气预测的重要组成部分，FourCastNet 使用大型集成带来的性能提升如下图所示。
　　感兴趣的读者可以阅读论文原文，了解更多研究细节。