On Neural Scaling Laws for Weather Emulation through Continual Training¶

会议: NeurIPS 2026
arXiv: 2603.25687
OpenReview: https://openreview.net/forum?id=u9jpPYohUA
代码: https://github.com/ShashankSubramanian/neural-scaling-weather
领域: 预训练
关键词: 神经网络缩放律, 气象模拟, 持续训练, 算力最优, Swin Transformer

一句话总结¶

本文采用极简 Swin Transformer 架构与“恒定学习率+周期性冷却”持续训练策略，系统构建了覆盖两个数量级算力（$6\times 10^{17} \sim 6\times 10^{19}$ FLOPs）的气象模拟神经缩放曲线并确定了算力最优平衡点，同时揭示了外推至超十亿参数时多轮训练引起的过拟合饱和限制。

研究背景与动机¶

近年来，以 FourCastNet、Pangu-Weather、GraphCast 为代表的数据驱动深度学习气象模拟系统在预报精度上已经全面比肩甚至超越了求解多尺度流体力学方程的传统数值天气预报（NWP），且在单张 GPU 上的推理速度快出数个数量级。然而，随着模型参数向数十亿级别跃迁以及时空分辨率不断提升，预训练这些气象大模型的算力开销急剧攀升。为了追求更优的预报技巧，学界提出了层出不穷的领域特异性架构改进、复杂的球谐几何变换以及定制化多目标损失函数。这一狂热趋势促使研究人员反思：当前的性能红利究竟源于那些精巧繁琐的领域架构特异性设计，还是遵循了大语言模型（LLM）领域中由“苦涩教训（Bitter Lesson）”所支配的规模法则（Scaling Laws）？

探索气象科学机器学习（SciML）神经缩放律的核心矛盾在于算力成本与架构混淆因素的交织。在传统 NLP 领域，Chinchilla 缩放律依赖从零训练不同模型规模和数据量的试验，并强依赖于与总训练轮次严格绑定的余弦学习率衰减（Cosine Schedule），这在处理高分辨率时空张量时会导致灾难性的计算浪费；更严重的是，已有关于气象模型缩放的研究要么依赖易受硬件波动干扰的 GPU 耗时而非 FLOPs，要么混杂了海况与污染物等多源异构数据，抑或在粗糙网格下过早得出缩放停滞的片面结论，无法回答“在固定算力预算下如何最优分配参数规模与数据样本”这一根本问题。

本文的切入视角是剔除所有领域特定的架构附加组件，采用通用视觉主干探索纯粹的规模扩展行为，并引入基于恒定学习率和快速冷却阶段（Cooldown）的持续训练框架以低成本绘制 IsoFLOP 曲线。核心 idea：以极简 Swin Transformer 为骨干网络，通过恒定学习率加末端 5% 迭代周期的快速冷却实现多算力梯度的无缝持续训练，在 $6\times 10^{17} \sim 6\times 10^{19}$ FLOPs 预算下首次拟合出天气预报的算力最优缩放幂律，并将冷却阶段复用于多步自回归（AR）或频谱修正（AMSE）下游任务对齐，同时论证了十亿参数外推在多轮数据重复下的过拟合边界。

方法详解¶

整体框架¶

本文构建的气象预报系统将全球大气连续演化建模为高维时空离散状态的多通道回归任务。在空间分辨率为 $0.25^\circ$ 的全球网格上，输入状态包含覆盖 13 个垂直气压层的高空物理量与关键地表变量，构成维度为 $71 \times 721 \times 1440$ 的高维物理张量。网络架构剥离了一切复杂的球谐变换或图神经网络拓扑，直接采用标准的移动窗口自注意力机制（Shifted Window Self-Attention, Swin Transformer），通过空间分块投影、三维球面坐标与归一化时间嵌入，输入到多层堆叠的等宽 Transformer 块中，最终以标准均方误差（MSE）单步预测向前推进 6 小时的大气状态。

为了在巨大的显存瓶颈下实现高效且可预测的模型缩放，整个实验流水线由两大核心支柱构成：一是用于大幅摊薄缩放探索算力成本的“恒定学习率 + 周期性冷却（Constant + Cooldown）”持续训练范式；二是解决高分辨率时空特征显存激增的二维空间并行（Spatial Domain Decomposition Parallelism）通信算法。训练完成后，短期的冷却阶段还可直接接入自回归多步展开或频谱调整损失，实现下游业务指标的免重新预训练高效对齐。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["全球大气状态输入 u_n<br/>71变量 × 721纬度 × 1440经度"] --> B["分块投影与静态球面坐标嵌入<br/>无参数化位置编码，Patch Size=4"]
    B --> C["二维空间域分解并行与分布式平移<br/>2D GPU 网格无缝切分中间激活"]
    C --> D["极简 Swin Transformer 主干<br/>去层级/去相对位置偏置，保留 QK-Norm"]
    D --> E["持续预训练引擎<br/>恒定学习率前推至目标算力预算"]
    E -->|分支1：缩放诊断| F["5% 周期退火冷却<br/>单步 MSE 拟合 IsoFLOP 算力最优曲线"]
    E -->|分支2：下游对齐| G["复用冷却阶段对齐<br/>4/8步自回归展开 或 AMSE 频谱损失"]
    F --> H["算力最优气象预报模型与外推边界诊断"]
    G --> H

关键设计¶

1. 极简 Swin 主干与显存感知二维空间并行：消除架构混淆并突破显存瓶颈

针对传统气象模型过度依赖复杂归纳偏置的问题，模型刻意移除了 Swin Transformer 中的多尺度分层下采样（Patch Merging）、相对位置偏置（Relative Position Bias）及自定义层，仅保留均匀深度的 Transformer 块，块内由局部窗口多头自注意力（W-MHSA）和多层感知机（MLP）构成，并结合 Pre-RMSNorm、残差连接以及保证大模型稳定性的 QK-Normalization。对于地球球面的网格曲率，模型直接将基于纬度、经度投射的地球三维欧氏坐标 $(x, y, z)$ 与归一化时间戳 $t$ 拼接入无偏置的轻量投影层作为位置编码，并在交替块的窗口平移掩码中加入东西向循环边界条件以保持地球自转维度的经度周期性。

在分辨率 $721 \times 1440$ 与 Patch 尺寸为 4 的设置下，单个时间步生成的序列长度高达 $64,800$。由于自注意力前向计算与梯度回传中的中间激活显存呈 $\mathcal{O}(B \cdot H \cdot W \cdot E)$ 增长，当引入多步自回归展开时，即便是全分片数据并行（FSDP）也无法缓解单卡显存溢出的瓶颈。为此，本文实现了一种正交于数据并行度 $n_d$ 的 $n_1 \times n_2$ 二维空间域分解并行（Spatial Parallelism）。为了支持 Swin 架构核心的窗口跨卡周期循环平移（Shifted Window Roll），算法利用 PyTorch 的 all_to_all_single 算子实现了一维环形邻居切片交换的分布式 Roll 操作，结合针对权重梯度的联合 AllReduce 通信钩子，将单卡激活显存严密压缩至 $\mathcal{O}(B \cdot H \cdot W \cdot E / (n_d n_1 n_2))$，使得 456.7M 参数量的模型及多步 AR 展开能够在 64 至 256 张 A100 GPU 上平稳训练。

2. 恒定学习率加周期冷却的持续训练：低成本绘制 IsoFLOP 曲线

在构建神经缩放律时，传统的 Chinchilla 方案必须对每一个模型参数规模和每一个训练步数从头训练（Scratch），因为余弦衰减学习率（Cosine LR）的性能最大化严格要求衰减周期匹配预设的总步长，使得跨预算探索的开销呈爆炸式累积。受现代语言模型持续训练理论启发，本文采用“恒定最大学习率 + 快速冷却”调度策略：模型在大规模预训练期间始终保持固定的学习率 $\eta_{\max}$，仅在达到特定 FLOPs 预算前的最后 5% 迭代步数内执行非线性快速退火至 0（采用 $1-\sqrt{t/T_{\text{cool}}}$ 退火曲线）。当需要探索更高的算力预算时，直接从冷却前的恒定学习率检查点恢复状态继续训练，无需重新启动。

实验严格证实，该策略在验证损失上系统性击败了同等总步长下从零训练的余弦学习率基线。不仅如此，消融研究表明，即使冷却周期短至总步数的 0.5%，验证损失也能迅速陡降，而在达到 5% 时性能收益即完全趋于饱和。这种恒定学习率持续训练机制使得同一个基础模型只需完整训练一次，即可在轨迹上定期“分叉”冷却出属于不同 IsoFLOP 预算的候选点，将原本无法承受的网格级大规模 SciML 缩放实验开销压缩至与单次线性扩展相当。

3. 下游任务适配对齐的冷却阶段复用：解耦缩放分析与复杂业务目标

在实际气象预报业务中，简单的单步 MSE 预训练往往无法满足下游评估的核心需求：一方面，自回归循环推演在多步展开中因误差累积和分布偏移，需要引入多步预测训练以抑制发散；另一方面，由于大气混沌效应以及 6 小时离散采样带来的高频欠采样，纯 MSE 目标倾向于对不确定性细节求统计均值，导致生成风场和湿度场在频域上严重平滑模糊（即高波数功率衰减）。若在整个大规模缩放预训练阶段直接混入这些复杂目标，将导致计算开销翻倍且无法解耦真正的架构缩放行为。

本文提出将持续训练的最后 5% 冷却阶段“重塑用途（Re-purposing Cooldown）”，作为轻量级的下游任务对齐接口。研究团队在冷却阶段分别探索了两种对齐模式：（1）多步自回归展开微调（4-step AR），强制模型在冷却期内以滚动预测误差进行优化，大幅压低长达 10 天预报周期的全局累计均方根误差（RMSE）；（2）球谐修正均方误差（Adjusted MSE, AMSE），将预测场与真实场的球面谐波展开系数分解为振幅误差与空间相干性惩罚： $$\text{AMSE}(u, \hat{u}) = \sum_{k} \left( \left(\sqrt{\text{PSD}_k(u)} - \sqrt{\text{PSD}_k(\hat{u})}\right)^2 + 2\max\left(\text{PSD}_k(u), \text{PSD}_k(\hat{u})\right)\left(1 - \text{Coh}_k(u, \hat{u})\right) \right)$$ 其中 $\text{PSD}_k$ 为第 $k$ 阶球谐度数下的功率谱密度，$\text{Coh}_k$ 为谱相干性度量。借助在冷却期的短程对齐，模型在无需重头预训练的前提下，既能在 AR 冷却下匹敌顶尖 SOTA 气象模型的长周期精度，又能在 AMSE 冷却下完美保留 ERA5 真实大气场的高波数能量，实现了基础模型规模缩放律与下游专用能力调优的完美解耦。

损失函数 / 训练策略¶

模型在主预训练阶段使用全局未加权的均方误差损失（MSE）： $$\mathcal{L}_{\text{pretrain}} = \frac{1}{B \cdot H \cdot W \cdot C} \sum_{b, h, w, c} \left( \mathcal{F}_\theta(u_n)_{b,h,w,c} - u_{n+1, b,h,w,c} \right)^2$$ 所有模型均使用全局 Batch Size 16 进行训练。优化器采用 AdamW，权重衰减固定为 $1\times 10^{-4}$，梯度全局范数裁剪阈值设为 1.0。学习率设定遵循参数规模扩展律：大于 100M 参数的模型采用 $\eta_{\max} = 5\times 10^{-4}$，而较小参数模型具备更高的学习率容忍度，峰值学习率随参数量的平方根倒数自适应缩放（最高达到 $1.5\times 10^{-3}$）。主干预训练全部采用 FP16 混合精度，数据读取利用 NVIDIA DALI 在 GPU 侧实现流水线重叠与均值方差归一化。

实验关键数据¶

主实验：算力最优缩放拟合与基准模型对比¶

实验在 ERA5 历史气象再分析数据集（1979-2016 为训练集，2017 为验证集，2020 为测试集）上展开，测试基于 WeatherBench 2 标准协议评估 732 个独立初始场在 10 天（240小时）范围内的预报技能。研究在 $6\times 10^{17}$ 至 $6\times 10^{19}$ FLOPs 跨度内构建了 4 条 IsoFLOP 曲线，各曲线包含 3.1M 到 456.7M 不等的 12 种模型规模，通过抛物线拟合提取最优极小值点，推导出数据规模与参数规模的最优算力分配幂律： $$S^\star(C) \propto C^{0.41}, \quad N^\star(C) \propto C^{0.59}$$ 下表展示了在 $6\times 10^{19}$ FLOPs 预算下的最优模型配置（203.8M 参数量，经过 4 步 AR 冷却）与传统欧洲中期天气预报中心数值模型 HRES 及代表性深度学习大模型 GraphCast 的预报误差（纬度面积加权 RMSE）对比：

评估物理变量	预报前置时间 (Lead Time)	传统数值预报 HRES (0.1°)	深度学习基准 GraphCast	本文 Compute-Optimal Swin (204M)	相对 HRES 提升
Z500 (位势高度, m²/s²)	24 小时 (1天)	85.2	58.1	56.4	-33.8%
Z500 (位势高度, m²/s²)	120 小时 (5天)	362.0	280.4	282.1	-22.1%
Z500 (位势高度, m²/s²)	240 小时 (10天)	815.3	695.2	718.5	-11.9%
T850 (温度, K)	24 小时 (1天)	1.15	0.82	0.80	-30.4%
T850 (温度, K)	120 小时 (5天)	2.48	1.88	1.86	-25.0%
T850 (温度, K)	240 小时 (10天)	3.82	3.22	3.35	-12.3%
T2M (地表2米温度, K)	120 小时 (5天)	1.88	1.45	1.42	-24.5%
U10M (地表10米纬向风, m/s)	120 小时 (5天)	3.42	2.68	2.74	-19.9%

注：Swin 仅使用 4 步 AR 冷却，而 GraphCast 在预训练后经过了 2~12 步长周期 AR 渐进微调；在推理耗时上，204M Swin 在单张 A100 上生成 10 天预报仅耗时 32 秒，较传统数值预报实现超 $100\times$ 吞吐加速。

消融实验：学习率调度与冷却策略机制分析¶

针对训练调度与下游对齐策略的有效性，研究在 115.1M 参数量 Swin 模型上进行了系统对比测试，量化不同调度方式以及不同冷却时长对验证集多步滚动误差（36小时平均 MSE）与频谱高波数保留能力的影响：

训练调度方案	冷却迭代占比 (%)	对齐损失类型	36小时滚动验证 MSE	高频波数功率谱匹配度 (PSD)	机制说明
Cosine Scheduler (Scratch)	0% (常规余弦衰减)	标准 MSE	0.00762	显著耗散衰减	传统全周期衰减基线
Constant + Cooldown	0.5%	标准 MSE	0.00741	显著耗散衰减	极速退火即大幅释放收敛潜力
Constant + Cooldown	2.0%	标准 MSE	0.00715	显著耗散衰减	验证误差随冷却步长平稳下行
Constant + Cooldown (默认)	5.0%	标准 MSE	0.00698	显著耗散衰减	最佳性价比，收益在此拐点饱和
Constant + Cooldown	10.0%	标准 MSE	0.00696	显著耗散衰减	继续拉长冷却步长无显著增益
Cooldown w/ 4-step AR	5.0%	4步自回归 MSE	0.00612	重度平滑平抑	压低滚动累积误差，牺牲高频纹理
Cooldown w/ AMSE	5.0%	修正球谐谱损失	0.00702	完美吻合真实 ERA5 频谱	保持高频涡旋与极端梯度真实感

关键发现¶

算力最优指数向参数量倾斜：在大气时空模拟中，推导出的算力最优缩放幂律 $N^\star \propto C^{0.59}$ 显著高于数据样本指数 $S^\star \propto C^{0.41}$，表明对于物理连续场，扩大模型容量以捕捉复杂的空间局部相互作用比盲目喂入重复采样的时间序列具有更高的算力收益。
十亿参数外推遭遇多轮过拟合饱和：当将推导出的经验缩放律外推一个数量级至 $2.25\times 10^{21}$ FLOPs 并训练 1.3B 参数模型时，虽然训练损失沿着理论轨迹降至 0.033，但验证损失却在 0.053 处停滞并早早出现平台期。原因在于 40 年 ERA5 数据仅包含约 35 万个独立采样，使得 1.3B 模型被迫训练了超过 13 个 Epoch，导致模型开始对气象样本产生记忆式过拟合。
高波数频谱质量随算力单调改善：随着模型沿算力最优轨迹缩放，虽然小模型也能拟合宏观低频长波环流，但高频小尺度涡旋往往产生虚假伪影（如台风风圈的块状断裂）；加大算力与模型规模能系统性抑制高频伪能生成，在超强台风（如 2019 年 Hagibis 登陆）的 3~5 天前期追踪中精准锁定气旋中心。

亮点与洞察¶

极简通用骨干击碎“领域架构神话”：证明无需任何球谐神经算子（FNO/SFNO）或复杂图网络（GNN）的特定归纳偏置，纯粹的通用视觉 Swin Transformer 依靠纯规模法则就能达到最先进气象基准的预报精度，印证了“苦涩的教训”在科学机器学习领域的普适性。
持续训练与下游任务对齐的优雅解耦：将恒定学习率持续训练与末端 5% 冷却阶段巧妙复用，既规避了 IsoFLOP 扫描中千百次“从头训练”的算力黑洞，又允许在免重新预训练的情况下自由切换 AR（保 RMSE 指标）与 AMSE（保频谱真实感）两种互斥的物理下游目标。
SciML 扩展的明确数据容量红线：首次在大气科学预报中定量捕获了“模型规模受限于自然历史数据量”的过拟合临界点，明确告诫社区在缺乏多源多模态地球科学数据接入之前，单凭在单分辨率单数据集上狂暴堆叠十亿以上参数将遭遇边际效益归零。

局限与展望¶

确定性预报的固有混沌天花板：当前全流程聚焦于确定性单点预测，而在 7~10 天以上的中长期天气推演中，大气动力学混沌性要求依赖集合概率预报（Ensemble Forecasting）；未来的缩放律应拓展至扩散模型（Diffusion）或 CRPS 概率评价体系下的泛化表现。
固定 Patch 尺寸的单维缩放：本研究在缩放中固定了 Patch Size=4（对应序列长度 64,800），未将空间分块粒度与输入序列长度作为一个独立的缩放自由度进行联合正交探索。
跨领域地球大模型的延展诉求：为了跨越多轮训练导致的过拟合瓶颈，未来的气象基础模型亟需打破单套 ERA5 边界，融合对流层数值模拟高分辨率切片、卫星雷达遥感与海-冰-气多圈层异构观测数据，在更广阔的数据分布上重塑缩放边界。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统性推导气象模拟领域算力最优缩放律，开创性地将持续训练冷却与气象物理对齐机制相融合。
实验充分度: ⭐⭐⭐⭐⭐ 从 3.1M 到 1.3B 参数量覆盖多个算力数量级，从 IsoFLOP 曲线到极端台风个例分析，证据链极其扎实。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰严密，对 SciML 领域迷信架构设计的反思引人深思，问题定义与推导逻辑环环相扣。
价值: ⭐⭐⭐⭐⭐ 为未来耗资巨大的地球系统大模型预训练提供了不可或缺的算力预算指南与过拟合边界预警。