Otter Weather: Skillful and Computationally Efficient Medium-Range Weather Forecasting¶
会议: NeurIPS 2026
arXiv: 2606.26421
OpenReview: https://openreview.net/forum?id=w4fZFKKSwS
代码: https://github.com/cambridge-mlg/otter
领域: 时间序列
关键词: 中期天气预报, Swin U-Net, 连续排位概率评分 (CRPS), Muon 优化器, 计算效率
一句话总结¶
Otter Weather 彻底剥离气象领域专有的复杂几何归纳偏置,采用标准二维 Swin-UNet 结合 2D RoPE、SwiGLU 与 Muon 矩阵正交优化器,仅需单卡不足 3.5 A100-天的极低训练成本即可超越数值预报基线,并基于 CRPS 与副产物深度集成推进概率预报的精度-算力帕累托前沿。
研究背景与动机¶
基于数据驱动的人工智能中长期全球天气预报正在深刻改变气象领域。以 GraphCast、Pangu-Weather、AIFS 和 Aurora 为代表的 SOTA 模型相继在预测技能上逼近甚至超越传统数值天气预报(NWP)系统。然而,这些突破性系统往往伴随着极高的计算门槛,通常需要由数百张 GPU 或 TPU 组成的庞大集群训练数百个计算日。这种高昂的计算壁垒不仅将无庞大算力支持的高校、中小型科研机构和发展中国家气象部门排除在外,损害了学术界的独立研究主权,更由于迭代周期极其冗长,严重制约了模型架构探索与快速实验创新。
这一高昂算力成本的根源在于当前主流模型所施加的沉重领域专有归纳偏置。现有前沿工作普遍假设地球大气的特殊几何特性必须通过专门的架构组件来建模,例如复杂的二十面体图神经网络(GNNs)、球面傅里叶神经算子(SFNO)或带有严格极地柱面注意力掩码的 3D 转换器。这些定制化组件虽然具备物理直觉,但对硬件极不友好,破坏了底层张量核心的并行优化,且无法直接复用大语言模型(LLM)和计算机视觉领域中高度优化的现代算子基础设施;而在概率预报方向上,基于扩散机制的生成模型(如 GenCast)在推理端需要数十步缓慢的迭代去噪,进一步加剧了计算负担。
本文切入的核心角度是打破“气象预测必须依赖复杂领域几何偏置”的固有假设,全面转向通用、高度硬件友好的现代机器学习组件。既然大语言模型与视觉 Transformer 的通用优化已历经千锤百炼,直接复用其核心进展并辅以严谨消融即可大幅压缩训练开销。核心 idea:构建基于恒定通道宽度的标准二维 Swin-UNet 架构,融入 2D RoPE、SwiGLU 与 Muon 矩阵优化器并移除极地注意力掩码,利用连续排位概率评分(CRPS)微调与超参搜索副产物构建“零额外开销”的深度集成,将单卡训练预算压至 3.5 A100-天内并重塑确定性与概率天气预报的精度-算力帕累托前沿。
方法详解¶
整体框架¶
Otter Weather 的整体预测流程旨在从历史大气状态序列自回归预测未来的全球天气演变。模型在空间分辨率为 \(1.5^\circ\)(网格尺寸 \(H=121, W=240\))的 ERA5 再分析数据集上运行,时间步长为 \(\delta t = 6\) 小时。模型输入包含过去 \(L=4\) 个时步的大气物理场,涵盖 13 个标准气压层上的 6 个高空变量(位势、温度、比湿、垂直速度、U/V 风速分量)、4 个地表变量(2 米温度、海平面气压、10 米 U/V 风速)以及 5 个静态地形特征,总计 87 个物理通道,同时拼接 32 维多尺度连续傅里叶时间编码。所有通道直接在特征维度展平,通过单层跨步卷积打块编码后送入对称的恒定宽度二维 Swin-UNet 骨干网,输出未来时步的残差增量 \(\Delta X\),并通过多步自回归展开生成中长期预报轨迹。在概率预报设定下,模型通过公平 CRPS 目标函数进行微调,并借助蒙特卡洛 Dropout(MC Dropout)和调参副产物形成深度集成。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
Input["多时序气象场输入<br/>87个物理变量 × 4历史时步 + 32维时间编码"] --> Patch["Patch Tokenizer<br/>跨步卷积投影 (P=2 / P=1, D=1536)"]
Patch --> Design1["恒定通道 Swin U-Net 与无掩码环面拓扑<br/>[2,8,4,8,2] 对称深度 / 恒定维度 D=1536 / 2D RoPE"]
Design1 --> Design2["现代 LLM 优化与正则化配方<br/>SwiGLU 前馈网络 / Muon 矩阵正交更新 / 强正则化"]
Design2 --> Design3["低成本 CRPS 概率回溯与超参副产物集成<br/>Fair-CRPS 准则 + MC Dropout + RFT 多学习率副产物集成"]
Design3 --> Output["高精度中期预报输出<br/>确定性残差更新 / 1-10天概率集合预报"]
关键设计¶
1. 恒定通道 Swin U-Net 骨干与无掩码环面拓扑:剔除领域定制偏置,释放硬件吞吐 针对以往气象模型在编码高低空变量时设计复杂非对称多分支结构、使用球面 GNN 或引入极地柱面注意力掩码导致硬件延迟剧增的痛点,Otter Weather 坚持最精简的二维视觉骨干。模型采用分层 Swin Transformer 构建对称 U-Net,下采样与上采样共 5 个阶段,各阶段深度配置为 \([2, 8, 4, 8, 2]\)。与常规随分辨率逐级加倍通道的做法不同,Otter 在整个 U-Net 内部维持完全恒定的 token 通道维度 \(D=1536\)(Base 配置),显著简化特征拼接与跳跃连接卷积。在空间建模上,窗口大小设为 \(15 \times 15\),交替采用标准窗口注意力(W-MSA)与移位窗口注意力(SW-MSA)。尤为激进的设计在于:完全移除地球两极的注意力掩码,将球面自然视作环面(torus)拓扑。消融证明,保留柱面掩码仅带来 0.17% 的边际增益,却引入额外 7% 计算开销并打碎标准高效注意力算子;配合引入的 2D 旋转位置编码(RoPE),网络可在极高计算吞吐下自适应捕获准确的球面相对几何关系。
2. 现代 LLM 优化与正则化配方:SwiGLU 激活、Muon 优化器与多重泛化约束 为了在极端压缩的计算预算(单卡几十小时)内实现与大模型可比的表征能力并避免过拟合,模型系统性吸纳了现代大语言模型的核心训练技术。前馈层全面舍弃标准 GeLU-MLP,改用带门控线性单元的 SwiGLU 结构: $\(\text{SwiGLU}(x) = (\text{SiLU}(x W_1) \otimes x W_2) W_3\)$ 为了在保持参数量与标准 MLP(4 倍扩展)相当的同时实现 Tensor Core 最优对齐,隐藏层扩展比缩减至 \(8/3 \approx 2.67\times\),并严格对齐至 128 的整数倍。在优化器层面,放弃传统的 AdamW,全面采用新兴的 Muon 优化器。Muon 将二维权重张量视作矩阵,通过两边动量变换和 Newton-Schulz 迭代实现参数更新的正交化,极大地拉平了不同特征方向上的学习速度,在小 epoch 条件下展现出显著快于 AdamW 的收敛效率。为抵消 Muon 在小批量下的矩阵运算开销,采用 4 步梯度累积(等效 batch size 为 16)。此外,配合高达 0.15 的权重衰减(Weight Decay)、\(p=0.10\) 的 Dropout 以及按 mini-batch 跳过残差块的随机深度(DropPath \(0.10\)),构建了极强且硬件无损的正则化屏障。
3. 低成本 CRPS 概率回溯与超参副产物深度集成:告别昂贵扩散去噪与从头重复训练 传统集合天气预报要么依赖从头独立多随机种子训练(单个模型耗费 \(>20\) A100-天),要么依赖 GenCast 等每步数百次去噪采样的扩散模型。Otter 提出极度高效的概率建模范式:直接复用确定性预训练权重,使用无偏的连续排位概率评分(Fair-CRPS)作为损失函数: $\(\text{CRPS}_{\text{fair}} = \frac{1}{M}\sum_{i=1}^{M}|X_{i}-y| - \frac{1}{2M(M-1)}\sum_{i=1}^{M}\sum_{j=1}^{M}|X_{i}-X_{j}|\)$ 在训练中仅需 \(M=2\) 个样本计算梯度,前一项约束预报准确性,后一项惩罚集合成员间的过聚拢。为了在不引入新参数的前提下注入随机性,Otter 舍弃了常规的自适应层归一化(AdaLN)噪声注入,转而直接在推理期激活模型自带的 MC Dropout(\(p=0.10\)),在测试时实现更高 CRPS 指标与极快的非迭代单次前向采样。而在深度集成构建上,作者发现无需重新耗费算力多轮独立训练,只需在自回归外推微调(Rollout Fine-Tuning, RFT)阶段收集不同学习率探索(\(1\times 10^{-6}, 2\times 10^{-6}, 5\times 10^{-6}\))下的调参副产物 Checkpoint 进行组合,即可以“零额外训练代价”获得泛化与离散度高度优秀的深度集成模型。
损失函数 / 训练策略¶
确定性阶段采用纬度加权与气压层加权的均方根误差(RMSE)损失函数: $\(\mathcal{L} = \sum_{v}\frac{1}{T}\sum_{t=1}^{T}\sqrt{\frac{1}{HW}\sum_{i=1}^{H}\sum_{j=1}^{W}w(\phi_{i})\bigl(\hat{X}_{i,j,v}^{t}-X_{i,j,v}^{t}\bigr)^{2}}\)$ 其中纬度权重 \(w(\phi_i) = \cos\phi_i / (\frac{1}{H}\sum_k \cos\phi_k)\) 用于校正经纬网格在极地的面积投影畸变,地表 2m 温度赋予高权重 10,其余变量赋予权重 1。训练分为两阶段:第一阶段单步预训练 20 个 epoch,学习率经 500 步线性预热至 \(2\times 10^{-4}\) 后余弦衰退至 \(2\times 10^{-6}\);第二阶段开展 RFT 微调,限制在 2007 年后数据上训练 1 个 epoch(约 19,000 步),采用渐进课程学习策略,预测步长从 24 小时(4 步)分 8 个阶段逐步放宽至 72 小时(12 步)。概率阶段则以单步 Fair-CRPS 微调 5 个 epoch,随后同样进行多步自回归 RFT 微调。
实验关键数据¶
主实验¶
在 ERA5 官方评测基准(WeatherBench 2 协议,\(1.5^\circ\) 分辨率)上,涵盖 Z500、Q700、T850、U850、V850、T2m、SP、U10m 和 V10m 九大核心变量。确定性指标汇报 24h-72h 聚合或 24h 相对 ECMWF 业务数值预报 IFS HRES 的 RMSE 改善率(%),概率指标汇报 1-10 天平均相对 IFS ENS 的 CRPS 改善率(%)。
| 模型 | 架构类型 | 空间分辨率 | 参数量 | 训练成本 (A100-天) | 确定性技能增益 (vs IFS HRES) |
|---|---|---|---|---|---|
| IFS HRES | 传统业务数值模式 (NWP) | 0.1° | – | 庞大超算集群 | 0.00% |
| Pangu-Weather | 3D 地球专用 Transformer | 0.25° | 256M | 1440 | +5.84% |
| FuXi | 级联 U-Transformer | 0.25° | 1.5B | 26 | +14.44% |
| NeuralGCM | 混合微分动力学 + 神经网络 | 0.7° | 31.1M | 8064 | +14.61% |
| GraphCast | 多网格复杂 GNN | 0.25° | 36.7M | 1344 | +16.32% |
| Keisler GNN | 二十面体 GNN | 1.0° | 6.7M | 5.5 | -15.27% |
| Stormer ENS | Vision Transformer (集合均值) | 1.5° | 300M | 128 | +6.11% |
| ArchesWeather-M×4 | Swin U-Net (4模型集合) | 1.5° | 84M | 20 | +8.86% |
| Otter Weather (Ours) | 恒定通道二维 Swin U-Net | 1.5° | 1.5B | 3.0 | +11.54% (24h: +9.6%) |
| Otter-XL (Ours) | 恒定通道二维 Swin U-Net (P=1) | 1.5° | 1.5B | 10.0 | +16.12% (24h: +14.9%) |
在概率预报设定下,单模型 Otter CRPS 仅需不到 8 A100-天即实现相比 IFS ENS +6.1% 的 CRPS 增益(SSR 为 0.94);而更大算力的 Otter-XL (Deep Ens) 耗费 29 A100-天达到 +9.7% 的 CRPS 改善与 0.95 的离散度-技能比(SSR),在耗费减少一个数量级的前提下以 2.2% 的明显优势超越前沿扩散模型 GenCast。
消融实验¶
基于非 RFT 基础配置在 24 小时预报时效下的相对 RMSE 性能变化及计算相对开销(以 Base 模型为基准 1.00×):
| 配置 / 变体 | 相对计算成本 | 24h RMSE 改善率 (%) | 说明 |
|---|---|---|---|
| Base 模型 (D=1536, [2,8,4], SwiGLU, 2D RoPE, Muon) | 1.00× | 0.00% | 基准配置 (非 RFT 状态) |
| 改变骨干结构: 浅而宽 D=2304, [1,4,1] | 1.00× | -0.05% | 同等算力下浅层网络表现更差 |
| 改变骨干结构: 深而窄 D=1280, [3,12,8] | 1.13× | +0.17% | 边际收益极小且增加耗时 |
| 移除 Dropout (Dropout = 0.0) | 1.00× | -17.61% | 出现灾难性过拟合,泛化骤降 |
| 移除 DropPath (DropPath = 0.0) | 1.09× | -0.96% | 稳定性与速度均受损 |
| 提高正则化 (Dropout=0.2, DropPath=0.2) | 0.93× | -0.51% | 过度约束导致欠拟合 |
| 缩小注意力窗口 (W=5 vs W=15) | 1.02× | -1.03% | 感受野不足影响大尺度气象系统 |
| 缩小 Patch 尺寸 (P=1, Otter-XL) | 3.25× | +5.28% | 密集 Token 显著拉升表征能力 |
| 激活函数替换: SwiGLU → GeLU | 1.00× | -2.28% | 门控机制对大模型特征筛选至关重要 |
| 增加极地注意力掩码 (Cylinder Mask) | 1.07× | +0.17% | 收益微弱且破坏算子优化,算力性价比低 |
| 位置编码替换: 移除 RoPE 仅留 Spherical | 0.99× | -1.18% | 证实 2D RoPE 的相对位置感知不可或缺 |
| 优化器替换: Muon (20 ep) → AdamW (23 ep) | 1.00× | -0.24% | 相同算力下 AdamW 收敛明显滞后 |
| 优化器替换: Muon (20 ep) → AdamW (35 ep) | 1.58× | +0.14% | AdamW 需多耗费 58% 算力才能追平 |
| 输入历史减少: 4 时步 → 2 时步 | 1.00× | -0.65% | 损失高阶时序导数与动态趋势信息 |
关键发现¶
- Dropout 是极低算力气象预报的生命线:完全移除 Dropout 会导致预报误差剧烈反弹 17.61%,表明在多变量气象场高维回归中,模型极易陷入局部特征拟合;强正则化是轻量模型取得 SOTA 的隐形基石。
- Muon 优化器在快速迭代中碾压 AdamW:Muon 的矩阵正交更新在有限步数内效率极其惊艳。AdamW 需要多花费近 60% 的训练周期才能勉强追平 Muon 在 20 个 epoch 建立的精度优势。
- 物理先验并非不可妥协:极地注意力掩码、专门的球面谐波位置编码与 GNN 邻域连接在实验中均未能展现出超越标准通用算子的显著优势;标准环面拓扑配合 2D RoPE 即可完美胜任。
- 跨动力系统泛化:在 The Well 基准的声学散射偏微分方程(PDE)任务中,Otter 零样本开箱即用,单步 VRMSE 达 0.0030,大幅超越 SOTA 物理基础模型 Walrus(0.0089),证明该范式具备跨领域物理时空模拟的通用潜力。
亮点与洞察¶
- 硬件友好的帕累托跃迁:拒绝为了追求物理归纳偏置而编写非标定制算子,坚持完全对齐现代 GPU Tensor Core 的通用结构,仅用 3.5 A100-天即跨越数值气象预报门槛,为中小型实验室提供了自主从头训练高精度气象大模型的范本。
- “零成本”深度集成策略:敏锐地指出研究人员在进行 RFT 学习率调参时本就会产生多个 Checkpoint,将这些调参中间产物组合直接构建深度集成,规避了反复多次独立训练的高昂开销。
- 通用时空建模的横向迁移性:同一套骨干和优化超参在声学散射波动方程上同样斩获超越前沿基础模型的战果,为流体力学及波动物理领域的轻量化神经算子设计指明了新路径。
局限与展望¶
- 空间分辨率受限于 1.5°:当前为验证架构与极速迭代,主要在 1.5° 分辨率下开展实验,尚未在高分辨率(如 0.25° 或更细网格)上进行千卡级全量对齐验证,极高分辨率下的显存与通道伸缩规律尚待确立。
- MC Dropout 存在轻度欠离散(Underdispersion):尽管概率模型 CRPS 表现极佳,但在短时效下集合的离散度-技能比(SSR)略低于理想值 1.0,集合预报的校准曲线仍有提升空间。
- 深度集成最佳分支节点未完全闭环:当前深度集成主要在 RFT 阶段分流,在预训练初期分支还是自回归微调期分支对校准与多样性的终极影响仍需系统性理论论证。
相关工作与启发¶
- vs GraphCast: GraphCast 采用多面体网格与深层消息传递 GNN,需要 1344 A100-天;Otter 采用标准二维 Swin Transformer 配合 Tensor Core 深度优化,在 1.5° 下以 1/130 的算力达到相当的预测水准。
- vs ArchesWeather / U-Cast: 同为高效天气预报模型,ArchesWeather 4 成员集合需 20 A100-天;Otter 通过引入 Muon 优化器与 2D RoPE,在 3.5 A100-天内便取得了超越其 4 成员集合的预测技巧。
- vs GenCast: GenCast 依赖高成本扩散去噪生成集合预报;Otter CRPS 采用确定性权重单向微调 + MC Dropout 单步采样,推理极快,且 Otter-XL (Deep Ens) 仅用 1/10 算力实现更优的 CRPS。
评分¶
- 新颖性: ⭐⭐⭐⭐ [打破领域几何偏置执念,系统引入 RoPE、SwiGLU、Muon 与零成本调参深度集成,构建新型高效气象建模范式]
- 实验充分度: ⭐⭐⭐⭐⭐ [完备的消融分析、多时效 RMSE/CRPS 评测、细致的概率标定考察,且包含声学 PDE 跨领域验证]
- 写作质量: ⭐⭐⭐⭐⭐ [逻辑严谨清晰,问题定义与计算成本透明,工程与理论分析洞见深刻]
- 价值: ⭐⭐⭐⭐⭐ [极大降低中长期天气与物理时空预测门槛,有力推动了气象 AI 的普及化与科研主权]