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

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.

• The paper shows that flow matching with linear transport paths enables accurate single-step closure sampling, retaining a relative error of about 13.3% compared with 11.1% at 100 steps while outperforming diffusion models.
• The paper demonstrates that reconstruction quality alone cannot ensure generative accuracy, whereas explicit metric-preserving latent regularization reduces latent generation error to roughly 2.9% and physical-space error to about 8.3%.
• The paper finds that metric-preserving latent flow matching improves a posteriori simulation accuracy to 4.01%, preserves uncertainty structure and the enstrophy spectrum, and generates ensembles up to seven times faster than physical-space models.

Motivation and problem setting

Closure modeling for coarse-grained dynamical systems requires approximating the effect of unresolved scales on resolved variables. Deterministic closures rest on a scale-separation assumption that fails when unresolved scales are far from equilibrium, motivating stochastic formulations in which the closure term is treated as a random object and its full conditional distribution $p(\mathbf{U} \mid \mathbf{V})$ is learned from data. Transport-based generative models—score-based diffusion, flow matching (FM), and stochastic interpolants (SI)—are natural candidates, but diffusion-based samplers require hundreds of iterative steps, which is prohibitive when a new closure sample must be drawn at every time step of an online simulation.

This paper addresses two coupled bottlenecks: sampling speed of transport-based generative models, and the geometric quality of the latent space in which sampling is performed. The central claims are that (i) flow matching with linear interpolation paths enables single-step sampling up to two orders of magnitude faster than iterative diffusion, and (ii) latent spaces trained only for reconstruction are geometrically distorted and catastrophically degrade generative fidelity, while both implicit regularization via joint training and explicit metric-preserving (MP) or geometry-aware (GA) regularizers restore the manifold structure needed for accurate stochastic closures.

Framework

The target system is a spatiotemporal PDE $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$, coarse-grained to resolved variables $\mathbf{V}$ with a closure term $\mathbf{U} = \mathcal{C}(\mathbf{V})$ capturing subgrid effects such as backscatter and dissipation. Rather than learning an explicit SPDE for $\mathbf{U}$, the framework learns the stationary conditional distribution $p(\mathbf{U}\mid\mathbf{V})$ directly with conditional transport-based generative models operating in a compressed latent space produced by convolutional autoencoders.

Three generative paradigms are compared:

• Score-based diffusion: a VE-SDE forward process reversed by integrating a learned score function; highly curved transport paths necessitate fine time discretization.
• Flow matching: conditional flow matching with linear interpolation paths and constant conditional velocities; training reduces to regression against $(\mathbf{x}_1 - \mathbf{x}_0)$ without ODE simulation.
• Stochastic interpolants: interpolants $$\mathbf{x}_\tau = \alpha_\tau \mathbf{x}_0 + \beta_\tau \mathbf{x}_1 + \sigma_\tau \mathbf{W}_\tau$$ with drift regression; notably, SI permits non-Gaussian source distributions, allowing direct transport from the conditioning variable's distribution to the closure distribution.

For latent-space structuring, two strategies are contrasted: end-to-end joint training (reconstruction + generative loss + KL regularization against latent collapse), which provides implicit but uncontrolled regularization; and two-phase training with explicit structural losses penalizing discrepancies between physical-space and latent-space distances over local neighborhoods, using either Euclidean distances (MP) or precomputed geodesic/manifold distances (GA).

Numerical setup

Experiments use 2D Kolmogorov flow: incompressible Navier–Stokes in vorticity form at $\nu = 10^{-3}$ on a periodic domain, driven by deterministic large-scale forcing plus spatially correlated white-in-time noise ($Q$-Wiener process, amplitude $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$0). The closure target is $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$1, combining nonlinear advection and stochastic forcing. Training data comprise 20,000 pairs from 100 pseudo-spectral DNS runs on $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$2 grids, downsampled to $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$3, split 18,000/2,000 train/test. Latent fields are $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$4; the generative backbone is a two-branch Fourier Neural Operator conditioned on encoded vorticity.

Physical-space baselines and the role of path geometry

With sufficient integration steps (100), all three paradigms achieve comparable ensemble-mean relative errors of roughly 8.6–9.0% and reproduce the analytically derived ground-truth field standard deviation of 0.02 within 10%, confirming that all paradigms capture both mean behavior and uncertainty magnitude.

Efficiency, however, differs sharply. Reducing integration steps degrades the diffusion model severely—a single-step P-DM sample yields a relative error of order $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$5, i.e., catastrophic failure—whereas P-FM retains essentially unchanged accuracy even at one step ($\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$6 vs. $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$7 at 100 steps). A straightness metric $\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$8 (ratio of endpoint distance to integrated path length) explains this: P-FM trajectories are nearly perfectly straight ($\partial_t \mathbf{v} = \mathcal{M}(\mathbf{v})$9) while P-DM paths are strongly curved ($\mathbf{V}$0). P-SI occupies an intermediate position, with single-step errors around 0.23–0.25. The paper also shows that diffusion can be partially rescued within a fixed step budget by initializing the reverse SDE at reduced variance (exploiting the narrowness of the conditional target) and applying adaptive time-stepping—but this is sampler engineering rather than a property of the paradigm itself. The implication is that for online closure sampling, where per-step cost multiplies across simulation time steps, path geometry is the primary determinant of practical viability.

Latent geometry determines generative success

The most consequential finding concerns autoencoder training strategy. Reconstruction-only autoencoders achieve the lowest reconstruction error of all tested configurations yet produce catastrophic generative failure: latent-space relative error near 30% and decoded physical-space error of about 33%, far worse than any physical-space baseline. Quantitative distortion metrics corroborate this—the NoReg latent space exhibits large Procrustes disparity and nearly a threefold increase in the conditional coefficient of variation relative to physical space. This establishes the paper's central thesis plainly: faithful reconstruction is necessary but insufficient; latent geometry governs downstream generative performance.

Both regularization routes repair this failure. Joint training reduces latent error more than fourfold relative to NoReg, reaching physical-space accuracy comparable to physical-space baselines, though its loss curve exhibits a non-stationary spike attributable to co-adapting autoencoder weights. Explicit regularization proves most effective and stable: MP-regularized models attain the lowest latent generation error (~2.9%) and a final physical-space error (~8.3%) matching or slightly surpassing the physical-space baselines. Notably, MP consistently outperforms GA. The authors attribute this to alignment with sampler mechanics—discretized SDE/ODE integration, interpolation, and gradient evaluation operate in Euclidean coordinates, so preserving local Euclidean distances matters more than preserving geodesics—and MP is also cheaper since it avoids geodesic estimation. This is a defensible mechanistic explanation, though it is supported here only empirically on a single benchmark.

A posteriori validation

Embedding the closures in forward Navier–Stokes simulations from $\mathbf{V}$1 to $\mathbf{V}$2 demonstrates operational value. The uncorrected solver accumulates 84.2% relative error by $\mathbf{V}$3; all generative closures keep final error below 13%. The best configuration—latent FM with MP regularization—achieves a 4.01% final ensemble-mean error, roughly a twofold improvement over the physical-space diffusion model. Ensemble-mean closures reduce error by 20–30% relative to single-sample stochastic trajectories, trading compute for variance reduction. The L-FM model also reproduces the spatial structure and magnitude of ensemble uncertainty and preserves the $\mathbf{V}$4 enstrophy-cascade spectrum up to the training resolution limit, indicating physically consistent subgrid variability rather than injected noise.

Computationally, latent-space models generate ensembles approximately seven times faster than physical-space counterparts, and one-step L-FM is twice as fast as iterative L-DM while being more accurate. Combined with the single-step capability established earlier, this substantiates the abstract's claim of up to two orders of magnitude speedup over iterative diffusion approaches.

Limitations and open questions

Several caveats bound these results. All evidence comes from a single 2D Kolmogorov flow with a prescribed additive noise whose variance admits an analytical reference; generalization to 3D turbulence, non-periodic domains, or systems lacking such analytical uncertainty targets remains untested. The claimed data efficiency ("without demanding a huge amount of training data") is asserted qualitatively rather than demonstrated through ablations over dataset size. The superiority of MP over GA rests on one benchmark, and the hyperparameters governing joint training ($\mathbf{V}$5, $\mathbf{V}$6) required grid search, so the robustness of the implicit-regularization route is unclear. The GA geodesic estimates depend on the manifold-distance approximation used, which is not analyzed for sensitivity. Finally, the closure update is applied every five physical time steps for efficiency; the sensitivity of accuracy to this update interval is not characterized.

Conclusion

This work systematically compares diffusion, flow matching, and stochastic interpolant paradigms for stochastic closure modeling and identifies flow matching in a metric-preserving regularized latent space as the fastest and most accurate configuration, enabling single-step sampling with a posteriori errors as low as ~4% over 20-second integrations and order-of-magnitude ensemble acceleration. Its principal methodological contribution is the demonstration that latent geometry—not reconstruction fidelity—is the controlling factor for latent generative closures, and that explicit Euclidean metric preservation is the most effective, stable, and efficient way to impose it.