[Paper Review] Climate Modeling with Neural Diffusion Equations

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

Model Reduction and Neural Networks56 references4 citations

TL;DR

This paper proposes Neural Diffusion Equations (NDE), a novel climate modeling framework that integrates the physical diffusion equation with neural ordinary differential equations (NODEs) to learn latent governing dynamics from climate data. By combining a differentiable diffusion process with a learnable uncertainty model and trainable heat capacity, NDE achieves state-of-the-art forecasting accuracy across synthetic and real-world datasets with fewer parameters and higher efficiency than existing baselines.

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.

Motivation & Objective

• To address the limitations of existing deep learning-based climate models that rely heavily on RNNs or GNNs without strong physical inductive biases.
• To improve climate factor forecasting by embedding the physically grounded diffusion equation into a neural ODE framework.
• To model real-world uncertainties in climate systems through a learnable neural network component within the ODE framework.
• To optimize heat capacity parameters from data rather than assuming fixed or predefined values.
• To achieve high predictive accuracy with reduced model complexity and efficient inference.

Proposed method

• Proposes a Neural Diffusion Equation (NDE) that models climate dynamics as a time-continuous process governed by a modified diffusion equation: $ \frac{d\bm{h}(t)}{dt} = -k\bm{L}\bm{h}(t) + f(\bm{h}(t), t; \bm{\theta}_f) $.
• Uses the symmetrically normalized graph Laplacian $ \bm{L} $ to encode spatial connectivity between grid nodes or weather stations in a graph-structured climate network.
• Introduces a learnable neural network $ f $ to model physical uncertainties and non-idealities in real-world climate systems, trained via mean squared error loss.
• Treats the heat capacity coefficient $ k $ as a learnable parameter, with variants including scalar, edge-class, and single-coefficient learning strategies.
• Employs a graph-based architecture where node features (e.g., temperature, pressure) are evolved over time using the NDE layer, followed by a decoding head for multi-step predictions.
• Applies the neural ODE framework to solve the integral form of the dynamics: $ \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 $.

Experimental results

Research questions

• RQ1 Can a physics-informed neural ODE framework that combines the diffusion equation with learnable uncertainty outperform standard RNN and GNN baselines in climate forecasting?
• RQ2 How does the inclusion of a learnable heat capacity coefficient affect forecasting accuracy and model generalization?
• RQ3 What is the impact of different uncertainty modeling strategies (e.g., with or without the neural network $ f $) on prediction performance?
• RQ4 How sensitive is the model’s performance to hyperparameters such as input history length $ P $, hidden dimension $ D $, and heat capacity initialization method?
• RQ5 Does the NDE framework achieve better efficiency (in terms of parameters and inference cost) while maintaining or improving accuracy?

Key findings

• NDE achieves the lowest mean squared error (MSE) across all datasets, with a 1-step prediction MSE of 0.2975 ± 0.0062 on the NOAA dataset, outperforming all 11 baselines.
• The ablation study shows that removing the uncertainty model $ f $ (NDE without $ f $) increases error to 0.3245 ± 0.0161 on LA, confirming its critical role in modeling real-world noise.
• The edge-class heat capacity learning method yields the best performance (0.2621 ± 0.0026 MSE on LA), significantly outperforming fixed (0.3657 ± 0.0117) and single-coefficient (0.2901 ± 0.0096) methods.
• The model achieves optimal performance at a hidden dimension $ D = 32 $, with a 25% improvement in MSE on the NOAA dataset compared to $ D = 16 $.
• NDE achieves high efficiency, placing near the bottom-left corner in model efficiency plots, indicating low parameter count and low inference cost while maintaining high accuracy.
• Sensitivity analysis shows that only one time step of past history ($ P = 1 $) is sufficient for optimal performance, indicating that the diffusion process is memory-efficient and does not require long historical sequences.