Light and Optimal Schrödinger Bridge Matching

要約

Schr\’odinger Bridges (SB) は、Entropic Optimal Transport (EOT) にも相互接続されている古典的な拡散モデルの有望な拡張機能として、最近 ML コミュニティの注目を集めています。
SB 用の最近のソルバーは、普及したブリッジ マッチング手順を活用しています。
このような手順は、分布間の輸送計画のみが与えられた分布間で質量を輸送する確率過程を回復することを目的としています。
特に、EOT 計画を考慮すると、これらの手順を適用して SB を解決できます。
この事実は、マッチングベースの SB ソルバーを生み出す最近の研究で大いに利用されています。
ここでの基礎は、EOT 計画を回復することです。最近の研究では、ヒューリスティックな近似 (ミニバッチ OT など) を使用するか、設計によりトレーニング中に誤差が蓄積する反復マッチング手順を確立しています。
私たちはこれらの制限に対処し、\textbf{最適な Schr\’odinger ブリッジ マッチング} と呼ぶ SB を学習するための新しい手順を提案します。
これは拡散プロセスの最適なパラメータ化を利用し、単一のブリッジ マッチング ステップで SB プロセス \textbf{(a)} を復元し、入力として任意のトランスポート プランを使用して \textbf{(b)} を証明します。
さらに、最適な橋梁マッチングの目標は、EOT/SB を学習するために最近発見されたエネルギーベース モデリング (EBM) の目標と一致することを示します。
この観察に触発されて、調整されたシュレディンガー ポテンシャルのガウス混合パラメータ化を使用して実際に最適なマッチングを実装するライト ソルバー (LightSB-M と呼ぶ) を開発しました。
さまざまな実際的なタスクにおけるソルバーのパフォーマンスを実験的に紹介します。
ソルバーのコードは 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 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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google