[论文解读] Climate Modeling with Neural Diffusion Equations

JeeHyun Hwang, Jeongwhan Choi+4 more|arXiv (Cornell University)|Nov 11, 2021

Model Reduction and Neural Networks参考文献 56被引用 4

一句话总结

该论文提出神经扩散方程（NDE），一种新颖的气候建模框架，将物理扩散方程与神经常微分方程（NODEs）相结合，从气候数据中学习潜在的控制动力学。通过将可微分扩散过程与可学习的不确定性模型及可训练的热容量相结合，NDE 在合成数据集和真实世界数据集上均实现了最先进的预测精度，且参数更少、效率更高，优于现有基线方法。

ABSTRACT

Owing to the remarkable development of deep learning technology, there have been a series of efforts to build deep learning-based climate models. Whereas most of them utilize recurrent neural networks and/or graph neural networks, we design a novel climate model based on the two concepts, the neural ordinary differential equation (NODE) and the diffusion equation. Many physical processes involving a Brownian motion of particles can be described by the diffusion equation and as a result, it is widely used for modeling climate. On the other hand, neural ordinary differential equations (NODEs) are to learn a latent governing equation of ODE from data. In our presented method, we combine them into a single framework and propose a concept, called neural diffusion equation (NDE). Our NDE, equipped with the diffusion equation and one more additional neural network to model inherent uncertainty, can learn an appropriate latent governing equation that best describes a given climate dataset. In our experiments with two real-world and one synthetic datasets and eleven baselines, our method consistently outperforms existing baselines by non-trivial margins.

研究动机与目标

• 解决现有基于深度学习的气候模型严重依赖RNN或GNN但缺乏强物理归纳偏置的局限性。
• 通过将物理上合理的扩散方程嵌入神经ODE框架，提升气候因子预测性能。
• 通过在ODE框架中引入可学习的神经网络组件，对气候系统中的真实世界不确定性进行建模。
• 从数据中优化热容量参数，而非假设固定或预定义的值。
• 在降低模型复杂度的同时实现高预测精度和高效的推理。

提出的方法

• 提出一种神经扩散方程（NDE），将气候动力学建模为由修改后的扩散方程控制的连续时间过程：$ \frac{d\bm{h}(t)}{dt} = -k\bm{L}\bm{h}(t) + f(\bm{h}(t), t; \bm{\theta}_f) $。
• 使用对称归一化的图拉普拉斯矩阵 $ \bm{L} $ 来编码图结构气候网络中网格节点或气象站之间的空间连通性。
• 引入一个可学习的神经网络 $ f $ 以建模现实气候系统中的物理不确定性与非理想性，通过均方误差损失进行训练。
• 将热容量系数 $ k $ 视为可学习参数，其变体包括标量、边类别和单系数学习策略。
• 采用基于图的架构，其中节点特征（如温度、气压）通过NDE层随时间演化，随后通过解码头实现多步预测。
• 应用神经ODE框架求解动力学的积分形式：$ \bm{h}(t_1) = \bm{h}(t_0) + \int_{t_0}^{t_1} \left( -k\bm{L}\bm{h}(t) + f(\bm{h}(t), t; \bm{\theta}_f) \right) dt $。

实验结果

研究问题

• RQ1 结合扩散方程与可学习不确定性的物理信息神经ODE框架，是否能在气候预测中超越标准RNN和GNN基线？
• RQ2 可学习热容量系数的引入如何影响预测精度与模型泛化能力？
• RQ3 不同的不确定性建模策略（例如，有或无神经网络 $ f $）对预测性能有何影响？
• RQ4 模型性能对超参数（如输入历史长度 $ P $、隐藏维度 $ D $ 和热容量初始化方法）的敏感性如何？
• RQ5 NDE框架是否在保持或提升精度的同时，实现了更高的效率（参数量与推理成本）？

主要发现

• NDE在所有数据集上均达到最低的均方误差（MSE），在NOAA数据集上的1步预测MSE为0.2975 ± 0.0062，优于全部11个基线模型。
• 消融研究显示，若移除不确定性模型 $ f $（即NDE without $ f $），在LA数据集上的误差上升至0.3245 ± 0.0161，证实其在建模真实世界噪声中的关键作用。
• 边类别热容量学习方法表现最佳（LA数据集上MSE为0.2621 ± 0.0026），显著优于固定方法（0.3657 ± 0.0117）和单系数方法（0.2901 ± 0.0096）。
• 模型在隐藏维度 $ D = 32 $ 时达到最优性能，相较于 $ D = 16 $，NOAA数据集上的MSE降低25%。
• NDE实现了高效率，在模型效率图中位于左下角附近，表明其参数量少、推理成本低，同时保持高精度。
• 敏感性分析表明，仅需一个时间步的历史信息（$ P = 1 $）即可实现最优性能，表明扩散过程具有记忆高效性，无需长历史序列。