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

TL;DR

This paper demonstrates that transport-based generative models in a latent space enable fast, physically faithful stochastic closure modeling for complex dynamical systems, significantly reducing sampling time compared to traditional diffusion models.

Contribution

It introduces a latent space flow matching approach for stochastic closure modeling, combining regularization techniques to preserve physical fidelity and topological features.

Findings

Single-step sampling is up to 100 times faster than iterative diffusion methods.

Latent space regularization preserves key topological information.

Fewer training data are needed due to topological inheritance.

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.