1 分钟内完成 15 天预报,英伟达/UC 伯克利等提出机器学习天气预报系统 FCN3,支持单卡极速推理 原创

发布于 2025-7-31 13:40
浏览
0收藏

自 20 世纪数值天气预报概念提出并逐步发展以来,数值天气预报(NWP)极大地革新了人类对大气现象的理解与预测方式。早期,受限于计算机性能,数值天气预报进展缓慢,直到 50 年代计算机技术取得突破,数值天气预报试验才初获成功,70 年代随着超级计算机能力增强,其开始广泛应用于业务。如今,数学建模的持续优化、计算能力的飞跃提升以及资料同化技术的不断改进,让 NWP 在天气预测、灾害防御、能源管理和气候研究等诸多方面都成为了不可或缺的关键工具。

但传统 NWP 模型一直面临着严峻挑战,因其基于流体力学和热力学方程组的数值求解,计算量极为庞大。在高分辨率预报和大规模集合预报时,传统 NWP 模型的计算成本急剧增加,难以满足快速、精准且大规模的概率集合预报需求,这严重限制了其在实际应用中的进一步拓展。

为解决这类问题,英伟达、美国劳伦斯伯克利国家实验室、加州大学伯克利分校、美国加州理工学院的联合研究团队,推出了 FourCastNet 3(FCN3),这是一个将球面信号处理与隐马尔可夫集合框架相结合的概率机器学习天气预报系统。

该模型的预报技巧超越传统 NWP 黄金标准 IFS-ENS,在中期预报上与领先的概率 ML 模型 GenCast  相当,且时间分辨率翻倍,基于单张 NVIDIA H100 GPU 可以在 60 秒内完成 15 天天气预报,速度为 GenCast 的 8 倍、 IFS-ENS 的 60 倍,支持单卡极速推理,4 分钟内即可生成 60 天、 0.25° 分辨率、 6 小时间隔的全球预报。

相关研究成果以「FourCastNet 3: A geometric approach to probabilistic machine-learning weather forecasting at scale」为题,发表于 arXiv 。

1 分钟内完成 15 天预报,英伟达/UC 伯克利等提出机器学习天气预报系统 FCN3,支持单卡极速推理-AI.x社区

论文地址:

https://arxiv.org/pdf/2507.12144

关注公众号,后台回复「FCN3」获取完整 PDF

更多 AI 前沿论文: https://go.hyper.ai/owxf6

ERA5 数据集,FourCastNet 3 训练的核心支撑

FourCastNet 3(FCN3)的核心训练数据来源于 ERA5 数据集,这是一份由欧洲中期天气预报中心(ECMWF)制作的、持续数十年的逐小时地球大气状态再分析数据。 ERA5 通过四维变分同化系统,将 1979 年以来的各类观测数据(包括探空仪、卫星、飞机、地面站、浮标等)与 IFS 模式状态相融合,生成了空间分辨率为 0.25°×0.25°(以 721×1440 经纬度网格表示)的全球大气场,数据总量约 39.5 TB 。

ERA5 的显著优势在于,其再分析过程始终基于同一 IFS 循环(如 CY41R2 及后续固定配置),这使得它在时间轴上能保持动力学的一致性,有效避免了业务分析因模式升级而产生的气候漂移,为机器学习模型提供了可重复、可追溯的「真实大气」基准。同时,它综合多种数据源并充分考虑各自的不确定性估计,能够对地球大气历史进行一致刻画,成为机器学习模型逼近行星尺度大气动力学的理想目标。

为训练 FCN3,研究人员从 ERA5 中选取了 72 个变量,涵盖 7 个表层变量和 5 个大气变量在 13 个等压面上的值。尽管最终训练的是 6 小时间隔的模型,但仍采用了 1980–2018 年的逐小时采样数据,以此最大化数据集规模并提升模型的泛化能力。

该数据集被明确划分为三个部分:1980–2016 年为训练集,2017 年为测试集,2018–2021 年为独立验证集(所有报告的指标均在 2020 年验证集上计算),这样的划分可有效防止时间泄漏。

在训练前,数据需经过规范化处理:输入与输出均按球面平均采用 z-score 或 min-max 归一化。其中,水汽相关量通过 min-max 归一化缩放到 [0,1],以满足非负约束;风场则假设均值为零,以总风速大小的标准差进行标准化,从而保留矢量方向信息。归一化常数的计算需先在球面上进行空间平均,再在整个训练集上进行时间平均。

正是这样经过精心选取、划分与处理的数据集,为 FCN3 在 1000+ GPU 上端到端学习全球大气的概率演化提供了坚实基础,保障了模型的有效训练和精准预报。

概率机器学习天气预报系统 FourCastNet 3

FourCastNet 3(FCN3)作为一款概率模型,由一个编码器、一个解码器和 8 个神经算子块组成,基于隐马尔可夫模型框架,给定某一时刻 tₙ 的 0.25° 网格大气状态 uₙ,模型能预测 6 小时后 tₙ₊₁ 的状态 uₙ₊₁=F_θ(uₙ, tₙ, zₙ),其中随机性通过随机噪声向量(noise variables)zₙ 引入——zₙ 来自多个具有不同空间-时间尺度的球面扩散过程采样,以此捕捉大气演化的不确定性。

1 分钟内完成 15 天预报,英伟达/UC 伯克利等提出机器学习天气预报系统 FCN3,支持单卡极速推理-AI.x社区

FourCastNet 3 模型示意图

在模型架构上,FCN3 采用球面神经算子设计,核心是局部与全局的球面群卷积——即在旋转群 SO(3) 作用下保持等变性的卷积。其中,全局卷积核在谱域参数化,借助球面卷积定理及球谐变换,类似经典伪谱方法;局部卷积则基于离散-连续(DISCO)卷积框架,用数值积分近似连续域卷积,支持各向异性滤波器,更贴合大气现象的几何特征。

整体架构分为编码器、处理器和解码器,其中编码器通过一层局部球面卷积将 721×1440 的输入/输出信号下采样到 360×720 高斯网格,嵌入维度为 641;处理器由若干采用 ConvNeXt 结构的球面神经算子块构成,实验表明 4 个局部块配 1 个全局块时预报技巧最佳,且省略层归一化以保留物理过程的绝对数值;解码器结合双线性球面插值与局部球面卷积上采样,恢复原始分辨率并抑制混叠。

值得注意的是,与多数机器学习天气模型预测「倾向」(预测与输入之差)不同,FCN3 直接预测下一时刻状态,有效抑制了高频伪影。此外,编码器与解码器均不跨通道混频,水汽通道经光滑样条输出激活函数处理,确保正值并减少高频噪声。

由于 FCN3 内部表征庞大,单张 GPU 显存难以承载,且自回归 rollout 需同时存放多次计算结果,内存压力显著。为此,如下图所示,研究团队采用混合并行策略实现可扩展训练:一方面借鉴传统数值方法的空间模型并行(域并行),通过空间域分解将模型和数据同时切分到不同 GPU,需将所有空间算法改写为分布式版本;另一方面结合数据并行,包括集合并行和批并行,因各成员在损失计算前互不依赖,仅需在损失阶段通信,效率极高。这些功能在 Makani 框架中实现,支撑扩展至数千张 GPU 。

1 分钟内完成 15 天预报,英伟达/UC 伯克利等提出机器学习天气预报系统 FCN3,支持单卡极速推理-AI.x社区

用于训练 FourCastNet 3 的模型和数据并行性的说明

训练过程分三阶段进行,其中初始预训练聚焦 6 小时预报技巧,利用 1980–2016 年 ERA5 训练集中的逐小时样本,构建以每个 UTC 整点为起点的 6 小时输入-目标对,在 NVIDIA Eos 超算 1,024 张 H100 上以批大小 16 、集合大小 16 训练 208,320 步,耗时 78 小时。

第二阶段预训练基于 6 小时初始场进行 4 步自回归 rollout,在 NERSC Perlmutter 系统 512 张 A100 上训练 5,040 步(每 840 步降低学习率),耗时 15 小时。

微调阶段针对 2012–2016 年的 6 小时样本,在 Eos 系统 256 张 H100 上进行,耗时 8 小时,以纠正潜在分布漂移并提升对近期数据的性能。因单卡 80GB 显存不足,训练中采用空间并行切片数据与模型,预训练切分为 4 份,微调因自回归需求更高切分为 16 份,最终实现了大模型的高效训练。

性能评估:FCN3 全面超越传统 NWP,且以极低成本追平当前最先进的扩散模型

为全面评估 FourCastNet 3(FCN3)的性能,研究团队从多个关键维度设计了实验,涵盖预报精度、计算效率、概率校准性及物理保真度等方面。如下图所示,在核心性能指标上,基于 2020 年(训练集之外)每 12 小时初始场的平均结果,FCN3 的连续分级概率评分(CRPS)与集合平均均方根误差(RMSE)表现优异,不仅全面超越传统物理数值天气预报的金标准 IFS-ENS,且与当前最佳数据驱动模型 GenCast 的差距已几乎可忽略。

在计算效率上,得益于一步式直接生成集合成员的设计,FCN3 在单张 NVIDIA H100 GPU 上仅需约 60 秒即可完成一次 15 天、 6 小时间隔、 0.25° 空间分辨率的预报,相较之下,GenCast 在 Cloud TPU v5 实例上完成同等时长预报(时间分辨率仅为 FCN3 的一半)需 8 分钟,而 IFS 在 96 颗 AMD Epyc Rome CPU 上以 9 km 业务分辨率运行约需 1 小时,若忽略硬件和分辨率差异,FCN3 相对 GenCast 提速约 8 倍,相对 IFS-ENS 提速约 60 倍。

1 分钟内完成 15 天预报,英伟达/UC 伯克利等提出机器学习天气预报系统 FCN3,支持单卡极速推理-AI.x社区

FourCastNet 3 的概率技能

考虑到 CRPS 与 RMSE 仅能进行逐点评估,无法衡量时空相关性,研究团队通过案例研究补充验证模型的物理保真度。如下图所示,以 2020 年 2 月 11 日 00 UTC 起报的温带风暴 Dennis 为例,预报其登陆爱尔兰及英伦诸岛前 48 小时的 850 hPa 风速与 500 hPa 位势高度的结果显示,FCN3 能再现这一天气事件,风速与气压场的协变关系合理,且 500 hPa 位势高度的角功率谱密度(PSD)保持正确斜率;即便预报时效延长至 30 天,角功率谱仍无衰减,预测始终保持锐利分辨率。

对 2020 年全年功率谱密度及其相对 ERA5 真值的相对误差分析表明,高波数区域误差始终有界(在 -0.2 至 0.2 之间),这得益于模型架构遵循的几何与信号处理原则,以及兼顾局部与全局分布的 CRPS 损失,促使模型学到正确的空间相关性。

1 分钟内完成 15 天预报,英伟达/UC 伯克利等提出机器学习天气预报系统 FCN3,支持单卡极速推理-AI.x社区

FourCastNet 3 对温带风暴 Dennis 的预测

相比之下,多数确定性机器学习天气模型存在高频信息衰减明显、预报结果模糊的问题;即便采用 CRPS 训练的混合模型 NeuralGCM,也在高频模态出现显著模糊;最新的概率模型 GenCast 与 AIFS-CRPS 则无法完整保留正确谱形,甚至出现高频模态积聚——这在传统数值天气预报中往往是模式发散的前兆。

综合对角谱、纬向谱及物理一致性的测试结果,FCN3 被证实是一款在概率技巧、计算效率、全球尺度上同时具备前所未有谱保真与物理真实性的天气模型,其预报在 60 天的次季节尺度仍能保持稳定,为次季节预报及大规模集合预报的发展铺平了道路。

1 分钟内完成 15 天预报,英伟达/UC 伯克利等提出机器学习天气预报系统 FCN3,支持单卡极速推理-AI.x社区


2020 年全年功率谱密度及其相对 ERA5 真值的相对误差

概率机器学习天气预报系统的突破与展望

事实上,全球产学研机构在概率机器学习天气预报系统领域的研究已相当深入,涌现出一系列颇具影响力的成果。

谷歌旗下的 DeepMind 团队推出的 GenCast,堪称该领域的标杆之作。作为基于条件扩散模型的概率天气模型,它能在 8 分钟内生成一组随机的 15 天全球预报,时间步长为 12 小时,分辨率达 0.25°,涵盖超过 80 种地表和大气变量。

在与全球顶尖的中期预报系统——欧洲中期天气预报中心(ECMWF)的集合预报(ENS)对比评估中,GenCast 在 1320 个评估指标中的 97.2% 都表现更优,其生成的边际和联合预报分布更为精准。

微软旗下的 Aurora AI 气象预报模型融合深度学习与大规模异构数据处理技术,不仅能精准预测天气,经微调后还可应用于洋流、空气质量等多种自然环境监测领域。其训练数据规模庞大,涵盖超过 100 万小时来自卫星、雷达、气象站及计算机模拟等多来源的气象与环境数据。

据微软研究团队测试数据显示,在 2022 至 2023 年全球热带气旋预测任务中,Aurora 的轨迹预测表现全面优于业界竞品和传统观测推理算法,同时在海浪预报和空气质量预测等复杂环境场景中也展现出高度精确性。

学术界的探索同样成果丰硕,诸多高校围绕概率机器学习天气预报系统开展深入研究并取得突破性进展。剑桥大学和艾伦图灵研究所的研究团队另辟蹊径,其开发的 Aardvark Weather 系统,是首个可在台式电脑上训练和运行、并能以单一 AI 模型取代天气预报流程所有步骤的系统,处理速度较传统方法快数千倍。

该系统可高效处理来自卫星、气象站和气象气球的多模态复杂数据,生成 10 天全球预报,在 4 个 NVIDIA A100 GPU 上,由观测数据生成完整预报仅需约一秒钟。

复旦大学伏羲团队提出的 FuXi Weather 系统,则是首个能够独立完成数据同化(DA)与循环预报的端到端机器学习全球天气预报框架。它通过融合多源卫星观测数据,以 0.25° 分辨率生成 10 天可靠预报,即便在非洲中部等观测稀疏地区,其表现甚至超越了欧洲中期天气预报中心(ECMWF)的高分辨率预报(HRES)。

这些探索与突破,不仅推动着概率机器学习天气预报系统向更高精度、更高效能、更广泛应用场景迈进,也为应对气候变化、减轻气象灾害影响、优化能源利用等全球性议题提供了强有力的技术支撑。随着技术的持续迭代与跨领域合作的深化,未来的概率机器学习天气预报系统将更精准地捕捉大气的复杂动态,筑起更坚实的气象保障防线。

参考文章:

1.https://mp.weixin.qq.com/s/fOEsBxHRxKR5yMXW9gr5Sw

2.https://mp.weixin.qq.com/s/T3IRjxUg5zL0X7_0a96F3Q

©著作权归作者所有,如需转载,请注明出处,否则将追究法律责任
收藏
回复
举报
回复
相关推荐