DeepMind 中期天气预报优势超过全球顶级气象台：1 台 TPU 1 分钟预测 10 天天气
原文标题：《DeepMind 通天了！AI 中期天气预报吊打全球顶级气象台：1 台 TPU1 分钟预测 10 天天气》
「中期天气预报」因为数据量太大，预测模型的质量一直是个难题。最近，DeepMind 和谷歌推出的全新机器学习模型，以超过 99% 的优势在天气预报模型中胜出。
众所周知，传统天气预报的可靠性，多少都有些一言难尽。
最近，DeepMind 和谷歌新研究出了一种基于机器学习的天气模拟器，可以在 60 秒内预测 10 天内的天气，而且准确率极高！
1. GraphCast 是一种基于图神经网络的自回归模型，性能优于世界上最准确的机器学习天气预报系统（中期天气预报）；
2. GraphCast 只需单台 Cloud TPU v4 设备，即可在 60 秒内生成 10 天内的天气预报（35GB 数据），分辨率高达 0.25°；
3. 通过在更大、更新、质量更高的数据上进行训练，可以进一步提升 GraphCast 预测的速度和准确性。
而在短期天气预报方面，DeepMind 曾于 2021 年 9 月在 Nature 上发文称，其生成模型能以 89% 的绝对优势吊打其他方法。
中期天气预报为什么那么难
「中期天气预报」通常是指对未来于 4 至 10 天内天气变化趋势的预报。其准确性，对于农业、建筑业、旅游业等行业的政策制定来说至关重要。
为此，国际领先的欧洲中期天气预报中心（ECMWF），每天都会提供多达四次的中期天气预报。
在中期天气预报的制作过程中，有两个关键的组成部分都需要利用大规模高性能计算（HPC）集群进行模拟：
通过分析由卫星、气象站、船舶等收集的当前和历史数据来预测天气状况，也就是「数据同化」；
通过数值天气预报（NWP）系统建立预测天气相关变量将如何随时间变化的模型。
然而，随着数据量的显著增加，NWP 模型却无法得到有效的扩展。
也就是说，虽然现在有大量的天气和气候观测档案，但我们却很难直接利用这些数据来提高预报模型的质量。
而改进 NWP 的方法，一般是由训练有素的专家手动创造更好的模型、算法和近似值，这个过程耗时耗力，成本高昂。
相比之下，基于机器学习方法可以利用更多、更高质量的可用数据来提高模型的准确性，而且计算预算通常要低得多。
GraphCast
在论文「GraphCast：中期全球天气准确预报学习」中，DeepMind 以「编码-处理-解码」的方式使用图神经网络（GNN）来创建一个自回归模型。
GraphCast 的三阶段模拟过程如下：
1. 使用从网格点到多网格的有向边的 GNN，将原始经纬度网格的输入数据映射到多网格上的学习特征中；
2. 一个深度 GNN 被用来在多网格上进行学习的信息传递，其中长距离的边允许信息在空间上有效传播；
3. 解码器将最终的多网格表示映射回经纬度网格，并执行任何必要的操作。
研究结果显示，GraphCast 的性能在 252 个变量中，有 99.2% 超过了现有最准确的机器学习天气预报模型；在 2760 个变量中，有 90% 超过了欧洲气象中心的高精度预报（ECMWF HRES Forecast）。
(a) 输入的天气状态是在高分辨率的纬度-经度-气压层网格上确定的。
(b) GraphCast 预测天气的下一个状态是纬度-经度-压力级网格。
(c) 通过迭代地将 GraphCast 应用于每个先前的预测状态，以产生一连串的状态，将天气表示为连续的提前量。
(d) GraphCast 架构的编码器组件将输入的局部区域（绿色方框）映射到多网格图的节点。
(e) 处理器组件使用所学的消息传递来更新每个多网格节点。
(f) 解码器组件将经过处理的多网格特征（紫色节点）映射到网格表示上。
ERA5 数据集
GraphCast 在 39 年（1979 年-2018 年）历史天气数据的语料库上进行了训练，即 ECMWF 的 ERA5 再分析数据集。
模型以 6 小时的时间步长，在 0.25° 经纬度分辨率下，对 5 个地表变量和 6 个大气变量进行 10 天的预测，每个变量在 37 个垂直压力层上，代表了特定地点和时间的天气状态。
如图 1a 所示，研究人员将时间指数 t 处的天气状态表示为
。
环绕地球的网格对应每个纬度、经度和压力级别的变量。表面和大气变量分别由放大视图中的黄色和蓝色框表示。
我们将
中对应于特定网格点𝑖（总共有 1,038,240 个）的变量子集称为
，并将 227 个目标变量中的每个变量𝑗称为
。
生成预测
GraphCast 将两个天气状态
作为输入，它们分别对应当前时间 t，和前一个时间 t-1，并预测下一个时间步长的天气状态（如图 1b 所示）。
为了生成 T-step 预测
，GraphCast 以自回归方式迭代上图的等式，将自己的预测作为输入，来预测后面的步长（即，预测步长 t+2，输入为
；预测步长 t + 3，输入为
。
图 1b、c 描述了这个过程。
架构
GraphCast 的核心架构在「编码-处理-解码」配置中使用 GNN，如图 1d、e、f 所示。
基于 GNN 的学习模拟器在学习流体和其他材料的复杂物理动力学方面非常有效，因为它们的表示和计算结构类似于有限元学习求解器。
GNN 的一个关键优势是，输入图的结构决定了表示的哪些部分通过学习的消息传递相互交互，从而允许在任何范围内进行任意模式的空间交互。
相比之下，卷积神经网络 (CNN) 仅限于计算局部 patch 内的交互（或者，在扩张卷积的情况下，有规律地跨越更长的范围）。
而 Transformer 虽然也可以完成任意的远程计算，但是在输入非常大的情况下，它们不能很好地扩展（要知道，GraphCast 的全局输入中有超过 100 万个网格点），因为计算中 all-to-all 的交互，会引起很复杂的二次记忆。
Transformer 的当代延伸通常会稀疏化可能的交互，以降低复杂性，这使它们实际上类似于 GNN。
通过引入 GraphCast 的内部多网格表示，研究人员利用 GNN 的能力，模拟了任意稀疏的交互方式。
它在全局范围内具有均匀的空间分辨率，并允许在少数消息传递步长内进行长距离互动。
要构造一个多网格，首先要将一个常规的二十面体（12 个节点和 20 个面）迭代 6 次，得到一个二十面体网格的层次结构，在最高分辨率下共有 40,962 个节点和 81,920 个面。
因为粗网格节点是细网格节点的子集，研究人员能够将网格层次结构中的各级边缘，叠加到最小分辨率的网格上。
这个过程产生了一个多尺度的网格集，粗边在多个尺度上弥合了长距离，细边捕捉了局部的相互作用。
图 1g 显示了每个单独的细化网格，而图 1e 显示了完整的多网格。
使用具有从网格点到多网格的定向边的 GNN，GraphCast 的编码器（图 1d）首先将原始经纬度网格的输入数据，映射为多网格上的学习特征。
然后，处理器（图 1e）使用一个 16 层的深度 GNN，在多网格上进行学习的信息传递，由于长距离的边缘，信息可以在空间上被有效传播。
然后，解码器（图 1f）使用具有定向边缘的 GNN，将最终的多网格表示映射回经纬度网格，并将该网格表示𝑌ˆ𝑡+𝑘与输入状态𝑋ˆ𝑡+𝑘相结合，形成输出预测，𝑋ˆ𝑡+𝑘+1 = 𝑋ˆ𝑡+𝑘 + 𝑌ˆ𝑡 +𝑘。
训练过程
GraphCast 被训练成在 12 步预测（3 天）中对 ERA5 目标进行目标函数最小化，使用的是梯度下降法。
目标函数如下 --
研究人员使用批处理并行技术，在 32 台 Cloud TPU v4 设备上花了大约 3 周时间对 GraphCast 进行了训练。
为了减少内存占用，研究人员还使用了复杂的梯度检查点策略和低精度的数值。
结果
结果显示，GraphCast 在 0.25° 分辨率的 10 天预报中，全面超越了 HRES 天气预报技术。
如图 4 所示，GraphCast（蓝线）在 10 个主要地表和大气变量上，都明显优于 HRES（黑线）。
此外，研究人员通过区域分析表明，这些结果在整个地球上是一致的。
根据评估的结果，GraphCast 在 2760 个变量、等级和前置时间（4 个地表变量，加上 5 个大气变量 ×13 个等级，历时 10 天，每天 4 个步长）中的 90.0% 表现优于 HRES。
研究人员表示，HRES 在高层大气级别上的表现往往比 GraphCast 好，特别是压力级别 50hPa，这并不奇怪，因为应用于 50hPa 或以下压力级别的总训练损失权重只占所有变量和级别总损失权重的 0.66%。
当排除 50hPa 水平时，GraphCast 在 2240 个目标中优于 HRES 的百分比为 96.6%；当排除 50 和 100hPa 水平时，1720 个目标中的百分比为 99.2%。
第 1 行显示 ERA5，第 2 行显示 HRES，第 3 行显示 GraphCast，第 4 行和第 5 行分别是 HRES 和 HRES-fc0、GraphCast 和 ERA5 之间的误差绝对值图。底部的图显示了 HRES 和 GraphCast 的 RMSE 水平。
自回归训练对预测的影响
当用较少的自回归步长训练时，模型在较短的前置时间内表现较好，而在较长的前置时间内表现较差。
随着自回归步数的增加，在较短的前置时间内性能变差，但在较长的前置时间内性能变好。
GraphCast 与顶级 ML 预测模型的性能比较
目前，基于 ViT 的 Pangu-Weather 代表了基于 ML 的天气预报的最新水平，其计算模式与 GNN 相似。
GraphCast 与 Pangu-Weather 的对比结果如图 8 所示。第 1 行和第 3 行显示 GraphCast（蓝线）、Pangu-Weather（红线）、HRES 对 HRES-fc0 的评价（黑线）和 HRES 对 ERA5 的评价的绝对 RMSE；第 2 行和第 4 行显示各模型之间相对于 Pangu-Weather 的归一化 RMSE 差分。
总结一下
GraphCast 模型在 10 天的预报中，在 6 小时步长和 0.25° 经纬度分辨率下，超过了目前最精确的确定性系统 --ECMWF 的 HRES。
针对 2760 个变量、压力等级和前置时间的组合进行评估的结果显示，GraphCast 模型在 90.0% 的指标上比 HRES 的 RMSE 低。
当排除了 100hPa 及以上的高层大气场时，GraphCast 在 1760 个目标中的 99.2% 表现优于 HRES。
此外，在 252 个目标中，GraphCast 有 99.2% 超过了之前最好的 ML 基线 --Pangu-Weather。
GraphCast 的一个关键创新是其新颖的「多网格」表征方法，这使得它能够捕捉到比传统的 NWP 方法更长的空间互动，从而支持更粗的原始时间步长。
这就是为什么 GraphCast 可以在一个 Cloud TPU v4 设备上以 6 小时为单位在 60 秒内生成准确的 10 天天气预报的部分原因。
参考资料：
https://arxiv.org/abs/2212.12794
本文来自微信公众号：新智元 （ID：AI_era），编辑：好困 Aeneas
