Data-Driven Stochastic Closure Modeling via Conditional Diffusion Model and Neural Operator

Introduction

Closure models are widely used in simulating complex multiscale dynamical systems such as turbulence and the earth system, for which direct numerical simulation that resolves all scales is often too expensive. For those systems without a clear scale separation, deterministic and local closure models usually lack enough generalization capability, which limits their performance in many real-world applications. In this work, we propose a data-driven modeling framework for constructing stochastic and non-local closure models via conditional diffusion model and neural operator. Specifically, the Fourier neural operator is incorporated into a score-based diffusion model, which serves as a data-driven stochastic closure model for complex dynamical systems governed by partial differential equations (PDEs). We also demonstrate how accelerated sampling methods can improve the efficiency of the data-driven stochastic closure model. The results show that the proposed methodology provides a systematic approach via generative machine learning techniques to construct data-driven stochastic closure models for multiscale dynamical systems with continuous spatiotemporal fields.

Closure in Complex Dynamical Systems

Consider a dynamical system whose general governing equation is given by: $\begin{array}{}\text{(1)}& \frac{\partial v}{\partial t}=M(v),\end{array}$ where $v$ represents the full system state and $M$ encompasses the dynamics derived from first principles, which can include highly nonlinear mappings, differential operators, and integral operators. In practice, resolving every detail of $v$ numerically is often impractical, which motivates the use of a reduced-order system: $\begin{array}{}\text{(2)}& \frac{\partial V}{\partial t}=\tilde{M}(V),\end{array}$ with $V:=\mathcal{K}(v)$ denoting a reduced-order representation of $v$, where the mapping $\mathcal{K}$ extracts the essential information for numerical simulation. Here, $\tilde{M}(V)$ is constructed based on domain knowledge to approximate the reduced dynamics $\mathcal{K}\circ M(v)$, resulting in a closed system in terms of $V$ that is tractable for practical computations.

To enhance the performance of traditional closure models $\tilde{M}$, researchers have recently introduced a machine-learning-based correction term $U$, aiming for the modified model $\tilde{M}(V)+U$ to more accurately capture the true reduced-order dynamics.

Conditional Diffusion Model Framework

In this work, we focus on a stochastic variant of the correction term $U$, enabled by recent advances in diffusion models. Diffusion models gradually add noise to data via a forward process, then learn to reverse the process to generate samples that mimic the original data distribution. In this work, we use the score-based diffusion model, which learns the gradient of the log probability density (the “score”) rather than the probability density itself. This formulation bypasses the need for a normalizing constant and makes the training process more tractable . Specifically, our goal is to approximate the distribution $p(U)$ and to sample from it efficiently. It is important to note that $U$ often depends on the current state $V$ as well as additional information—such as the temporal history of $V$ or spatially non-local data. Hence, we denote all such auxiliary information by $\mathrm{y}$ and concentrate on a conditional diffusion model framework targeting the approximation and efficient sampling of $p(U\mid \mathrm{y})$. By adopting this stochastic viewpoint, the model can naturally account for the uncertainty and variability inherent in unresolved scales. This probabilistic treatment is also especially important in systems where past states influence current dynamics—a phenomenon known as non-Markovian behavior.

Figure 1. Schematic plot of the entire workflow.

Leveraging Fourier Neural Operators

One of the major innovations in this work is the incorporation of Fourier Neural Operators (FNOs) to build the score function. Unlike traditional neural networks that operate on discretized data, FNOs work directly with continuous spatiotemporal fields. This confers two major advantages:

1. Resolution Invariance: FNOs allow the model to generalize across different grid resolutions. This is particularly useful when training data might be available at one resolution while simulations require finer scales

2. Non-Local Modeling: Because FNOs operate in the Fourier domain, they inherently capture global interactions within the field. This makes them well-suited to model non-local effects, which are critical when the closure term depends on long-range interactions in the system.

The neural operator architecture seamlessly integrates multiple data inputs—such as the current vorticity field and sparse measurements of the true closure term—into the score-based diffusion framework.

More to come.