[論文レビュー] Likelihood Training of Schr\\"odinger Bridge using Forward-Backward SDEs Theory

Tianrong Chen, Guan-Horng Liu+1 more|arXiv (Cornell University)|Oct 21, 2021

Generative Adversarial Networks and Image Synthesis被引用数 11

ひとこと要約

本稿では、前向き・後向き確率微分方程式（FBSDE）理論を用いて、シュレーディンガー・ブリッジ（SB）モデルのための新しい尤度訓練フレームワーク、SB-FBSDEを提案する。SBの正確な対数尤度表現を導出し、線形拡散過程の特殊ケースとしてのSGM尤度目的関数を一般化することで、原理的かつ現代的な訓練手法を可能にした。評価時にランジュバン補正を必要とせず、MNIST、CelebA、CIFAR10で最先端のFIDスコアを達成した。

ABSTRACT

Schr\\"odinger Bridge (SB) is an entropy-regularized optimal transport problem that has received increasing attention in deep generative modeling for its mathematical flexibility compared to the Scored-based Generative Model (SGM). However, it remains unclear whether the optimization principle of SB relates to the modern training of deep generative models, which often rely on constructing log-likelihood objectives.This raises questions on the suitability of SB models as a principled alternative for generative applications. In this work, we present a novel computational framework for likelihood training of SB models grounded on Forward-Backward Stochastic Differential Equations Theory - a mathematical methodology appeared in stochastic optimal control that transforms the optimality condition of SB into a set of SDEs. Crucially, these SDEs can be used to construct the likelihood objectives for SB that, surprisingly, generalizes the ones for SGM as special cases. This leads to a new optimization principle that inherits the same SB optimality yet without losing applications of modern generative training techniques, and we show that the resulting training algorithm achieves comparable results on generating realistic images on MNIST, CelebA, and CIFAR10. Our code is available at https://github.com/ghliu/SB-FBSDE.

研究の動機と目的

• シュレーディンガー・ブリッジ（SB）最適化と、対数尤度目的関数に依存する現代的な深層生成モデル、特にスコアベース生成モデル（SGM）との間の原理的つながりを確立すること。
• SBモデルがSGMの理論的裏付けのある代替手段として機能可能かどうかの曖昧さを解消し、SBの訓練を尤度フレームワークに根ざさせること。
• FBSDE理論を用いて、SBモデルの尤度に基づく訓練を可能にする計算フレームワークを構築し、従来の手作業的またはIPFに基づく最適化の制限を克服すること。
• SGM尤度目的関数を完全非線形拡散過程に一般化することで、SBとSGMを共通の尤度原理の下に統一すること。
• 本尤度フレームワークを用いて訓練されたSBモデルが、標準ベンチマークで競争的または優れた画像生成性能を達成することを示すこと。

提案手法

• 本手法は、確率的最適制御理論を用いて、PDEに基づく最適性条件を扱いやすいSDEに変換することで、SBの最適性条件を前向き・後向きSDE（FBSDE）の連立系として導出する。
• FBSDE連立方程式を解くことで、SBの正確な対数尤度表現を構築し、拡散が線形の場合のSGM尤度目的関数の特殊ケースとして一般化される。
• 従来の平均マッチング回帰手法に代わり、発散に基づく目的関数（例：KL発散）を用いて前向きおよび後向き過程の訓練を行う。
• 標準的なディープラーニング技術を用いて、前向きおよび後向きSDEの共同最適化を可能にし、SGMにインspiredされたランジュバン補正器を refine に用いる。
• 導出された尤度目的関数を用いて、前向きおよび後向き方策の最適化を交互に実行することで、安定的かつスケーラブルな訓練を実現する。
• 本手法はPyTorchベースのコードベースに実装され、U-Net や NCSN++ といった標準アーキテクチャを用いて、MNIST、CelebA、CIFAR10で評価された。

実験結果

リサーチクエスチョン

• RQ1 シュレーディンガー・ブリッジは、SGM尤度フレームワークを一般化する尤度目的関数を用いて訓練可能か？
• RQ2 SB最適性のFBSDEに基づく定式化は、現代的な深層生成モデル技術との間でどのように原理的つながりを確立するか？
• RQ3 FBSDEを用いたSBの尤度訓練は、平均マッチング回帰に基づく従来のSB手法と比較して、サンプルの質および多様性が向上するか？
• RQ4 導出された尤度目的関数は、ランジュバン補正器のような現代的な訓練改善手法をSBモデルに統合するために使用可能か？
• RQ5 標準的な画像生成ベンチマーク（MNIST、CelebA、CIFAR10）における、尤度訓練されたSBモデルの定量的性能はいかほどか？

主な発見

• SB-FBSDEフレームワークは、シュレーディンガー・ブリッジの正確な対数尤度表現を導出し、これはSGM尤度目的関数の特殊ケースとして一般化される。これにより、両モデルクラスの間の根本的な理論的つながりが確立された。
• CIFAR10では、10,000枚の生成サンプルを用いて、Fréchet Inception Distance（FID）が11.85に達し、従来の最適輸送モデルを著しく上回った。
• SGM回帰の50,000ステップに加え、5,000ステップの発散ベース訓練を実施した後でも、FIDは33.68から13.35に低下し、急速な性能向上が確認された。
• 定性的な結果では、特にCelebAにおける顔貌や背景の多様性において、従来のSB手法よりも高い画像多様性と視覚的質を示した。
• ランジュワン補正器をFID評価で使用せずとも、MNIST、CelebA、CIFAR10で最先端の性能を達成しており、優れた内在的サンプル品質を示している。
• 本フレームワークにより、ランジュワン補正器のような現代的な訓練技術の利用が可能となり、アーキテクチャやデータセットサイズにかかわらず一貫した性能向上が得られた。