Light and Optimal Schrödinger Bridge Matching
Nikita Gushchin Sergei Kholkin Evgeny Burnaev Alexander Korotin
Schr\"odinger Bridges (SB) have recently gained the attention of the ML
community as a promising extension of classic diffusion models which is also
interconnected to the Entropic Optimal Transport (EOT). Recent solvers for SB
exploit the pervasive bridge matching procedures. Such procedures aim to
recover a stochastic process transporting the mass between distributions given
only a transport plan between them. In particular, given the EOT plan, these
procedures can be adapted to solve SB. This fact is heavily exploited by recent
works giving rives to matching-based SB solvers. The cornerstone here is
recovering the EOT plan: recent works either use heuristical approximations
(e.g., the minibatch OT) or establish iterative matching procedures which by
the design accumulate the error during the training. We address these
limitations and propose a novel procedure to learn SB which we call the
\textbf{optimal Schr\"odinger bridge matching}. It exploits the optimal
parameterization of the diffusion process and provably recovers the SB process
\textbf{(a)} with a single bridge matching step and \textbf{(b)} with arbitrary
transport plan as the input. Furthermore, we show that the optimal bridge
matching objective coincides with the recently discovered energy-based modeling
(EBM) objectives to learn EOT/SB. Inspired by this observation, we develop a
light solver (which we call LightSB-M) to implement optimal matching in
practice using the Gaussian mixture parameterization of the Schr\"odinger
potential. We experimentally showcase the performance of our solver in a range
of practical tasks. The code for the LightSB-M solver can be found at
\url{https://github.com/SKholkin/LightSB-Matching}.