Light and Optimal Schrödinger Bridge Matching
- URL: http://arxiv.org/abs/2402.03207v2
- Date: Tue, 30 Jul 2024 13:43:36 GMT
- Title: Light and Optimal Schrödinger Bridge Matching
- Authors: Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, Alexander Korotin,
- Abstract summary: We propose a novel procedure to learn Schr"odinger Bridges (SB) which we call the textbf Schr"odinger bridge matching.
We show that the optimal bridge matching objective coincides with the recently discovered energy-based modeling (EBM) objectives to learn EOT/SB.
We develop a light solver (which we call LightSB-M) to implement optimal matching in practice using the mixture parameterization of the Schr"odinger potential.
- Score: 67.93806073192938
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- Abstract: 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 rise 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 adjusted Schr\"odinger potential. We experimentally showcase the performance of our solver in a range of practical tasks. The code for our solver can be found at https://github.com/SKholkin/LightSB-Matching.
Related papers
- Adversarial Schrödinger Bridge Matching [66.39774923893103]
Iterative Markovian Fitting (IMF) procedure alternates between Markovian and reciprocal projections of continuous-time processes.
We propose a novel Discrete-time IMF (D-IMF) procedure in which learning of processes is replaced by learning just a few transition probabilities in discrete time.
We show that our D-IMF procedure can provide the same quality of unpaired domain translation as the IMF, using only several generation steps instead of hundreds.
arXiv Detail & Related papers (2024-05-23T11:29:33Z) - Generalized Schrödinger Bridge Matching [54.171931505066]
Generalized Schr"odinger Bridge (GSB) problem setup is prevalent in many scientific areas both within and without machine learning.
We propose Generalized Schr"odinger Bridge Matching (GSBM), a new matching algorithm inspired by recent advances.
We show that such a generalization can be cast as solving conditional optimal control, for which variational approximations can be used.
arXiv Detail & Related papers (2023-10-03T17:42:11Z) - Light Schrödinger Bridge [72.88707358656869]
We develop a lightweight, simulation-free and theoretically justified Schr"odinger Bridges solver.
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.
arXiv Detail & Related papers (2023-10-02T13:06:45Z) - Building the Bridge of Schr\"odinger: A Continuous Entropic Optimal
Transport Benchmark [96.06787302688595]
We propose a novel way to create pairs of probability distributions for which the ground truth OT solution is known by the construction.
We use these benchmark pairs to test how well existing neural EOT/SB solvers actually compute the EOT solution.
arXiv Detail & Related papers (2023-06-16T20:03:36Z) - Diffusion Schr\"odinger Bridge Matching [36.95088080680221]
We introduce Iterative Markovian Fitting (IMF) and Diffusion Schr"odinger Bridge Matching (DSBM)
IMF is a new methodology for solving SB problems, and DSBM is a novel numerical algorithm for computing IMF iterates.
We demonstrate the performance of DSBM on a variety of problems.
arXiv Detail & Related papers (2023-03-29T16:59:22Z) - The Schr\"odinger Bridge between Gaussian Measures has a Closed Form [101.79851806388699]
We focus on the dynamic formulation of OT, also known as the Schr"odinger bridge (SB) problem.
In this paper, we provide closed-form expressions for SBs between Gaussian measures.
arXiv Detail & Related papers (2022-02-11T15:59:01Z) - Likelihood Training of Schr\"odinger Bridge using Forward-Backward SDEs
Theory [29.82841891919951]
It remains unclear whether the optimization principle of SB relates to the modern training of deep generative models.
We present a novel computational framework for likelihood training of SB models grounded on Forward-Backward Theory.
We show that the resulting training achieves comparable results on generating realistic images on MNIST, CelebA, and CIFAR10.
arXiv Detail & Related papers (2021-10-21T17:18:59Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.