🔬 physics

Stochastic and Non-local Closure Modeling for Nonlinear Dynamical Systems via Latent Score-based Generative Models

This paper proposes a latent score-based generative AI framework that jointly trains convolutional autoencoders with conditional diffusion models to efficiently learn stochastic, non-local closure laws for nonlinear dynamical systems, achieving significant computational acceleration while maintaining predictive accuracy comparable to standard physical-space diffusion models.

Original authors: Xinghao Dong, Huchen Yang, Jin-Long Wu

Published 2026-06-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xinghao Dong, Huchen Yang, Jin-Long Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Too Much Detail" Trap

Imagine you are trying to predict the weather for a whole city. To do this perfectly, you would need to track every single molecule of air, every drop of water, and every gust of wind. In the world of physics and engineering, this is called "resolving all scales."

The problem is that this is impossible for computers. It's like trying to count every grain of sand on a beach to predict the tide; it would take a supercomputer millions of years just to finish the calculation.

To solve this, scientists use "closure models." Think of a closure model as a shortcut. Instead of tracking every single grain of sand, you just track the big waves and use a rule of thumb to guess what the tiny ripples are doing.

The Old Way vs. The New Problem

For a long time, these shortcuts were deterministic (they always gave the exact same answer) and local (they only looked at what was happening right next to them).

The Flaw: Real-world systems (like turbulent water or air) are messy. They are full of randomness and long-distance connections (a ripple here might affect a wave miles away). The old shortcuts were too rigid. They couldn't capture the "chaos" or the "long-range gossip" of the system.

Recently, scientists started using AI to fix this. Specifically, they used a type of AI called a Diffusion Model.

The Analogy: Imagine a blurry photo that slowly gets clearer. A diffusion model works in reverse: it starts with pure static (noise) and slowly "denoises" it until it reveals a realistic picture of the missing physics. This is great because it captures randomness and long-range effects perfectly.

But there's a catch: This AI is incredibly slow. It has to take thousands of tiny steps to turn static into a picture. If you try to use this AI to fix a weather simulation step-by-step, your computer would freeze. It's like trying to paint a masterpiece by hand for every single frame of a movie; the movie would never finish.

The Solution: The "Latent Space" Shortcut

The authors of this paper proposed a clever two-step solution to make the AI fast and accurate.

1. The Compression (The Autoencoder)

First, they built a "compressor" (an Autoencoder).

The Analogy: Imagine you have a giant, high-definition 4K movie file. It's huge and hard to send. You compress it into a tiny, low-resolution thumbnail.
In the paper: They take the complex, high-resolution physics data and squish it down into a tiny, low-dimensional "latent space" (a compressed summary). The AI doesn't need to work on the giant 4K file; it just works on the tiny thumbnail.

2. The Generator (The Diffusion Model)

Next, they trained the Diffusion Model to work only on these tiny thumbnails.

The Analogy: Instead of painting the whole 4K movie frame by frame, the AI just paints the tiny thumbnail. Once the thumbnail is perfect, they "uncompress" it back into the big 4K movie.
The Result: Because the AI is working on a tiny thumbnail, it is 10 times faster than working on the full movie, but it still produces a high-quality result.

The Secret Sauce: "Joint Training"

Here is where the paper makes its biggest breakthrough.

Usually, scientists do this in two separate phases:

Phase A: Train the compressor to make the best possible thumbnail.
Phase B: Train the AI to generate thumbnails based on that compressor.

The Problem: The best thumbnail for saving space isn't always the best thumbnail for generating new images. It's like taking a photo that looks great when printed on a billboard, but when you try to use it as a puzzle piece, the edges don't fit the AI's needs. The AI struggles to learn the rules because the "thumbnail" is too weirdly shaped for it.

The Innovation: The authors trained the compressor and the AI together (Joint Training).

The Analogy: Instead of the photographer and the painter working in separate rooms, they sit at the same table. The photographer learns to take photos that are specifically easy for the painter to work with, while the painter learns to paint those specific photos.
The Outcome: The "thumbnail" (latent space) is now perfectly tuned for both compression and generation. The AI learns the rules much faster and makes fewer mistakes.

The Results: Speed Without Sacrifice

The team tested this on a simulation of swirling fluid (turbulence).

Accuracy: The new method produced results just as accurate as the slow, full-size AI. It captured the randomness and the big-picture patterns perfectly.
Speed: Because they were working in the "compressed" space, they got a 10x speedup.
Uncertainty: Crucially, the model didn't just give one answer; it gave a range of possible answers (an ensemble). This allows scientists to say, "There is a 90% chance the storm will go this way," which is vital for safety and planning.

Summary

The paper introduces a new way to model complex, chaotic systems (like turbulence) by:

Compressing the data into a tiny summary.
Using a Diffusion AI to generate the missing physics within that summary.
Training the compressor and the AI together so they work in perfect harmony.

This allows scientists to run simulations that are 10 times faster without losing the accuracy needed to understand complex, chaotic real-world phenomena.

Technical Summary: Stochastic and Non-local Closure Modeling via Latent Score-Based Generative Models

Problem Statement
Computational mechanics often deals with complex nonlinear dynamical systems, such as turbulent flows, where explicitly resolving all relevant scales is computationally prohibitive. Classical closure modeling approaches (e.g., RANS, LES) rely on deterministic and local assumptions to approximate subgrid-scale phenomena. However, these assumptions are often too restrictive for systems lacking clear scale separation, failing to capture uncertainties, regime transitions, and non-local spatial correlations. While recent advances in generative AI, specifically diffusion models, offer a promising framework for stochastic and non-local closure modeling by treating closure terms as random objects, their practical application is limited by prohibitive computational inference costs. Standard diffusion models require hundreds to thousands of network evaluations per time step, making them infeasible for integration into numerical simulations where closure terms must be evaluated at every step.

Methodology
The authors propose a Latent Conditional Diffusion Model (L-CDM) framework to address the trade-off between computational efficiency and generative accuracy. The methodology consists of three core components:

Latent Space Formulation: The framework employs convolutional autoencoders to compress high-dimensional physical fields (the system state $V$ and the closure term $U$ ) into low-dimensional latent representations ( $z_V$ and $z_U$ ). This dimensionality reduction is critical for reducing the computational cost of the reverse sampling process.
Non-local Modeling: To capture long-range spatial dependencies inherent in complex systems, the score function $s_\theta$ within the latent space is constructed using a Fourier Neural Operator (FNO). The FNO performs learned convolutions in the Fourier domain, ensuring global receptive fields. Additionally, the autoencoder architecture incorporates deep hierarchical layers and explicit self-attention modules to further model non-local interactions.
Joint Training Strategy: A key methodological innovation is the shift from the conventional two-phase training pipeline (training autoencoders and diffusion models separately) to an end-to-end joint optimization scheme.
- Conventional Approach: Autoencoders are trained solely to minimize reconstruction error, often resulting in a latent space that is difficult for diffusion models to sample from effectively.
- Proposed Approach: The autoencoder and the latent diffusion model are trained simultaneously using a composite loss function: $L_{joint} = L_{AE} + \lambda_{score} L_{score} + \lambda_{KL} L_{KL}$ $L_{j o in t} = L_{A E} + λ_{scor e} L_{scor e} + λ_{K L} L_{K L}$ .
  - $L_{AE}$ ensures faithful reconstruction of physical fields.
  - $L_{score}$ trains the diffusion model to match the conditional distribution $p(z_U | z_V)$ .
  - $L_{KL}$ (Kullback-Leibler divergence) acts as a regularizer to prevent latent space collapse, ensuring the latent distribution remains sufficiently stochastic and non-degenerate.

The sampling process utilizes an adaptive time-stepping scheme (Euler-Maruyama integration) to further accelerate the reverse diffusion process within the latent space.

Key Contributions

Latent Diffusion Framework: The paper introduces a stochastic closure modeling framework that maps high-dimensional fields to a latent manifold, performing conditional Langevin sampling there to achieve orders-of-magnitude faster inference compared to physical-space diffusion models.
Joint Optimization: The authors demonstrate that the standard two-phase training pipeline compromises the accuracy of closure-term generation. They propose a joint training scheme that harmonizes reconstruction and score-matching objectives, significantly improving generative accuracy while retaining the efficiency of latent-space sampling.
Non-local Stochastic Modeling: By integrating FNOs and attention mechanisms within a latent diffusion architecture, the framework effectively captures non-local spatial dependencies and stochastic behaviors that deterministic local models miss.

Numerical Results
The framework was evaluated on a two-dimensional Kolmogorov flow governed by the stochastic Navier-Stokes equations, where the goal was to model the unknown stochastic nonlinear convection term.

Accuracy vs. Efficiency:
- Physical-Space Baseline (P-CDM): Achieved high accuracy but required ~2.64 seconds to generate 1,000 samples.
- Two-Phase Latent Model: Achieved a 5.5× speedup (~0.48 seconds) but suffered from significant accuracy degradation (higher MSE and relative Frobenius norm errors) compared to the physical-space baseline.
- Jointly Trained Latent Model (Joint L-CDM): Achieved a ~~5.4× speedup (~~0.49 seconds) while maintaining predictive accuracy comparable to the physical-space baseline. The joint training corrected the distributional divergence observed in the two-phase approach.
Ensemble Simulations: When integrated into numerical solvers for 20-second simulations, the Joint L-CDM enabled ensemble generation with a ~10× computational acceleration compared to the P-CDM (1,236s vs. 8,652s) while maintaining relative errors below 10% for ensemble simulations.
Uncertainty Quantification: The model successfully reproduced not only the mean resolved dynamics but also the spatial structure and magnitude of the intrinsic uncertainty (standard deviation) of the closure term, outperforming deterministic approaches and the two-phase latent model in capturing physical variance.
Spectral Analysis: Energy spectrum analysis showed that the Joint L-CDM closely matched the ground truth up to wavenumbers of order $10^2$ , preserving the forward enstrophy cascade characteristic of 2D turbulence.

Significance and Claims
The paper claims that this work bridges the gap between the representational expressiveness of generative AI and the computational feasibility required for practical scientific computing. By jointly optimizing the latent representation and the generative model, the authors demonstrate that it is possible to achieve orders-of-magnitude computational acceleration in stochastic closure modeling without sacrificing predictive accuracy. This makes practical, large-scale ensemble-based uncertainty quantification feasible for complex dynamical systems where it was previously prohibitively expensive. The authors note that while the framework is architecturally extensible to 3D systems, the primary challenge remains the computational cost of high-fidelity 3D data, further underscoring the necessity of the latent-space approach. They also acknowledge that sensitivity to out-of-distribution data (e.g., extreme events) remains an area for future investigation.