One Coupling to Rule Them All: Optimal Transport as the Unifying Geometry of Diffusion Models, Flow Matching, and Reasoning in Deep Generative Models
Mohammad Sajjad Ghaemi
Abstract
This paper presents a unified theoretical framework based on Wasserstein geometry and Optimal Transport (OT) that explains when generative models memorize, generalize, or reason, alongside theoretical metrics and implications for Flow Matching (FM), diffusion models, and autoregressive (AR) variants that are widely used across benchmarks. We establish three formal identifications: (1) memorization corresponds to a transport coupling concentrated on training-data pairs, with a coupling gap lower-bounded by the empirical Wasserstein convergence rate $\Omega(n^{-2/d})$; (2) generalization corresponds to learning a coupling that approximates the Wasserstein-geodesic (Brenier) map, which is independent of any specific training example by definition; and (3) compositional reasoning in AR models is equivalent to correctly propagating a multi-marginal OT plan through the sequential factorization, with the compositional gap. The exposure bias inherent in AR generation is characterized by the Wasserstein covariate shift occurring between the teacher-forced and free-running conditional marginal distributions. Subsequently, five geometrical evaluation metrics are derived from these identifications to demonstrate that holistic FM guidance of AR models creates a mathematical correction to the compositional gap, rather than a mere heuristic trick. Moreover, this framework has extensive applications across diverse domains, including scientific discovery, proteomic sequence design, and structured prediction.
Chat is not available.
Successful Page Load