Schrödinger Bridge with Quadratic State Cost is Exactly Solvable

Abstract

Schrödinger bridge is a diffusion process that steers a given distribution to another in a prescribed time while minimizing the effort to do so. It can be seen as the stochastic dynamical version of the optimal mass transport, and has growing applications in generative diffusion models and stochastic optimal control. {\black{We say a Schrödinger bridge is ``exactly solvable'' if the associated uncontrolled Markov kernel is available in closed form, since then the bridge can be numerically computed using dynamic Sinkhorn recursion for arbitrary endpoint distributions with finite second moments.}} In this work, we propose a regularized variant of the Schrödinger bridge with a quadratic state cost-to-go that incentivizes the optimal sample paths to stay close to a nominal level. Unlike the conventional Schrödinger bridge, the regularization induces a state-dependent rate of killing and creation of probability mass, and its solution requires determining the Markov kernel of a reaction-diffusion partial differential equation. We derive this Markov kernel in closed form, {\black{showing that the regularized Schrödinger bridge is exactly solvable, even for non-Gaussian endpoints. This advances the state-of-the-art because closed form Markov kernel for the regularized Schrödinger bridge is available in existing literature only for Gaussian endpoints}}. Our solution recovers the heat kernel in the vanishing regularization (i.e., diffusion without reaction) limit, thereby recovering the solution of the conventional Schrödinger bridge {\black{as a special case}}. We deduce properties of the new kernel and explain its connections with certain exactly solvable models in quantum mechanics.

Figures (6)

• Figure 1: Sample paths for 1D SB with state cost $\frac{1}{2}Qx^2$ for fixed $Q\geq 0$, time horizon $[t_0,t_1]=[0,1]$, and endpoint PDFs $\rho_0 = \sum_{\mu_0\in\{-1,1\}}\frac{1}{2}\mathcal{N}(\mu_0,0.05^2)$, $\rho_1=\mathcal{N}(0,0.5^2)$.
• Figure 1: Dynamic Sinkhorn recursion to solve SBP with linear convergence in Hilbert metric. In this work, the $\mathcal{L}_{\rm{forward}}, \mathcal{L}_{\rm{backward}}$ are as in \ref{['FactorPDEs']}.
• Figure 1: Solution to PDE \ref{['FactorPDEForward']} in $n=2$ dimensions with $t_0 = 0$ and initial condition $\widehat{\varphi}_{0}(\cdot)=1$. All subfigures are over the spatial domain $\left[-7,7\right]^{2}$.
• Figure 1: Solution to PDE \ref{['FactorPDEForward']} in $n=2$ dimensions with $t_0 = 0$ and initial condition $\widehat{\varphi}_{0}(\cdot)=\frac{1}{ 2\pi}e^{-\frac{1}{2}\|\cdot\|^2}$. All subfigures are over the spatial domain $\left[-5,5\right]^{2}$.
• Figure 2: Solution to PDE \ref{['FactorPDEForward']} in $n=2$ dimensions with $t_0 = 0$ and initial condition $\widehat{\varphi}_{0}(\cdot)=\rho_{\infty}$. All subfigures are over the spatial domain $\left[-0.25,1.25\right]^{2}$.

Theorems & Definitions (19)

• Lemma 2.1: A series sum for exponential of negative quadratic
• Proof 1
• Lemma 2.2: Central identity of QFT
• Theorem 3.1: Existence-Uniqueness of solution for SB with state cost-to-go
• Theorem 3.2: Minimizer of SB problem with state cost-to-go
• Remark 3.3
• Remark 3.4
• Lemma 4.1: Solution of a parametric second order nonlinear ODE
• Proposition 4.2: IVP solution for \ref{['FactorPDEForwardNew']} with positive diagonal $\bm{D}$ as series sum
• Theorem 4.3: Closed form kernel $\kappa_{++}$ for IVP solution of \ref{['FactorPDEForwardNew']} with positive diagonal $\bm{D}$