Light and Optimal Schrödinger Bridge Matching

要約

Schr’odinger Bridges (SB)は、古典的な拡散モデルの有望な拡張として、またエントロピー最適輸送(Entropic Optimal Transport (EOT))にも関連するものとして、最近MLコミュニティの注目を集めている。SBに対する最近のソルバーは、広く普及しているブリッジマッチング手続きを利用している。このような手続きは、分布間の輸送計画のみが与えられた場合に、分布間で質量を輸送する確率過程を回復することを目的としている。特に、EOT計画が与えられれば、これらの手続きはSBを解くのに適応できる。この事実は、マッチングに基づくSBソルバーを提供する最近の研究によって大いに利用されている。ここでの要はEOT計画を回復することである。最近の研究では、ヒューリスティックな近似（例えばミニバッチOT）を用いるか、あるいは設計上学習中に誤差を蓄積する反復マッチング手順を確立している。我々はこれらの限界に対処し、SBを学習する新しい手順を提案する。これは、拡散過程の最適なパラメータ化を利用し、1回のブリッジマッチングでSB過程を回復することを証明する。さらに、最適な橋梁マッチング目的が、最近発見されたEOT/SBを学習するためのエネルギーベースモデリング（EBM）目的と一致することを示す。この観察に触発されて、我々は、シュルオディンガーポテンシャルのガウス混合パラメタリゼーションを用いて、実際に最適マッチングを実装するためのライトソルバ（LightSB-Mと呼ぶ）を開発する。我々は、様々な実用的なタスクにおいて、我々のソルバーの性能を実験的に示す。LightSB-Mソルバーのコードは \url{https://github.com/SKholkin/LightSB-Matching} にあります。

要約(オリジナル)

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}.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL