Schrödinger bridge for generative AI: Soft-constrained formulation and convergence analysis
Jin Ma, Ying Tan, Renyuan Xu
TL;DR
The paper reframes generative modeling as solving a soft-constrained Schrödinger bridge problem (SCSBP) by replacing hard terminal equality constraints with a penalty $kG(\cdot)$, yielding a McKean–Vlasov stochastic control formulation. It proves the existence of optimal policies for every penalty level $k$ and shows that, as $k\to\infty$, these policies and the associated value functions converge linearly to those of the classical SBP, employing Doob's $h$-transform, Schrödinger-potential stability, $\Gamma$-convergence, and a Schauder-fixed-point argument. For the delta-initial case, the optimal policy has an explicit feedback form $\widehat{\alpha}^k_t = \nabla \log h^k(t,X^{\widehat{\alpha}^k}_t)$ with $h^k$ tied to the density $f_k$; the paper provides precise rate bounds $\int_0^T|\widehat{\alpha}^k_t-\widehat{\alpha}_t| dt \le C/k$ and $|J_{\varepsilon}(\widehat{\alpha}^k) - J_{\varepsilon}(\widehat{\alpha})| \le C/k$ (via early stopping). It then extends these results to general initial measures using a fixed-point framework on a convex, compact set, establishing continuity of Schrödinger potentials and convergence guarantees. These contributions yield quantitative convergence guarantees for soft-constraint regularization in SBP and offer a robust, flexible foundation for data generation, fine-tuning, and transfer learning in generative AI.
Abstract
Generative AI can be framed as the problem of learning a model that maps simple reference measures into complex data distributions, and it has recently found a strong connection to the classical theory of the Schrödinger bridge problems (SBPs) due partly to their common nature of interpolating between prescribed marginals via entropy-regularized stochastic dynamics. However, the classical SBP enforces hard terminal constraints, which often leads to instability in practical implementations, especially in high-dimensional or data-scarce regimes. To address this challenge, we follow the idea of the so-called soft-constrained Schrödinger bridge problem (SCSBP), in which the terminal constraint is replaced by a general penalty function. This relaxation leads to a more flexible stochastic control formulation of McKean-Vlasov type. We establish the existence of optimal solutions for all penalty levels and prove that, as the penalty grows, both the controls and value functions converge to those of the classical SBP at a linear rate. Our analysis builds on Doob's h-transform representations, the stability results of Schrödinger potentials, Gamma-convergence, and a novel fixed-point argument that couples an optimization problem over the space of measures with an auxiliary entropic optimal transport problem. These results not only provide the first quantitative convergence guarantees for soft-constrained bridges but also shed light on how penalty regularization enables robust generative modeling, fine-tuning, and transfer learning.