Generalized Schrödinger Bridge Matching

要約

拡散モデルやフローモデルを学習するための最新の分布マッチングアルゴリズムは、2つの境界分布間のマージナル分布の時間発展を直接規定する。本研究では、一般化された分布マッチングの設定を考える。そこでは、これらのマージン分布は、タスク固有の目的関数の解として暗黙的に記述されるだけである。この問題設定は、Generalized Schr’odinger Bridge (GSB)として知られ、機械学習内外の多くの科学分野で広く用いられている。我々は、運動エネルギー最小化を超えて一般化し、タスク固有の状態コストを考慮した、最近の進歩に触発された新しいマッチングアルゴリズムである一般化シュルアディンガーブリッジマッチング（GSBM）を提案する。このような一般化は、条件付き確率的最適制御を解くことであり、そのために効率的な変分近似を用いることができる。GSB問題を解く先行手法と比較して、我々のGSBMアルゴリズムは、学習中常に境界分布間の実現可能な輸送写像を保持することにより、安定した収束と著しく改善されたスケーラビリティを実現する。群衆ナビゲーション、意見脱分極、LiDAR多様体、画像領域移動を含む、広範な実験セットアップにおいて、我々の主張を実証的に検証する。我々の研究は、タスク固有の最適性構造で強化された拡散モデルを訓練するための新しいアルゴリズムの機会をもたらす。

要約(オリジナル)

Modern distribution matching algorithms for training diffusion or flow models directly prescribe the time evolution of the marginal distributions between two boundary distributions. In this work, we consider a generalized distribution matching setup, where these marginals are only implicitly described as a solution to some task-specific objective function. The problem setup, known as the Generalized Schr\’odinger Bridge (GSB), appears prevalently 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, generalizing them beyond kinetic energy minimization and to account for task-specific state costs. We show that such a generalization can be cast as solving conditional stochastic optimal control, for which efficient variational approximations can be used, and further debiased with the aid of path integral theory. Compared to prior methods for solving GSB problems, our GSBM algorithm always preserves a feasible transport map between the boundary distributions throughout training, thereby enabling stable convergence and significantly improved scalability. We empirically validate our claims on an extensive suite of experimental setups, including crowd navigation, opinion depolarization, LiDAR manifolds, and image domain transfer. Our work brings new algorithmic opportunities for training diffusion models enhanced with task-specific optimality structures.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL