Light Schrödinger Bridge

TL;DR

LightSB addresses the need for a simple, fast baseline solver for continuous Schrödinger Bridges by marrying a Gaussian-mixture parameterization of the Schrödinger potential with an energy-based view of the log-Schrödinger potentials. The method yields a tractable objective, a closed-form conditional plan, and a diffusion with an explicit drift, enabling simulation-free training on CPU and rapid inference. Theoretical guarantees include a universal approximation property for SBs and finite-sample generalization bounds, while experiments across 2D toys, SB/EOT benchmarks, single-cell data, and latent-space image translation demonstrate strong performance with minimal computational cost. The approach offers a practical, accessible alternative to heavier neural solvers and could establish LightSB as a standard, easy-to-use SB baseline for moderate-dimensional problems.

Abstract

Despite the recent advances in the field of computational Schrödinger Bridges (SB), most existing SB solvers are still heavy-weighted and require complex optimization of several neural networks. It turns out that there is no principal solver which plays the role of simple-yet-effective baseline for SB just like, e.g., $k$-means method in clustering, logistic regression in classification or Sinkhorn algorithm in discrete optimal transport. We address this issue and propose a novel fast and simple SB solver. Our development is a smart combination of two ideas which recently appeared in the field: (a) parameterization of the Schrödinger potentials with sum-exp quadratic functions and (b) viewing the log-Schrödinger potentials as the energy functions. We show that combined together these ideas yield a lightweight, simulation-free and theoretically justified SB solver with a simple straightforward optimization objective. As a result, it allows solving SB in moderate dimensions in a matter of minutes on CPU without a painful hyperparameter selection. Our light solver resembles the Gaussian mixture model which is widely used for density estimation. Inspired by this similarity, we also prove an important theoretical result showing that our light solver is a universal approximator of SBs. Furthemore, we conduct the analysis of the generalization error of our light solver. The code for our solver can be found at https://github.com/ngushchin/LightSB

Figures (6)

• Figure 1: Unpaired male$\rightarrow$female translation by our LightSB solver applied in the latent space of ALAE for 1024x1024 FFHQ images. Our LightSB converges on 4 cpu cores in less than 1 minute.
• Figure 2: The process $T_{\theta}$ learned with LightSB (ours) in Gaussian$\!\rightarrow\!$Swiss roll example.
• Figure 3: Unpaired translation by our LightSB solver applied in the latent space of ALAE for 1024x1024 FFHQ images. Our LightSB converges on 4 cpu cores in less than 1 minute.
• Figure 4: Convergence speed comparison on MSCI dataset, starting day $3$, ending day $7$ and evaluation day $4$.
• Figure 5: LightSB Adult$\rightarrow$Child for different $\epsilon$.

Theorems & Definitions (14)

• Proposition 3.1: Feasible reformulation of the KL minimization
• Proposition 3.2: Tractable form of plan's components
• Proposition 3.3: Properties of $T_{\theta}$
• Theorem 3.4: Gaussian mixture parameterization for the adjusted potential provides the universal approximation of Schrödinger bridges
• Theorem A.1: Bound for the statistical error
• proof : Proof of Proposition \ref{['prop-feasible']}
• proof : Proof of Proposition \ref{['prop-explicit-form']}
• proof : Proof of Proposition \ref{['prop-drift-closed-form']}
• proof : Proof of Theorem \ref{['thm-universal-approximation']}
• Proposition B.1: Rademacher bound on the statistical error