NeuralGCM：AI天气预报进入下半场？

最新推荐文章于 2024-05-29 14:23:25 发布

AI蜗牛车

最新推荐文章于 2024-05-29 14:23:25 发布

阅读量851

点赞数 2

文章标签：人工智能

原文链接：https://mp.weixin.qq.com/s?__biz=MzA4ODUxNjUzMQ==&mid=2247503036&idx=1&sn=4c219ee9d74885f9be1c175cd4f7a222&chksm=911c257265f92c21fdd598dc0dae912a31b0443fcceff7a5d09e806d68788c5ee4e5b0e11e4d&scene=126&sessionid=0

版权

无意间看到Google又在ArXiv上发布了一篇基于AI的天气预报论文NeuralGCM（我为什么要说“又”？），着实被结果震惊了：

中短期预报技巧与GraphCast和EC相当，同时大幅缓解AI模型的物理一致性问题。
在预报时效最多5天的数据集上训练的模型可以稳定地模拟年代际气候变化。
包含微分求解器和机器学习参数化方案的端到端模型，结构与数值模式高度一致。

尽管仍是以ERA5为基准的AI预报模型，但以上每一点似乎都是重大突破。一起看看他们是如何做到的。

1. NeuralGCM是什么：聊聊NeuralGCM和之前的AI气象大模型的区别。

2. NeuralGCM的设计：非大气科学背景的读者敬请只看加粗部分或略过。

3. NeuralGCM的评估：展示部分NeuralGCM的评估结果。基于AI的气候模拟来了！

4. 讨论：喜闻乐见的吹水环节。

NeuralGCM是什么

当看到AI模型被冠以GCM（General Circulation Models，大气环流模式）的字眼时，我意识到这可能是一篇不一样的论文。这里试图为不同背景的读者解释清楚NeuralGCM的特别之处。

数值模式主要由两个部分组成：动力框架和物理过程参数化。前者用方程组描述大气变化的过程和不同物理量间的关系；后者用来计算一些前者无法描述的物理过程。例如，由于数值模式在求解动力方程组时需要将其离散化（在空间上表示为网格，ERA5和ECMWF预报就都是格点数据），必然遗漏了小于网格分辨率的物理过程，需要参数化方案来简单模拟。

近期发表的盘古、GraphCast等AI模型实际上省略了这样的分解，将整个预报流程简化为图像预测任务，即输入上一时刻的格点大气状态，预报下一时刻的状态（见之前的公众号文章）。然而在此之前，已经有研究试验了更深入的数值模式-机器学习融合策略。一个典型的做法是用机器学习模型来代替部分参数化方案。其思路可以概括为：一些高分辨率的数值模式往往能够“分辨”较小尺度的物理过程，而受限于计算效率，实际业务中又常使用粗分辨率的数值模式。利用前者模拟的输入输出来训练一个基于机器学习的参数化方案嵌入到后者中，即可得到一个在某方面效果接近前者而速度不慢于甚至快于后者的混合模型。

例如：

机器学习对流参数化：doi.org/10.1073/pnas.1810286115
机器学习边界层参数化：doi.org/10.1029/2022MS003034

这样做的显著问题是，机器学习模型的训练是个独立的过程，其泛化性将在实际预报过程中受到很大挑战——训练集中反映的统计关系和实际应用场景可能有较大的区别。如果将这种范式称为“离线训练”，那么NeuralGCM的训练就是“在线”（online-training）的——代替参数化方案的模型与动力框架求解器共同参与训练过程，组成一个整体直接拟合目标：这个端到端模型的输入和输出就是ERA5的大气场，与盘古、GraphCast等一致。

不得不提其技术实现：作者利用Google近两年发展的机器学习框架JAX，构建了基本方程组的自动微分模块和代表物理过程的神经网络优化模块。由于导数计算本就是JAX的拿手好戏，二者得以融为一体。

简而言之，NeuralGCM是一个具有与数值模式类似的时间积分方案，由动力框架和神经网络组成，具有物理一致性且能长期稳定迭代预报的混合模型。

NeuralGCM的设计

NeuralGCM与数值模式的结构已非常接近，可以拆解为下文几个部分：

# 输入输出

模型输入分为三部分：外强迫（如海温，对气候模拟尤为重要），随机噪音（用于生成集合预报）和ERA5的大气场。输出则是预报的大气场。

着重解释下图中的“Learned encoder”和“Learned decoder”。由于ERA5的数据是在等压面上提供的，而模型为了做谱展开，采用的是σ垂直坐标系（与大部分数值模式类似。简而言之，模型内部计算时垂直分层与ERA5不一致）。因此需要将变量垂直线性插值到σ坐标系上（encoder），预报完成后再插值回来（decoder）。但这种插值有一定误差，会导致虚假重力波等问题，因此这里的encoder和decoder不仅是线性插值，还包含一个可学习的神经网络，用于订正插值结果。这里的神经网络是个“单柱”模型，即每个格点独立订正垂直方向的插值。

# 动力框架

动力框架包含的基本方程组如下图所示：

左边分别为散度ζ、涡度δ、温度T、比湿q和表面气压Ps对时间的导数（水平风速可以从ζ、δ推导）。NeuralGCM的动力部分目标即是求解这些方程左侧项，来预报一个时间步长后这些量的变化，即前文结构图中的”Dynamic tendencies“。

此外，诊断方程如下：

方程左侧为诊断量，通俗理解，即这些量通过其他量来求得，而不是通过求其对时间的导数来更新。其中涉及垂直速度，气压，位势高度，虚温等变量。

特别一提，对于降水，在NeuralGCM中以P-E形式出现：

方程左侧为格点上的降水（Precipitation）减去蒸发（Evaporation），其值等于该格点上空水汽变化的垂直积分，而NeuralGCM预报了水汽，因此可以诊断P-E的值（论文Fig.H21与ERA5进行了对比）。

其方程简化、半隐式时间积分方案、波数截断等细节请参照原文附录。这里只再提两点：

训练了T62L32，T125L32，T253L32三种分辨率的模型，分别对应2.8°、1.4°和0.7°的输出分辨率。
三种分辨率的步长分别为12、6、3.75分钟。

# 物理过程参数化

参数化方案同样为单柱模型，不同格点间共享参数，由于其结构为全连接神经网络，其整体可以理解为1x1的卷积操作。

模型输入包括上文基本方程组左侧的预报量及其水平梯度，还包括结构图中的Forcings，如地形、时间、经纬度、太阳辐射、海冰密集度和海表温度等（根据下垫面类别和模式分辨率有所不同，附录C.1、C4.2有介绍）。对于集合预报，输入还包括一组高斯噪声。

模型输出与动力部分类似，为预报量的变化，即前文结构图中的“Physics tendencies"。换句话说，基本方程和神经网络预报的变化项共同作用于预报量的演变。

此外还有两点细节：

Physics tendencies每隔30～60分钟更新一次，大于dynamic tendencies的步长；
作者将physics tendencies输出缩放为观测抽样中tendencies的0.01倍。即手动控制使Physics tendencies对整体预报的影响幅度大约在1%左右。

如此，我们简要提炼了NeuralGCM的架构。

# 训练过程

这里主要谈NeuralGCM的优化设计，仍然将关键部分加粗。

类似GraphCast，NeuralGCM训练开始时只预报6小时并计算loss；逐步增加到3~5天为止。当然，由于这里类似数值模式时间积分方案的设计，即使是6小时，也包含数十至上百个时间步的迭代。
在等压面和等σ面上同时计算loss。在等σ面上计算loss可能是为了拆分任务，排除前文encoder和decoder的干扰，明确encoder和decoder的目标为垂直插值。
对不同变量赋予不同权重，以使训练过程不至于被个别变量主导。
由于预报时间越长，误差本身必然越大，因此对更长的预报时间赋予更小的权重，使得不同预报时间的loss接近。

NeuralGCM的损失函数分为三类：

Accuracy loss：使用MSE使预测更准确。随着预报时效的延长，大气的可预报性降低，如果模型预报出了某个小尺度波动，但位置出现了偏差，则观测和预报的两个位置的MSE都会增大（称为“双重惩罚”问题，这也是很多现有AI模型模糊化的根源之一），然而实际这样的预报是有意义的（比如预报海淀下雨，结果朝阳下了，总比压根没报雨好些）。因此，作者在计算MSE loss前，滤去了认为在不同预报时效上已没有可预报性的小尺度变化（即过滤了高波数分量），以避免双重惩罚。
Sharpness loss：另一方面，为了解决模糊问题，作者直接在loss中计算各高度层各变量不同波长分量的能量与观测的差异，即强迫模型预报不同尺度的变化。这一项针对性地解决之前ECMWF一篇论文《On some limitations of data-driven weather forecasting models》中提及的问题。
Bias loss：该项在计算MSE前，先对整个batch和各预报时长做平均，因此不针对个例，而是调整不同高度层、不同变量、不同波长的振幅。

此外，集合预报模型采用了CRPS（连续分级概率评分）作为损失项。

在训练的最后阶段，作者还对decoder部分进行了微调，以解决谱空间与格点空间相互转换带来的高频噪声问题（附录G.5）。

NeuralGCM的评估

之前解读AI模型的文章几乎不涉及模型评估结果，但NeuralGCM有些特殊，值得一起看看。

# 中短期预报

NeuralGCM在短期预报上与GraphCast相当，而中期集合预报甚至超越ECMWF-ENS。此外还有如下优势：

有效缓解了AI模型中期预报的模糊问题。无论是确定性预报和集合预报成员的可视化，还是功率谱分析的结果都表明，NeuralGCM在更长的预报时间上始终保持清晰，远优于GraphCast。这也有利于极端事件的预报。
具有更强的物理一致性和可解释性。从降水预报的空间分布，地转风平衡和模型的泛化性角度看，NeuralGCM的预报相比GraphCast更加稳定。这显然也得益于动力框架提供的保障。

上一期推送的文末提到要继续讨论AI模型物理一致性的问题，因此NeuralGCM在中短期预报上评估的细节将在下一期一起讨论。这里我们再着重介绍下NeuralGCM的气候模拟能力。

# 气候模拟

尽管NeuralGCM只在3~5天预报时长上训练，但可以进行稳定的气候模拟。作者采用1.4°版本，以2019年37个不同时间数据作为初始场，其中35个提供了稳定了2年模拟（不太清楚作者对不稳定的定义，以及什么原因导致了不稳定）。

从下图a可以看出，NeuralGCM模拟的全球850hPa平均温度不仅准确地模拟了季节循环，而且趋势与观测一致，远优于气候平均（1990-2019年平均）。NeuralGCM在可降水量、动能等指标上也有类似表现，对哈得莱环流、季风环流等气候现象的空间分布和季节循环也与观测一致。

不仅如此，NeuralGCM还可以用于气候变化预估。37次模拟中的22个在40年模拟上保持稳定。这里的模拟形式类似AMIP。对此稍作解释：

GraphCast、盘古等模型可以称为“大气模式”，因为只考虑了大气本身的变化规律。而从长期（月、季、年际甚至年代际）角度，大气的变化取决于其他圈层，特别是下垫面的影响。例如大家可能耳熟能详的厄尔尼诺，就是赤道中东太平洋海表面温度的异常，是大气季节预测的重要预测因子。GraphCast、盘古的输入是大气的“初始条件”；而要进行长期模拟，关键则是把握“边界条件”，即其他圈层（大气的边界）对大气的影响。当然，这里长期模拟的目的不再是预报某次降水、降温、台风等，而是从统计上预估气候特征的变化，如全球平均气温的长期趋势。

图片来源：https://www.energy.gov/science/doe-explainsearth-system-and-climate-models

正因如此，气候变化预估不仅需要大气模式，还需要模拟海洋、冰雪、陆面、人类活动等的变化。而AMIP（Atmospheric Model Intercomparison Project，大气模式比较计划），则仅旨在测试大气模式本身的性能。其方式是以过去数十年观测到的海表面温度、海冰等作为“外强迫”，观察大气模式是否能在此基础上模拟过去的气候变化。

前文展示的NeuralGCM结构中，“Forcings"即为这类外强迫。2.8°版本的NeuralGCM模拟的1981年以来全球平均气温变化展示了与ERA5一致的增温趋势，其年际震荡也与ERA5相近（见上上张图b）。显然，NeuralGCM有效学习了海温等外强迫对气候变化的影响。不仅如此，图c展示了不同高度层的40年温度变化线性趋势。NeuralGCM的集合平均与ERA5十分接近，甚至优于AMIP的大气数值模式（当然，这样的比较不太公平，毕竟NeuralGCM是在ERA5上训练的）。

对于NeuralGCM的解读就到这里。原文共67页，本文必然遗漏了不少细节，但希望可以为有兴趣研读的朋友当作索引。

讨论

通篇看来，NeuralGCM最大的不同是以动力框架提供的物理关系为核心，进而从源头上解决了AI模型的最大问题：物理一致性。其带来的好处不言而喻：

不必如上篇推送中提及的模型一样在误差累积问题上做太多文章。相反，可以像数值模式一样根据空间分辨率设置相应的时间积分步长。即使是神经网络部分，也可以说拟合的是物理过程而不是统计关系。
保证输出的大气环流具备物理可解释性，与观测到的天气、气候现象始终保持一致。在物理方程的约束下，仅用3~5天预报时间的数据即可训练出在几十年模拟上保持稳定的模型。

当然，由于其监督学习的本质没有改变，其中仍然存在一些问题：

仍然以ERA5为基准训练模型。类似数值模式的动力框架置入NeuralGCM后，要超越数值模式的预报，某种程度上其成功的核心仍然是替代参数化方案的神经网络部分。同时，ERA5也依赖数值模式，则NeuralGCM的学习目标其实限于ERA5对观测的同化。换句话说，NeuralGCM的本质仍然是利用神经网络拟合观测，且上限仍受同化效果的约束。
仍存在大量“魔法数字”的设置。如不同变量损失的平衡，参数化方案输出的放缩，不同损失项的权重等。这些参数的影响很难评估。当然，对于次网格尺度过程而言，这些经验性超参数即使在数值模式中也不可避免。

进一步，我们可以从不同角度理解NeuralGCM的设计：

从AI角度，NeuralGCM将代表物理过程的基本方程组嵌入神经网络，训练出受物理约束的AI预报模型。
从数值模式角度，则是用神经网络代替数值模式中的参数化方案。
甚至，我们可以将NeuralGCM中的神经网络视作“订正模型”，只是将订正直接嵌入了数值模式的时间积分过程中，针对每一步的结果订正，而不是订正最终的输出。

更折衷的理解是，NeuralGCM将“统计拟合”的范畴从天气变化整体压缩到了更小的范围内，从而既以动力框架为主体，反映物理过程；又以神经网络替代繁杂的参数化方案。这充分利用了二者的不同特性：对可以清晰表达的过程用求解方程的形式解决，而对诸多无法明确表达的过程，则笼而统之用神经网络模拟，看起来是个非常合理的混合方案。类似我们定框架、做实验、写论文，最后请ChatGPT补充细节、润色文笔。

NeuralGCM原文中有一句话：

ML atmospheric models also require considerably less code, for example GraphCast has 5417 lines vs 376578 lines for NOAA’s FV3 atmospheric model.

乍一看这种比较方式有点无厘头，甚至Fortran也躺枪了（数值模式基本是Fortran写的），但通读全论文后，发现代码行数某种程度上代表着迭代难度的降低。一直以来，数值模式技术是大气科学“皇冠上的明珠”。也正因为其复杂性，数值模式的进步可能正变得越发艰难。即便只考虑工程实现，几十万行代码也让人望而却步。如果机器学习与数值模式的深度融合之路能够走通，更低的成本和更快的速度，将会给数值试验带来极大便利。

数值模式“静悄悄的革命”进程一直是以年代为单位考量的。而自从AI技术引入大气科学后，突破性（至少从形式上）的研究层出不穷，目不暇接，不用说跟上前沿发表论文，就连写公众号都有些力不从心……每项研究都是如此“不完美”，以至于可以花大量篇幅指出其中的不足，但他们的进化速度又是如此之快，以至于我能想到的每个突破口都有团队在攻克。不得不说这种“小步快跑”的崭新节奏成为大气科学研究领域一道新的风景。

正如我们在第一篇推送中就提到的：要让AI与学科研究相辅相成；用神经网络压缩的不仅是数据的规律，还包括人类对物质世界的理解和刻画。Google与ECMWF的这次合作，为AI天气预报的“下半场”打开了大门。