cs.LGMay 20, 2026

On the Regularity and Generalization of One-Step Wasserstein-guided Generative Models for PDE-Induced Measures

Authors: Likun LinZhongjian WangJack XinZhiwen Zhang

Organizations: Department of Mathematics, The University of Hong Kong, Pokfulam Road, Hong Kong SAR, P.R. China. · Division of Mathematical Sciences, School of Physical and Mathematical Sciences, Nanyang Technological University, Singapore 637371. · Department of Mathematics, University of California at Irvine, Irvine, CA 92697, USA. · Materials Innovation Institute for Life Sciences and Energy (MILES), HKU-SIRI, Shenzhen, 518045, P.R. China.

Abstract

Despite the remarkable empirical success of generative models, the available theory on their statistical accuracy in scientific computing remains largely pessimistic. This paper develops a theoretical framework for understanding the regularity of transport maps and the generalization properties of one-step Wasserstein-guided generative models for PDE-induced probability measures. We consider normalized target densities associated with linear elliptic and parabolic equations on bounded domains, as well as diffusion and Fokker--Planck equations on the torus. Under standard structural assumptions, we prove that these target measures satisfy doubling conditions. By combining this fact with regularity theory for optimal transport between doubling measures, we show that the optimal transport map from a uniform source measure to the target measure is Hölder continuous. This regularity yields an approximation-theoretic justification for one-step generative models that learn PDE-induced distributions via a single pushforward map. As a representative instance, we study DeepParticle and derive excess-risk bounds characterizing the discrepancy between the learned map and the population-optimal map. We also establish a robustness estimate under target shift and illustrate the theory with experiments which support the derived rates.

Explore similar work

Oct 2, 2024stat.ML

Robustness and Structure Preservation in Flow-Based Generative Models via Wasserstein Path-Space Divergences

We introduce a novel Wasserstein-1 (W1W_1) path-space divergence for stochastic and deterministic dynamics and establish a Wasserstein Uncertainty Propagation (WUP) theorem that bounds the W1W_1 distance between terminal distributions by the proposed divergence, equivalently characterized by a weighted L2L^2 discrepancy between the underlying drifts and the W1W_1 distance between their initial measures. A key ingredient is a probabilistic framework combining adjoint Feynman-Kac representations with synchronous coupling (and reflection coupling on bounded domains), yielding Wasserstein stability estimates beyond existing PDE- and Girsanov-based approaches. The framework accommodates time-varying and possibly degenerate diffusion coefficients, empirical and singular measures, and remains valid in the deterministic limit of flow matching. Unlike KL-based uncertainty quantification bounds, it does not require absolute continuity of path measures and therefore remains well-defined in singular settings. As consequences of the WUP theorem, we derive W1W_1 robustness and generalization bounds for score-based generative models and flow matching at both population and finite-sample levels. We further specialize the framework to group-symmetric targets, providing the first error analysis of equivariant flow-based models and the first quantitative comparison between data augmentation and equivariant inductive bias. Our analysis identifies a symmetry-aware Wasserstein path-space divergence that quantifies the model-form error induced by non-equivariant parametrizations. We prove that this error cannot be removed by additional data or training and vanishes only under equivariant architectures, establishing a precise theoretical advantage of equivariant inductive bias over data augmentation. Numerical experiments on group-symmetric Gaussian mixtures corroborate the theory.
Ziyu Chen, Markos A. Katsoulakis, Benjamin J. Zhang
Jun 29, 2026cs.LG

The Fundamental Limits of Valid Transport Map Estimation

Many modern generative modeling methods, including diffusion models, normalizing flows, and flow matching, estimate transport maps or plans between distributions without explicitly targeting an optimal transport (OT) map. In applications like generative modeling, the transport cost itself is irrelevant, and this makes it natural to target maps which are more tractable from either a statistical or computational standpoint. In this short note, we formalize the task of estimating any valid transport map in a rigorous minimax framework. One consequence of this framing is that it yields sample complexity lower bounds for any method whose learned object is evaluated as a transport map or plan, including flow matching and diffusion-based generative models, in settings where direct analysis would be challenging due to the analytic complexity of the methods and their target maps. We observe that, under standard, though strong, stability assumptions from the OT literature, estimating any valid transport map is statistically as hard as estimating the OT map. We complement these results with some examples showing that when these stability assumptions fail, alternative transport maps can be learned substantially more accurately than the OT map. Our minimax framing provides a rigorous foundation for understanding the statistical limits of modern transport-based generative methods and clarifies when targeting sub-optimal maps can provide real statistical advantages.
Sivaraman Balakrishnan
May 29, 2026cs.LG

A Unifying View of Variational Generative Wasserstein Flows

Many modern generative models can be viewed as minimizing divergences between probability distributions, yet they rely on different algorithmic and geometric principles. Wasserstein gradient flows provide a continuous-time formulation for optimizing over distributions, and can be approximated through their implicit discretization via the Jordan-Kinderlehrer-Otto (JKO) scheme. In this work, we present a unified theoretical framework for generative modeling based on Wasserstein gradient flows, which we refer to as Generative Wasserstein Flows (GWF). We show that a broad class of existing methods can be derived as instances of parametric JKO schemes for ff-divergence objectives, and we establish equivalences between several recently proposed algorithms. We extend this framework beyond f-divergence to Integral Probability Metrics and squared Maximum Mean Discrepancy, deriving new JKO-based generative algorithms, and clarifying their connections with GANs. We study empirically the impact of the JKO regularization for a wide set of objectives. Finally, we analyze parametric Wasserstein flows, where the dynamics are restricted to distributions induced by parametrized maps.
Paul Caucheteux, Clément Bonet, Anna Korba