跳转至

Scaling Laws of Global Weather Models

会议: ICML 2026
arXiv: 2602.22962
代码: https://github.com/spcl/scaling-laws-weather-model
领域: 地球科学 / 天气预报 / Scaling Laws
关键词: 天气预报模型, 缩放定律, 计算最优, 宽度优先, ERA5

一句话总结

本文在统一的训练/评测协议下,对 5 个主流数据驱动天气模型(Aurora、AIFS、Pangu、GraphCast、SFNO)做了首个跨模型的缩放定律分析,发现天气模型偏爱"加宽而非加深"、计算预算应优先投给更多训练数据而非更大模型,且不同气象变量的缩放行为差异巨大——这些规律与 NLP/视觉的缩放定律截然不同。

研究背景与动机

领域现状:数据驱动的神经网络天气预报(GraphCast、Pangu、Aurora 等)正在快速逼近甚至超越传统数值天气预报(NWP)。随着高分辨率再分析数据(如 ERA5)与大规模训练基础设施的普及,这些模型即将进入"靠堆规模往上冲"的阶段。

现有痛点:NLP 和视觉领域早已有成熟的缩放定律(Kaplan、Chinchilla)指导"该把算力投给参数还是数据",但天气预报领域几乎没有系统性的跨模型缩放研究。现有的少数工作只针对单一模型、单一配置,无法回答"不同架构的缩放效率谁更高、规律是否通用"。

核心矛盾:直接把语言/视觉的缩放定律套到天气模型上并不安全。大气是一个有可预报性物理上限的混沌系统;天气模型要同时预测上百个物理性质各异、难度不同的相关变量(温度、风速、气压等);而且 GNN、Transformer、傅里叶算子这几类架构对空间分辨率的依赖方式根本不同,缩放行为很可能各走各路。

本文目标:在完全统一的实验条件下,刻画验证损失 \(\mathcal{L}\) 如何随模型规模 \(N\)、数据量 \(D\)、算力预算 \(C\) 以及模型形状(宽 vs 深)变化,并回答"固定算力下该怎么分配资源"。

切入角度:作者不发明新模型,而是把 5 个代表性模型放进同一个"风洞"里——同样的 ERA5 数据、同样的分辨率、统一的验证损失函数——逐一变动 \(N/D/C\)/形状,用 log-log 线性回归拟合幂律系数,从而把架构差异从外部噪声中剥离出来。

核心 idea:用"统一标准化的跨模型缩放实验"来揭示天气预报独有的缩放规律,把缩放定律从"单模型经验"上升为"领域级设计指南"。

方法详解

整体框架

本文本质是一项大规模实证研究而非新方法,整体可以理解为一条"标准化风洞 → 多维扫描 → 幂律拟合 → 提炼规律"的分析流水线。输入是 ERA5 再分析数据与 5 个待测模型,输出是一组缩放定律系数和 4 条可操作的设计结论。

具体地:所有模型都在 ERA5(1979–2020 训练、2021 验证、\(0.25^\circ\times0.25^\circ\) 全球网格、6 小时间隔)上做 6 小时预报;通过取各模型输入/目标变量的最大公共集、统一学习率调度、统一权重初始化等手段把"非架构因素"压到最小;再分别沿 \(N\)(参数量)、\(D\)(累计训练样本量,按 scaling 文献惯例指"已处理的样本总量"而非固定数据集大小)、\(C\)(算力)和模型形状四个轴扫描,每个配置拟合验证损失对各因素的幂律:

\[\mathcal{L}(D)=\alpha D^{-\beta},\quad \mathcal{L}(N)=\gamma N^{-\delta},\quad \mathcal{L}(C)=\lambda C^{-\epsilon}\]

其中 \(\beta,\delta\) 越大表示随数据/参数增长损失下降越快、长期潜力越好。整个实验耗费超过 43 万 GPU 小时。下面四个"关键设计"其实是这条流水线上的四个核心分析维度与对应发现。

关键设计

1. 统一标准化评测协议:把架构差异从噪声里剥出来

跨模型比较最大的陷阱是各模型默认的训练损失、变量加权、数据加载方式都不一样,直接比损失等于在比"苹果和橘子"。本文的方法学基石是强制对齐:所有模型用同一个验证损失 \(\mathcal{L}\)——对预测与真值的平方差 \((\hat{x}-x)^2\) 在空间网格与大气变量上做加权平均,每个变量按自身标准差归一化,每个网格按归一化面积加权(处理地球球面几何),高空变量再额外按气压层加权(与 GraphCast 一致,越高气压层越重要)。同时统一分辨率、统一学习率调度、统一初始化。这样后续观察到的任何性能差异,才能可信地归因于架构本身,而不是优化技巧或数据分布漂移。作者还试了 Maximal Update Parameterization(\(\mu\)P)来稳定跨尺度训练,但发现 \(\mu\)P 对 AIFS 无效,说明这套对齐协议是逐模型调出来的而非照搬。

2. 数据/参数双幂律拟合:用 \(\beta\) 而非 \(\alpha\) 判断长期潜力

针对 \(D\),本文固定参数量、变动数据量拟合 \(\mathcal{L}(D)=\alpha D^{-\beta}\)。关键洞察是:截距 \(\alpha\) 不能直接比,因为模型的损失排名会随 \(D\) 改变——例如 GraphCast 在 \(D=30\) TB 时损失最低,但到 \(D=100\) TB 时被 Aurora 反超;真正稳健的长期指标是斜率 \(\beta\)。结果 Aurora 的 \(\beta\approx0.51\) 显著高于其余模型(\(0.30\)\(0.46\)),意味着它从更多数据里榨取信息的效率最高(数据增 10 倍验证损失最多降 3.2 倍)。针对 \(N\),固定数据量拟合 \(\mathcal{L}(N)=\gamma N^{-\delta}\),五个模型都呈稳定幂律,GraphCast 在同等或更小 \(N\) 下损失最低、参数效率最佳。一个有意思的发现是 Pangu 存在"阈值效应":15 TB 时 \(\delta\) 偏低(加参数没用),到 30 TB 时 \(\delta\) 上升,说明参数缩放被数据量卡脖子——只有数据够多,加参数才划算。

3. 宽度优于深度:天气预报靠表征容量而非深层非线性

这是与 NLP 缩放定律最尖锐的分歧。Kaplan 等发现语言模型在固定 \(N\) 时性能几乎不依赖形状(宽深可换),但本文在固定参数预算下对每个模型各取一宽一窄两种配置(参数量大致相同、形状不同)对比,发现所有模型的宽配置都一致优于窄配置。更极端的是,GraphCast 和 SFNO 在深度为 1 时依然表现良好。作者用几何深度学习(线性图模型可媲美深层模型)和神经算子理论(单层谱/注意力即可逼近复杂全局优化步骤)来解释:6 小时短程天气动力学可能用近似线性的模型就能逼近,深层非线性堆叠是冗余的。结论是设计天气模型时,固定参数预算下应优先加宽——更大的表征容量比更多的非线性变换更值钱。

4. 计算最优分配:固定算力时把资源投给数据而非参数

固定算力 \(C\) 时,把验证损失对 \(D\) 拟合成抛物线(Chinchilla 式做法),抛物线最低点给出 \(N\)\(D\) 的最优配比,形式化为 \(N_{\text{opt}}\propto C^{a}\)\(D_{\text{opt}}\propto C^{b}\)。一个微妙之处是 NLP 的约束 \(a+b=1\) 在天气模型里并不普遍成立:GraphCast/AIFS 的骨干跑在与输入分辨率无关的隐图上,违反了 \(C\approx 6ND\) 的标准假设(本文固定分辨率,故 \(C\sim N\cdot B\));而 Transformer 类的 patch 机制把有效 \(D\) 缩了 \(p^2\) 倍,得到 \(C\approx 6ND/p^2\),此时 \(a+b=1\) 才保持。尽管各模型最优配比不同,抛物线行为普遍存在,且趋势一致指向:随算力增长,优先延长训练(喂更多数据)比增大模型更能降低预报误差——宁可用小模型训久一点,也别用大模型训不充分。

一个完整示例

不适用。本文是缩放定律的实证分析,核心是跨模型扫描与幂律拟合,没有需要"走一遍流程"的单样本推理过程,故省略。

实验关键数据

主实验

五个模型在 ERA5 上统一评测,数据缩放系数 \(\beta\)(越大越好)与硬件利用率对比如下:

模型 参数量 骨干 \(\beta\uparrow\)(数据缩放) 备注
Aurora 1.3B Swin Transformer 0.51 数据缩放效率最佳
AIFS 255M Graph Transformer 0.46 球面图、降低经纬畸变
Pangu 276M Swin Transformer 0.43 参数缩放有数据量阈值效应
GraphCast 36.7M GNN 0.36 参数效率最佳但硬件利用率低
SFNO 433M 球面傅里叶算子 0.34 深度=1 仍表现良好

硬件利用率分析

缩放定律通常把算力当静态资源、忽略 wall-clock,但本文实测发现"验证损失缩放好"不等于"算力用得好":

模型 Tflop/s GPU 峰值 (Tflop/s) GPU 利用率 (%)
Aurora 368 989 (32-bit) 37.2
AIFS 33.7 1979 (16-bit) 1.70
GraphCast 10.15 989 (32-bit) 1.03
Pangu 3.25 989 (32-bit) 0.33
SFNO 0.215 989 (32-bit) 0.022

关键发现

  • 数据缩放冠军 ≠ 参数缩放冠军:Aurora 数据缩放最强(\(\beta\approx0.51\)),但参数效率冠军是 GraphCast(同等 \(N\) 损失最低);两者优势互补、不可兼得。
  • 理论效率 ≠ 实际效率:GraphCast 参数效率最优,但其 GNN 消息传递受内存带宽限制,H100 上单精度利用率仅 1.03%,而 Aurora 达 37.2%——约 36 倍差距,说明工程实现对实际部署效率至关重要。
  • 变量异质性:不同气象变量缩放行为差异显著。10m 风速 u 分量(10U)上 Aurora RMSE 全程最低,但 2m 温度(2T)上 GraphCast 反超所有模型。因此"全变量加权损失"只是模型性能的粗略指标,必须看变量级缩放。

亮点与洞察

  • "风洞式"统一协议是这篇论文最大的方法学贡献:跨模型缩放比较历来因损失函数/加权不一致而难以可信,本文用统一验证损失 + 统一训练配置把架构因素单独拎出来,这套思路可迁移到任何"多模型横评"的领域。
  • 用斜率 \(\beta\) 而非截距 \(\alpha\) 判断长期潜力:很巧妙地点破了"在某个数据量下谁损失低"是会翻盘的,真正该看的是缩放速度——这对资源有限、只能在小规模上做实验再外推的团队尤其有价值。
  • "宽优于深"颠覆了 NLP 直觉:在固定参数预算下加宽比加深更划算,甚至深度=1 都能打,这个结论可以直接改变下一代天气模型的架构设计取向。
  • 把硬件利用率纳入缩放讨论:指出经典缩放定律把算力当静态资源、忽略 wall-clock 是一个盲点——36 倍的利用率差距说明"缩放曲线漂亮"和"训练跑得快"是两回事。

局限与展望

  • 算力/显存约束限制了扫描范围:最大可探索宽度仅 512、Aurora 最小宽度受单注意力头大小(64)限制,导致部分模型(Aurora、GraphCast)的计算最优抛物线只看得到左半边(尚未收敛),右半边需要更小宽度才能补全,外推存在不确定性。
  • 只覆盖 6 小时短程预报:"深度=1 即可"的结论很可能只对 6 小时这种近线性的短程动力学成立,更长 lead time 或自回归滚动预报是否仍偏爱宽浅结构尚未验证。
  • \(\mu\)P 等跨尺度稳定技术未普适:作者发现 \(\mu\)P 对 AIFS 无效,说明这套缩放结论部分依赖逐模型手调的超参,可复现性需谨慎。
  • 改进方向:可引入变量级的混合建模或变量自适应调参,以及把 wall-clock/利用率正式纳入计算最优框架,让缩放定律更贴近真实部署。

相关工作与启发

  • vs Kaplan / Chinchilla(NLP 缩放定律):他们建立了 \(N\propto C^{0.73}\)\(D\propto C^{0.27}\) 的语言模型计算最优法则且性能几乎不依赖形状;本文发现天气模型 \(a+b=1\) 不普遍成立、且强烈偏爱宽度,证明缩放规律是领域特异的,不能照搬。
  • vs 单模型天气缩放研究(如 Nguyen 等、Couairon 等):以往多聚焦单一模型/单一配置;本文首次做跨 5 个模型的统一横评,把缩放从"个例经验"提升为"架构选型指南"。
  • vs 经典缩放定律的算力假设:经典工作把 compute 当静态资源、忽略硬件利用率与 wall-clock;本文用 Aurora 对 GraphCast 的 36× 利用率差距指出这一盲点,启发后续把工程效率纳入缩放分析。

评分

  • 新颖性: ⭐⭐⭐⭐ 首个跨模型天气缩放分析,"宽优于深"是反直觉的领域级发现。
  • 实验充分度: ⭐⭐⭐⭐⭐ 5 模型 × 多维扫描、43 万 GPU 小时、统一协议,覆盖面扎实。
  • 写作质量: ⭐⭐⭐⭐ 逻辑清晰、图表丰富,但部分抛物线只有半边、结论外推需谨慎。
  • 价值: ⭐⭐⭐⭐⭐ 直接为下一代天气模型的架构与资源分配提供可操作指南。