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

TL;DR

It is shown 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.

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 (8)

• 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$.
• 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.