Skip to content

ClimODE: Climate and Weather Forecasting with Physics-informed Neural ODEs

Status: completed

Authors: Yogesh Verma, Markus Heinonen, Vikas Garg

Venue / Year: ICLR 2024 (Oral)

Affiliations: Aalto University, YaiYai Ltd, MIT CSAIL

Links: arXiv | OpenReview | GitHub | Project

Tags: [[physics-informed-ml]] [[neural-ode]] [[weather-forecasting]] [[climate-modeling]] [[continuity-equation]] [[uncertainty-quantification]]

One-Sentence Summary

ClimODE 将天气建模为连续时间平流 PDE(continuity equation),通过二阶 Neural ODE 学习全球天气输运的速度场,以仅 2.8M 参数在 ERA5 全球与区域预报中超越 ClimaX、FourCastNet 等数据驱动方法,并附带不确定性估计。

Problem

Core Challenge

气候与天气预报需要在长时间尺度上精确预测多个气象变量(温度、风速、位势等)的空间演化。核心挑战包括:

  • 物理守恒:大气输运遵循质量守恒(continuity equation),离散时间自回归模型天然违反该约束
  • 长时间稳定性:多步预报误差累积导致预报崩溃
  • 不确定性量化:确定性深度学习方法无法提供预测置信度
  • 计算效率:NWP(如 ECMWF IFS)需要超算资源,难以部署

Motivation

论文指出现有深度学习天气模型存在三个关键缺陷:

  1. 缺乏物理约束:Transformer/GNN 方法(GraphCast、Pangu-Weather、ClimaX)作为纯数据驱动黑箱,不尊重质量守恒
  2. 离散时间跳跃:自回归预测违反连续时间物理动力学,导致长 horizon 预报不稳定
  3. 无不确定性估计:所有主要竞品(FourCastNet、ClimaX、GraphCast、Pangu-Weather)均为点预测

动机论证具有说服力:将连续时间 PDE(continuity equation)作为归纳偏置是统计力学的自然选择,而非追热点。

Method

Core Innovation

ClimODE 的核心思想是将天气演化建模为受连续性方程约束的 Neural ODE:

  1. 平流方程:天气量 uku_k 的变化由速度场 vkv_k 驱动,保证质量守恒
  2. 二阶流:参数化速度变化率 v˙\dot{v} 而非速度本身,提升表达能力
  3. 发射模型:高斯发射输出偏差 μ\mu 和方差 σ2\sigma^2,同时解决源项和不确定性

Architecture

ComponentFunctionKey Design
Convolution Network fconvf_{\mathrm{conv}}局部天气输运ResNet with 3x3 conv, [5,3,2] residual blocks, [128,64,out] hidden dims
Attention Conv Network fattf_{\mathrm{att}}全球远程关联KQV dot-product attention, K/V stride=2, Q stride=1
Flow Velocity vkv_k二阶 ODE 状态每个气象量独立的 2D 速度场
Emission Model gg源项与不确定性[3,2,2] residual blocks, 输出 μk\mu_kσk\sigma_k
Initial Velocity Inference初速度估计最小二乘 + 高斯 RBF 先验,Adam lr=2, 200 epochs
Spatiotemporal Embedding ψ\psi时间周期 + 空间位置三角函数编码(日、年、经纬度)

Key Equations

连续性方程(核心动力学):

u˙k(x,t)=vk(x,t)uk(x,t)uk(x,t)vk(x,t) \dot{u}_k(x,t) = -v_k(x,t) \cdot \nabla u_k(x,t) - u_k(x,t) \nabla \cdot v_k(x,t)

其中第一项为 transport(量的空间移动),第二项为 compression(流的可压缩性)。

二阶流(速度加速度):

v˙k(x,t)=fθ(u(t),u(t),v(t),ψ) \dot{v}_k(x,t) = f_\theta\big(u(t), \nabla u(t), v(t), \psi\big)

混合网络(局部 + 全局):

fθ=fconv+γfatt f_\theta = f_{\mathrm{conv}} + \gamma f_{\mathrm{att}}

其中 γ\gamma 为可学习标量。

一阶 ODE 系统(Method of Lines):

[u(t)v(t)]=[u(t0)v(t0)]+t0t[(ukvk)fθ()]dτ \begin{bmatrix} u(t) \\ v(t) \end{bmatrix} = \begin{bmatrix} u(t_0) \\ v(t_0) \end{bmatrix} + \int_{t_0}^{t} \begin{bmatrix} -\nabla \cdot (u_k v_k) \\ f_\theta(\cdots) \end{bmatrix} d\tau

发射模型(不确定性):

ukobs(x,t)N(uk(x,t)+μk(x,t),  σk2(x,t)) u_k^{\mathrm{obs}}(x,t) \sim \mathcal{N}\big(u_k(x,t) + \mu_k(x,t),\; \sigma_k^2(x,t)\big)

损失函数(负对数似然):

L(θ;D)=1NKHWi=1N[logN(yiu(ti)+μ(ti),  diagσ2(ti))+logN(σ(ti)0,  λσ2I)] \mathcal{L}(\theta; \mathcal{D}) = -\frac{1}{NKHW} \sum_{i=1}^{N} \left[ \log \mathcal{N}\big(y_i \mid u(t_i) + \mu(t_i),\; \mathrm{diag}\,\sigma^2(t_i)\big) + \log \mathcal{N}\big(\sigma(t_i) \mid 0,\; \lambda_\sigma^2 I\big) \right]

其中 λσ1\lambda_\sigma^{-1} 通过 cosine annealing 衰减以移除正则化效应。

初速度推断(预处理):

v^k(t)=argminvk(t)u˙~k(t)+vk(t)u~k(t)+uk(t)vk(x,t)22+αvk(t)K2 \hat{v}_k(t) = \arg\min_{v_k(t)} \left\| \tilde{\dot{u}}_k(t) + v_k(t) \cdot \nabla \tilde{u}_k(t) + u_k(t) \nabla \cdot v_k(x,t) \right\|_2^2 + \alpha \|v_k(t)\|_K^2

Architecture Deep Dive

整体 Pipeline

  1. 输入:当前状态 u(t0)RK×H×Wu(t_0) \in \mathbb{R}^{K \times H \times W}K=5K=5 气象量,H=32,W=64H=32, W=64
  2. 初速度推断:利用过去 3 帧通过最小二乘估计 vk(t0)v_k(t_0)
  3. ODE 前向求解:将 [u;v;ψ]R(3K+C)×H×W[u; v; \psi] \in \mathbb{R}^{(3K+C) \times H \times W} 输入 fθf_\theta,用 Euler solver 以 1h 步长积分
  4. 发射输出gg 输出 μk,σk\mu_k, \sigma_k,构造高斯似然

Convolution Network fconvf_{\mathrm{conv}}

  • 结构:ResNet,[5,3,2] residual blocks,hidden dim [128,64,out_channels]
  • 卷积参数:kernel=3, stride=1, padding=1
  • 边界处理:X 轴 circular padding(国际日期变更线),Y 轴 reflection padding(极地)
  • 输入[u,u,v,ψ]R(5K+C)×H×W[u, \nabla u, v, \psi] \in \mathbb{R}^{(5K + C) \times H \times W}
  • 输出v˙kR2K×H×W\dot{v}_k \in \mathbb{R}^{2K \times H \times W}
  • Dropout:0.1

Attention Convolutional Network fattf_{\mathrm{att}}

  • Key/Value:2 层 CNN,stride=2,将 H×WH \times W 空间下采样 4 倍
  • Query:2 层 CNN,stride=1,保持空间分辨率
  • Attention 计算β=softmax(QK)V\beta = \mathrm{softmax}(QK^\top)V
  • 后处理1×11 \times 1 卷积映射回输出通道
  • 设计意图K/VK/V 下采样降低计算量,QQ 保持全分辨率

Emission Model gg

  • 结构:ResNet,[3,2,2] residual blocks,hidden dim [128,64,out_channels]
  • 输入[u(x,t),ψ(x,t)][u(x,t), \psi(x,t)]
  • 输出μk(x,t)\mu_k(x,t)(偏差)和 σk(x,t)\sigma_k(x,t)(不确定性)
  • 功能:捕获昼夜循环等源项,同时量化 aleatoric + epistemic 不确定性

时空嵌入 ψ(x,t)\psi(x,t)

  • 日周期sin(2πt),cos(2πt)\sin(2\pi t), \cos(2\pi t)
  • 年周期sin(2πt/365),cos(2πt/365)\sin(2\pi t / 365), \cos(2\pi t / 365)
  • 空间{sin,cos}×{h,w}\{\sin, \cos\} \times \{h, w\},球坐标 sin(h)cos(w),sin(h)sin(w)\sin(h)\cos(w), \sin(h)\sin(w)
  • 联合ψ(t)×ψ(x)\psi(t) \times \psi(x)(叉积)
  • 静态ψ(h),ψ(w),lsm,oro\psi(h), \psi(w), \mathrm{lsm}, \mathrm{oro}

算法流程

训练

  1. 预处理:对每个 t0t_0,利用 u(t02),u(t01),u(t0)u(t_0-2), u(t_0-1), u(t_0) 通过 cubic spline 拟合 u˙~k\tilde{\dot{u}}_k
  2. 初速度推断:Adam (lr=2) 优化 200 epochs 得到 v^k(t0)\hat{v}_k(t_0)
  3. ODE 前向:将 [u;v;ψ][u; v; \psi] batch 为单张图,Euler solver 1h 步长积分至目标时刻
  4. 发射输出:gg 输出 μ,σ\mu, \sigma
  5. 损失计算:负对数似然 + 方差高斯先验
  6. 反向传播:通过 ODE solver adjoint 方法回传梯度
  7. 优化器:Cosine Annealing LR scheduler,300 epochs,单 V100 GPU

推理

  • 与训练相同流程,但额外输出 σ\sigma 作为不确定性估计
  • CRPS 评估概率预测质量

Evidence

Claims → Evidence 映射

Claim类型支撑证据强度风险
ClimODE 超越 ClimaX/FCN/NODE 全球预报EmpiricalTable 6: RMSE/ACC 全变量全 lead-time 优于分辨率仅 5.625°,未与 Pangu/GraphCast 比较
2.8M 参数,一个量级更小EmpiricalTable 1: ClimaX 107M, Pangu 256M, GraphCast 37M参数量不直接等于推理效率
质量守恒约束有效EmpiricalAppendix H: ukdx\int u_k dx 恒定至 101210^{-12}仅验证了闭系统部分
不确定性估计有效EmpiricalFig 5: CRPS 随 lead-time 合理增长缺乏与其他不确定性方法的定量对比
区域预报有效EmpiricalTable 2: 三区域五变量全面优于仅三个区域,代表性有限
气候预报有效EmpiricalFig 5: 月度预报优于 FourCastNet仅与 FCN 对比,无 ClimaX
Advection 是最大贡献组件EmpiricalFig 7 消融:NODE → +Adv 提升最大消融在 unweighted RMSE 上进行
单 GPU 可训练EmpiricalAppendix D.4: 单 32GB V100仅限 5.625° 分辨率

Key Results

全球预报(Table 6):

VariableLead-TimeNODEClimaXFCNIFSClimODE
z24h877.8364.9333.051.0193.4$\pm$16.3
t24h3.352.171.830.871.55$\pm$0.18
t2m24h3.862.371.681.021.40$\pm$0.09
u1024h4.102.492.331.112.01$\pm$0.10
v1024h4.072.482.391.332.04$\pm$0.10

关键发现:ClimODE 在所有气象变量和 lead-time 上优于 NODE/ClimaX/FCN,但仍显著落后于 IFS(金标准)。在 t2m 上表现最接近 IFS。

消融实验

组件作用效果
NODE(baseline)自由形式二阶 ODE最差
+Adv(平流)物理约束最大提升
+Att(注意力)远程关联最小提升
+Emission源项+不确定性显著提升

Limitations

  1. 分辨率限制:仅在 5.625°(约 625km)上验证,远低于 GraphCast(0.25°)和 Pangu-Weather(0.25°),未证明在高分辨率下的可扩展性
  2. 关键竞品缺失:Pangu-Weather 和 GraphCast 因"代码不可用"而未对比,但两者均为同期或更早的强竞品
  3. 初速度推断瓶颈vk(t0)v_k(t_0) 的预处理需要 Adam 优化 200 epochs,显著增加推理延迟,论文未报告推理时间
  4. 源项建模简化:发射模型 gg 仅输出高斯偏差和方差,无法捕获降水等强非线性和间歇性过程
  5. 闭系统假设:核心平流方程假设质量守恒,但实际天气存在辐射源汇,仅通过发射模型事后补偿
  6. 气候预报对比不足:月度预报仅与 FourCastNet 对比,未与 ClimaX 等更强 baseline 比较
  7. 长期气候适应性:论文承认无法可靠预测气候变化情景下的天气模式
  8. IFS 差距显著:尽管超越所有 neural baselines,ClimODE 与 IFS 的 RMSE 差距仍然很大(如 z 24h: 193 vs 51)

Critical Assessment

Strengths

  • 物理归纳偏置清晰:将 continuity equation 嵌入 Neural ODE 的动机清晰、物理合理,不是简单的 "加个 loss 正则化"
  • 参数效率惊人:2.8M 参数 vs ClimaX 107M / Pangu 256M,在单 GPU 上训练,对计算资源有限的团队极具吸引力
  • 质量守恒验证严格ukdx\int u_k dx 恒定至 101210^{-12},是真正实现了约束而非近似
  • 不确定性估计:在深度天气预报中罕见地提供了 μ\muσ\sigma,且 μ\mu 物理上可解释(昼夜循环)
  • 消融实验充分:每个组件的贡献被清晰量化

Weaknesses & Risks

重要缺陷

  • 缺失关键 baseline:Pangu-Weather(Nature 2023)和 GraphCast(Science 2023)是同期最重要的竞品,以"代码不可用"为由跳过不够充分。这两者参数量虽大但在各自论文中展示了 SOTA 性能,ClimODE 的 "state of the art" 声明需要与它们直接对比
  • 初速度推断的推理成本:论文强调单 GPU 训练,但未报告推理时间。初速度推断需要对每个新输入运行 Adam 200 epochs,这在业务预报场景中可能不可接受
  • 分辨率差距:5.625° 分辨率在天气预报领域属于粗分辨率。ClimODE 的性能优势是否能保持到 1.5° 或 0.25° 是未解的关键问题

可补救问题

  • 消融实验中 NODE baseline 的 ResNet 结构与 ClimODE 不完全一致(NODE 用 [5,3,2],emission 用 [3,2,2]),控制变量不够严谨
  • Table 6 中 IFS 的多个 lead-time 标记为 N/A,影响完整性
  • 区域预报仅选三个区域(北美、南美、澳洲),未覆盖欧洲、亚洲、非洲

潜在风险

  • 闭系统假设可能在极端天气事件(如台风、热浪)中失效,因为源汇项变得主导
  • 发射模型假设高斯不确定性,但天气预报中不确定性通常是非高斯的(如降水的零膨胀分布)

Reviewer Feedback

SourceKey Points
ICLR 2024 Oral被选为 Oral,说明审稿人对物理约束 + 参数效率的认可度很高
OpenReview无法访问(403),具体审稿意见未知

Innovation Score: 7/10

理由

  • 将 continuity equation 嵌入 Neural ODE 是有概念深度的贡献,不是简单的架构堆叠
  • 2.8M 参数击败 107M 的 ClimaX 是一个令人印象深刻的结果
  • 但方法本身(Neural ODE + PDE constraint + attention + emission)的各个组件都是已知的,创新在于组合方式和物理动机
  • 不确定性估计是高斯发射模型,技术含量有限
  • 分辨率和 baseline 缺失限制了结论的可信度

Reusable Ideas

  1. PDE-constrained Neural ODE for conservation laws:对于任何满足守恒律的时空预测问题(流体力学、交通流、生态扩散),可以类似地将 PDE 嵌入 ODE 结构
  2. 初速度推断 as preprocessing:利用 PDE 本身的恒等式从观测量反推隐变量(如速度),避免引入额外编码器
  3. Method of Lines + batched ODE:将 PDE 离散为位置级 ODE 但共享同一 fθf_\theta,整个地球可 batch 为单张图求解
  4. Emission model for source terms:当核心动力学是闭系统但实际存在源汇时,用概率发射模型补偿,同时提供不确定性
  5. γ\gamma-gated hybrid networkf=flocal+γfglobalf = f_{\mathrm{local}} + \gamma f_{\mathrm{global}} 的门控混合是一种简洁的局部-全局融合方式

直接竞品

方法架构参数量优势vs ClimODE
ClimaX (ICML 2023)ViT107MFoundation model, 灵活输入ClimODE 更准、更小、有 UQ
FourCastNet (2022)AFNON/A高分辨率 (0.25°)ClimODE 更准但分辨率低
GraphCast (Science 2023)GNN37M0.25°, 多变量未直接对比
Pangu-Weather (Nature 2023)3D ViT256M0.25°, SOTA at submission未直接对比
NowcastNet (Nature 2023)Physics+GenN/A极端降水 nowcasting不同任务(precipitation vs general)

方法族

  • Neural ODEs:Chen et al. (2018) 的原始框架,ClimODE 使用其二阶扩展 + PDE 约束
  • Physics-Informed Neural Networks (PINNs):Raissi et al. (2019),但 PINNs 通常在 loss 中加入 PDE 残差,ClimODE 将 PDE 直接嵌入动力学结构
  • Hamiltonian/Lagrangian Neural Networks:Greydanus et al. (2019), Cranmer et al. (2020),类似物理约束但用于力学系统
  • Neural PDEs:FNO (Li et al. 2021), DeepONet (Lu et al. 2021) 学习 PDE 算子,但 ClimODE 学习 PDE 的时域演化

Personal Notes

ClimODE 在 AI for Science 领域的定位很清晰:不追求最大模型,而是用最小模型 + 正确物理实现最大效果。这在当前 LLM/大模型时代是一个有价值的立场。

值得关注的方向

  • 将 ClimODE 扩展到高分辨率(1.5° 或 0.25°),验证 scaling law
  • 替换高斯发射为非参数或混合分布,以捕获降水等间歇性变量
  • 初速度推断的加速(如用 amortized inference network 替代逐点优化)
  • 与 NowcastNet 的物理约束结合:NowcastNet 的 evolution operator 是显式 advection + residual,ClimODE 是连续时间 advection ODE,两者可以互补

跟进决策:值得参考其 PDE-constrained Neural ODE 的设计范式,但直接在高分辨率天气预报场景使用可能受限于分辨率和初速度推断效率。

科研品味三维评分

创新性 (Novelty):        ★★★★☆ (4/5)
  - 将 continuity equation 嵌入 Neural ODE 有概念深度,但各组件(Neural ODE、attention、emission model)均为已有技术

严谨性 (Rigor):         ★★★☆☆ (3/5)
  - 质量守恒验证严格,消融充分
  - 缺失 Pangu-Weather/GraphCast 对比严重削弱 "SOTA" 声明
  - 初速度推断的推理成本未报告

影响力 (Impact):        ★★★★☆ (4/5)
  - 2.8M 参数的效率极具吸引力,开源代码
  - ICLR Oral 说明社区认可度高
  - 但粗分辨率限制了实际业务应用

一句话总结

ClimODE 在 AI 天气预报方向通过 continuity equation + Neural ODE 实现了参数高效且物理一致的连续时间预报,相比 ClimaX/FCN 的优势是更小模型 + 质量守恒 + 不确定性,劣势是粗分辨率 + 缺失关键竞品对比。在 ICLR 2024 Oral 的标准下,这是一篇概念清晰但实验验证仍有提升空间的扎实工作。

行动建议

如果我是作者的 advisor,我会建议:
- [x] 接受发表(ICLR Oral 已证明质量)
- [x] 鼓励继续深化:扩展到高分辨率 + 与 Pangu/GraphCast 对比
- [ ] 缓一缓:如果无法补充关键 baseline 对比

Static research notes built with VitePress and KaTeX.