[Paper Review] Capturing missing physics in climate model parameterizations using neural differential equations

Ali J. Ramadhan, John C. Marshall+9 more|arXiv (Cornell University)|Oct 23, 2020

Meteorological Phenomena and Simulations17 citations

TL;DR

This paper proposes using neural differential equations (NDEs) to improve climate model parameterizations by learning missing physics from high-resolution simulations. The NDE framework enhances a base parameterization (e.g., convective adjustment) with data-driven flux corrections, achieving superior accuracy and stability over traditional models like KPP, particularly in capturing nonlocal and up-gradient fluxes.

ABSTRACT

We explore how neural differential equations (NDEs) may be trained on highly resolved fluid-dynamical models of unresolved scales providing an ideal framework for data-driven parameterizations in climate models. NDEs overcome some of the limitations of traditional neural networks (NNs) in fluid dynamical applications in that they can readily incorporate conservation laws and boundary conditions and are stable when integrated over time. We advocate a method that employs a 'residual' approach, in which the NN is used to improve upon an existing parameterization through the representation of residual fluxes which are not captured by the base parameterization. This reduces the amount of training required and providing a method for capturing up-gradient and nonlocal fluxes. As an illustrative example, we consider the parameterization of free convection of the oceanic boundary layer triggered by buoyancy loss at the surface. We demonstrate that a simple parameterization of the process - convective adjustment - can be improved upon by training a NDE against highly resolved explicit models, to capture entrainment fluxes at the base of the well-mixed layer, fluxes that convective adjustment itself cannot represent. The augmented parameterization outperforms existing commonly used parameterizations such as the K-Profile Parameterization (KPP). We showcase that the NDE performs well independent of the time-stepper and that an online training approach using differentiable simulation via the Julia scientific machine learning software stack improves accuracy by an order-of-magnitude. We conclude that NDEs provide an exciting route forward to the development of representations of sub-grid-scale processes for climate science, opening up myriad new opportunities.

Motivation & Objective

• To address the persistent challenge of missing physics in climate model parameterizations, particularly in unresolved sub-grid-scale processes like oceanic boundary layer turbulence.
• To develop a data-driven parameterization framework that improves upon existing models by learning residual fluxes not captured by base parameterizations.
• To leverage neural differential equations (NDEs) to ensure conservation laws and long-term stability in fluid dynamical simulations.
• To demonstrate that NDEs can capture nonlocal and up-gradient fluxes—phenomena often missed by traditional parameterizations like convective adjustment or KPP.
• To validate the method using a realistic oceanic free convection case, showing improved fidelity over established parameterizations.

Proposed method

• A residual parameterization approach is used, where a neural network models the flux corrections (residuals) that are not represented by an existing base parameterization, such as convective adjustment.
• Neural differential equations (NDEs) are trained on high-resolution explicit simulations of the oceanic boundary layer to learn the dynamics of unresolved processes like entrainment fluxes at the mixed layer base.
• The NDE framework incorporates conservation laws and boundary conditions inherently through its differential equation formulation, ensuring physical consistency and numerical stability over long integrations.
• Online training via differentiable simulation using the Julia scientific machine learning stack (e.g., DiffEqFlux.jl) enables end-to-end backpropagation through the time integration, improving training accuracy by an order of magnitude.
• The method is tested with multiple neural network architectures (dense and convolutional) and evaluated using two loss functions: one based on fluxes (L1) and one on time series (L2), with training conducted on 9 high-resolution simulations.
• The K-Profile Parameterization (KPP) is used as a baseline, and its parameters are optimized via differential evolution to establish a fair comparison with the NDE-enhanced model.

Experimental results

Research questions

• RQ1 Can neural differential equations (NDEs) effectively learn and represent missing sub-grid scale fluxes—particularly nonlocal and up-gradient fluxes—beyond what is captured by standard parameterizations like convective adjustment?
• RQ2 How does the performance of an NDE-augmented parameterization compare to established models such as KPP in simulating oceanic free convection?
• RQ3 To what extent does the NDE framework maintain conservation laws and numerical stability over long-time integrations compared to traditional neural networks?
• RQ4 Does online training via differentiable simulation significantly improve the accuracy of the NDE-based parameterization compared to standard training?
• RQ5 Can the NDE-based parameterization generalize across different initial conditions and remain robust to changes in time-stepping schemes?

Key findings

• The NDE-augmented parameterization outperforms the standard KPP model, reducing the loss function from 1.26×10⁻⁴ (with reference parameters) to 5.33×10⁻⁵ after optimization, representing a 2.4× improvement.
• The NDE framework successfully captures entrainment fluxes at the base of the well-mixed layer—processes that convective adjustment alone cannot represent—demonstrating its ability to model nonlocal and up-gradient fluxes.
• The NDE model achieves a 10× improvement in accuracy when trained using online differentiable simulation via the Julia stack, compared to standard training, highlighting the benefit of end-to-end backpropagation.
• The NDE-based parameterization is stable over long integrations and independent of the time-stepping scheme, a key advantage over traditional neural networks that often suffer from instability.
• The method maintains physical consistency by incorporating conservation laws and boundary conditions directly into the differential equation structure, ensuring physically plausible behavior.
• The NDE framework generalizes well across different training and test conditions, with multiple neural network architectures (dense and convolutional) showing consistent performance, indicating robustness to architectural choices.