Synergizing Transport-Based Generative Models and Latent Geometry for Stochastic Closure Modeling

Xinghao Dong; Huchen Yang; Jin-long Wu

cs.LGFeb 19, 2026

Synergizing Transport-Based Generative Models and Latent Geometry for Stochastic Closure Modeling

Authors: Xinghao Dong, Huchen Yang, Jin-long Wu

Organizations: Department of Mechanical Engineering, University of Wisconsin–Madison, Madison, WI 53706

Abstract

Diffusion models recently developed for generative AI tasks can produce high-quality samples while still maintaining diversity among samples to promote mode coverage, providing a promising path for learning stochastic closure models. Compared to other types of generative AI models, such as GANs and VAEs, the sampling speed is known as a key disadvantage of diffusion models. By systematically comparing transport-based generative models on a numerical example of 2D Kolmogorov flows, we show that flow matching in a lower-dimensional latent space is suited for fast sampling of stochastic closure models, enabling single-step sampling that is up to two orders of magnitude faster than iterative diffusion-based approaches. To control the latent space distortion and thus ensure the physical fidelity of the sampled closure term, we compare the implicit regularization offered by a joint training scheme against two explicit regularizers: metric-preserving (MP) and geometry-aware (GA) constraints. Besides offering a faster sampling speed, both explicitly and implicitly regularized latent spaces inherit the key topological information from the lower-dimensional manifold of the original complex dynamical system, which enables the learning of stochastic closure models without demanding a huge amount of training data.

Figures & tables

Figure 1: Qualitative comparison of stochastic closure samples from physical-space models. This figure assesses the performance of conditional generation. Each column corresponds to a different model: the ground truth, P-DM, P-FM, and P-SI with two different priors. Each row displays an independent, random sample of the closure term $H$ , all generated for the same input vorticity field $\omega$ .

Sample steps	P-DM		P-FM		P-SI
	P-DM		P-FM		Gaussian priors		Empirical priors
	$D_{\mathrm{MSE}}$	$D_{\mathrm{RE}}$	$D_{\mathrm{MSE}}$	$D_{\mathrm{RE}}$	$D_{\mathrm{MSE}}$	$D_{\mathrm{RE}}$	$D_{\mathrm{MSE}}$	$D_{\mathrm{RE}}$
100	7.945e-04	1.197e-01	7.228e-04	1.112e-01	8.776e-04	1.254e-01	8.576e-04	1.244e-01
50	7.431e-03	3.158e-01	8.025e-04	1.197e-01	8.736e-04	1.250e-01	8.641e-04	1.249e-01
10	3.130e-02	7.563e-01	8.796e-04	1.200e-01	8.971e-04	1.269e-01	9.390e-04	1.298e-01
1	8.262e+02	2.851e+02	9.846e-04	1.334e-01	3.434e-03	2.502e-01	2.845e-03	2.265e-01

Table 2: Sampling accuracy vs. number of integration steps. Linear-interpolation–based methods (P-FM, P-SI) remain stable with far fewer steps, while the P-DM model degrades sharply under coarse discretization.

Sample steps	P-DM		P-FM		P-SI
Sample steps	P-DM		P-FM		Gaussian priors		Empirical priors
100	1.278e-01	$\pm$ 2.467e-03	9.995e-01	$\pm$ 1.881e-04	1.929e-01	$\pm$ 3.938e-03	1.857e-01	$\pm$ 1.357e-02
50	1.717e-01	$\pm$ 3.293e-03	9.995e-01	$\pm$ 1.718e-04	2.602e-01	$\pm$ 5.120e-03	2.508e-01	$\pm$ 1.770e-02
10	2.553e-01	$\pm$ 4.854e-03	9.997e-01	$\pm$ 1.153e-04	4.524e-01	$\pm$ 8.342e-03	4.373e-01	$\pm$ 2.716e-02
1	1.000e+00	$\pm$ 0.000e+00	1.000e+00	$\pm$ 0.000e+00	1.000e+00	$\pm$ 0.000e+00	1.000e+00	$\pm$ 0.000e+00

Table 3: Trajectory straightness (

\mathcal{S}

) for different sampling methods. Higher values (

\to 1

) indicate straighter paths. P-FM’s near-perfect straightness explains its robustness to coarse time discretization.

Figure 2: Visualization of 10-step sampling trajectories. Each row shows intermediate states of the generated field for a different sampling strategy, with time $\tau$ evolving according to the model’s process. The comparison between the standard P-DM (Row 1) and an adaptive P-DM with reduced initial variance (Row 2) highlights how sampler design can overcome path curvature. The smooth evolution of P-FM (Row 3) and P-SI (Rows 4-5) visually confirms their straighter transport paths.

Figure 3: Normalized Flow Matching (FM) training loss (log-scale). The explicitly regularized models (MP, GA) achieve the lowest and smoothest loss trajectories, indicating that a well-structured latent space simplifies the generative learning task. The Joint model exhibits a non-stationary spike, while the NoReg model converges to a much higher loss.

Figure 4: t-SNE visualizations of latent space structure. Each column represents a different training strategy. The unregularized latent space (Column 2) is visibly distorted compared to the physical space ground truth (Column 1). Joint training (Columns 3-6) and explicit regularization (Columns 7-8) produce far more coherent structures.

Figure 5: Temporal evolution of relative simulation error ( $D_{\mathrm{RE}}$ ). Comparison of stochastic trajectories (dotted lines, representing the mean over 1000 runs) and ensemble-mean predictions (dashed/solid lines) for various closure models. All generative closures significantly outperform the uncorrected baseline (not shown, error reaches 0.84), with the L-FM model achieving the lowest error.

Figure 6: Qualitative comparison of vorticity field evolution. Snapshots from $t=30$ to $t=50$ for the high-fidelity ground truth, three generative closure models, and the uncorrected simulation. The closure-corrected simulations successfully capture the fine-scale flow structures, whereas the uncorrected simulation diverges and develops spurious features.

Figure 7: Spatial distribution of simulation uncertainty. The pixel-wise standard deviation is computed across an ensemble of 1000 stochastic simulations. The top row shows the ground truth variability, while the bottom row shows the variability captured by the L-FM closure model. The close agreement in both structure and magnitude demonstrates the model’s ability to reproduce the physical uncertainty of the system.

Figure 8: Vorticity energy spectra at different time instances. Spectra from simulations using various ensemble-mean closures are compared against the ground truth. All closure-equipped models correctly reproduce the energy distribution and maintain the characteristic $k^{-3}$ slope of the enstrophy cascade, confirming their physical consistency.

Diffusion Model Sampling

Diffusion model sampling studies how to generate high-quality samples efficiently by designing reverse processes, solvers, and inference-time controls.

See all 79 papers

Stochastic and Non-local Closure Modeling for Nonlinear Dynamical Systems via Latent Score-based Generative Models
Jun 25, 2025Xinghao Dong, Huchen Yang, Jin-Long WuDynamical Systems Latent Diffusion Models
Department of Mechanical Engineering, University of Wisconsin-Madison, Madison, WI, 53706, USA
StocBench: A Benchmark for Generative Modeling of Stochastic Dynamics
Aug 23, 2026Sebastian Pfister, Benjamin Holzschuh, Nils ThuereyGenerative Modeling Probabilistic Forecasting
School of Computation, Information and Technology, Technical University of Munich, Germany.
A Unified Measure-Theoretic View of Diffusion, Score-Based, and Flow Matching Generative Models
May 7, 2026Aditya Ranganath, Mukesh SinghalFlow Matching Generative Modeling
Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA 94551, USA · Department of Electrical Engineering University of California Merced, CA 95343, USA

Jun 25, 2025J/K move · Enter open · S save

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

Xinghao Dong, Huchen Yang, Jin-Long Wu

Department of Mechanical Engineering, University of Wisconsin-Madison, Madison, WI, 53706, USA

Dynamical Systems Latent Diffusion Models Conditional Diffusion Models Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

We propose a latent score-based generative AI framework for learning stochastic, non-local closure models and constitutive laws in nonlinear dynamical systems of computational mechanics. This work addresses a key challenge of modeling complex multiscale dynamical systems without a clear scale separation, for which numerically resolving all scales is prohibitively expensive, e.g., for engineering turbulent flows. While classical closure modeling methods leverage domain knowledge to approximate subgrid-scale phenomena, their deterministic and local assumptions can be too restrictive in regimes lacking a clear scale separation. Recent developments of diffusion-based stochastic models have shown promise in the context of closure modeling, but their prohibitive computational inference cost limits practical applications in many real-world settings. This work addresses this limitation by jointly training convolutional autoencoders with conditional diffusion models in latent space, significantly reducing the dimensionality of the sampling process while preserving essential physical characteristics. Numerical results demonstrate that the joint training approach helps discover a proper latent space that not only guarantees small reconstruction errors but also ensures good performance of the diffusion model in the latent space. When integrated into numerical simulations, the proposed stochastic modeling framework via latent conditional diffusion models achieves significant computational acceleration while maintaining comparable predictive accuracy to standard diffusion models in physical space.

Abstract

Diffusion Model Sampling

Explore similar work

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

StocBench: A Benchmark for Generative Modeling of Stochastic Dynamics

A Unified Measure-Theoretic View of Diffusion, Score-Based, and Flow Matching Generative Models