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 ω .
Sample steps
P-DM
P-FM
P-SI
Gaussian priors
Empirical priors
DMSE
DRE
DMSE
DRE
DMSE
DRE
DMSE
DRE
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
Gaussian priors
Empirical priors
100
1.278e-01
± 2.467e-03
9.995e-01
± 1.881e-04
1.929e-01
± 3.938e-03
1.857e-01
± 1.357e-02
50
1.717e-01
± 3.293e-03
9.995e-01
± 1.718e-04
2.602e-01
± 5.120e-03
2.508e-01
± 1.770e-02
10
2.553e-01
± 4.854e-03
9.997e-01
± 1.153e-04
4.524e-01
± 8.342e-03
4.373e-01
± 2.716e-02
1
1.000e+00
± 0.000e+00
1.000e+00
± 0.000e+00
1.000e+00
± 0.000e+00
1.000e+00
± 0.000e+00
Table 3: Trajectory straightness ( S ) for different sampling methods. Higher values ( →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 τ 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 ( DRE ). 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 studies how to generate high-quality samples efficiently by designing reverse processes, solvers, and inference-time controls.
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.
Xinghao Dong, Huchen Yang, Jin-Long Wu
Department of Mechanical Engineering, University of Wisconsin-Madison, Madison, WI, 53706, USA
We benchmark transport-based generative models as well as distillation-based few-step methods for the probabilistic forecasting of stochastic fluid flows, with a particular focus on performance under limited inference budgets. All methods are evaluated on a two-dimensional Kolmogorov flow with stochastic forcing. We measure one-step distributional accuracy against large simulated reference ensembles and assess whether the invariant measure is preserved during autoregressive rollouts via the enstrophy spectrum. On the stochastic task, flow matching achieves the most accurate one-step conditional distribution at high inference budgets, while the second-order exponential integrator DPM-2 is strongest at very low NFE. Few-step distillation methods are competitive with the multi-step methods and preserve the enstrophy spectrum particularly well. A deterministic control task, in which the forcing over the prediction interval is observed, separates aleatoric from epistemic uncertainty. Model performance does not translate between the two settings: the distilled models are competitive on the stochastic task but least accurate on the control task. While stochastic diffusion samplers such as DDPM better preserve the enstrophy spectrum during rollouts in the stochastic setting, deterministic samplers such as DDIM and DPM-2 show better spectral preservation in the deterministic setting.
Sebastian Pfister, Benjamin Holzschuh, Nils Thuerey
School of Computation, Information and Technology, Technical University of Munich, Germany.
We survey continuous-time generative modeling methods based on transporting a simple reference distribution to a data distribution via stochastic or deterministic dynamics. We present a unified framework in which diffusion models, score-based generative models, and flow matching are instances of learning a time-dependent vector field that induces a family of marginals (ρt)t∈[0,1] governed by continuity and Fokker-Planck equations. Such a unified theory is timely because these methods are converging methodologically, yet fragmented notation and competing derivations continue to obscure their shared structure and the practical tradeoffs governing sampling, stability, and computation. Within this framework, we (i) derive reverse-time sampling for diffusion and score-based models as controlled stochastic dynamics, (ii) show that the probability flow ODE yields identical marginals and connects diffusion to likelihood-based normalizing flows, and (iii) interpret flow matching as direct regression of the velocity field under a chosen interpolation, clarifying when it coincides with or differs from score-based training. We compare objectives, sampling schemes, and discretization errors under unified notation, discuss connections to Schrodinger bridges and entropic optimal transport, and summarize theoretical guarantees and open problems on approximation, stability, and scalability.
Aditya Ranganath, Mukesh Singhal
Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA 94551, USA · Department of Electrical Engineering University of California Merced, CA 95343, USA