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

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.

💡 Research Summary

The paper addresses a fundamental limitation of diffusion‑based generative models—slow sampling—by introducing a transport‑based approach that operates in a low‑dimensional latent space. Using the two‑dimensional Kolmogorov flow as a testbed, the authors demonstrate that “flow matching” (a continuous probability‑flow formulation) can replace the multi‑step denoising process of traditional diffusion models with a single‑step ODE integration, achieving up to two orders of magnitude faster generation.

To preserve the physical fidelity of the sampled closure terms, the study compares three regularization strategies for the latent space. The first is an implicit joint‑training scheme that simultaneously optimizes the generative network and the stochastic closure model, encouraging the latent representation to inherit the original system’s topology without explicit constraints. The second and third are explicit regularizers: a Metric‑Preserving (MP) loss that aligns pairwise distances in latent space with those in the high‑dimensional state space, and a Geometry‑Aware (GA) loss that embeds Riemannian curvature information of the original dynamics into the latent manifold. All three strategies successfully limit latent distortion, but the joint‑training approach stands out for its simplicity and lack of extra hyper‑parameters.

Quantitative evaluation includes Kullback‑Leibler divergence, Wasserstein distance, energy spectra, and entropy measures. The flow‑matching model matches the statistical properties of the reference Kolmogorov flow almost perfectly, while its sampling time drops from several seconds (or tens of seconds for standard diffusion) to milliseconds. Moreover, the model requires far fewer training samples—on the order of a few thousand—to achieve comparable performance, highlighting its data‑efficiency relative to GANs or VAEs.

The authors also analyze the topological preservation of the latent space. Both explicit regularizers (MP and GA) and the implicit joint scheme retain key manifold structures, enabling the learned stochastic closure term to respect conservation laws and symmetries inherent in the original fluid dynamics. This preservation is crucial for downstream tasks such as sub‑grid modeling, where physical consistency cannot be compromised.

In summary, the paper makes three principal contributions: (1) it demonstrates that transport‑based flow matching in a reduced latent space can dramatically accelerate sampling for stochastic closure modeling; (2) it provides a systematic comparison of implicit and explicit latent‑space regularizations, showing that all can maintain the essential geometry of the underlying dynamical system; and (3) it validates that high‑quality stochastic closure models can be learned with modest data volumes. The methodology is broadly applicable to any high‑dimensional physical system where fast, physically consistent generative sampling is required, including climate modeling, plasma physics, and turbulent flow simulations.