Light and Optimal Schrödinger Bridge Matching
arxiv(2024)
摘要
Schrödinger 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
optimal Schrödinger bridge matching. It exploits the optimal
parameterization of the diffusion process and provably recovers the SB process
(a) with a single bridge matching step and (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ödinger
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
.
更多查看译文
AI 理解论文
溯源树
样例
生成溯源树,研究论文发展脉络
Chat Paper
正在生成论文摘要