(ICML2023) I2SB: Image-to-Image Schrödinger Bridge

Transcript

1. “Schrödinger Bridge”による画像変換 2025/4/23@論文読み会 Shumpei Takezaki (D2, Uchida Lab.)
2. • I2SB: Image-to-Image Schrödinger Bridge • 画像変換タスクのための新しい条件付き拡散モデルである Image-to-Image Schrödinger Bridgeの提案 紹介する論文 1 Inpainting JPED restoration ICML2023 に採択！
3. • あくまで画像を条件とした”画像生成”に過ぎない 従来の条件付き拡散モデルによる画像変換 2 𝜖𝜃 ×T 条件付け 構造情報などを損失 𝑡 = 1 𝑡 = 0 変換前 変換後
4. • Schrödinger Bridgeと拡散モデルを融合したモデルI2SBを提案 “直接”画像変換をおこなう拡散モデルがほしい！ 3 𝜖𝜃 ×T 𝑡 = 1 𝑡 = 0 変換前 変換後
5. • 拡散過程に基づく生成モデル 拡散モデル (DM) 4 𝑋0 ~𝑝0 (𝑋0 ) 𝑋1 ~𝒩(0, 𝐼) d𝑋𝑡 = 𝑓𝑡 𝑋𝑡 d𝑡 + 𝛽𝑡 d𝑊𝑡 d𝑋𝑡 = 𝑓𝑡 − 𝛽𝑡 𝛁 log 𝑝 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d ഥ 𝑊𝑡 SDEによる定式化 d𝑋𝑡 = 𝑓𝑡 𝑋𝑡 d𝑡 + 𝛽𝑡 d𝑊𝑡 d𝑋𝑡 = 𝑓𝑡 − 𝛽𝑡 𝛁 log 𝑝 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d ഥ 𝑊𝑡 モデル𝜖𝜃 のための損失関数 第１回論文読み回で紹介！ = 𝜖𝜃 𝑋𝑡 , 𝑡 − 𝑋0−𝑋𝑡 𝜎𝑡 𝜖𝜃 𝑋𝑡 , 𝑡 − 𝜎𝑡𝛁 log 𝑝 𝑋𝑡,𝑡|𝑋0 𝛁 log 𝑝 𝑋𝑡 , 𝑡|𝑋0 = 𝑋0 − 𝑋𝑡 𝜎𝑡 2 スコア関数𝛁 log 𝑝を求める スコア関数が解析的に書ける → NNで近似可能！
6. • エントロピー正則化された最適輸送モデル Schrödinger Bridge (SB) 5 𝑋0 ~𝑝𝒜 (𝑋0 ) 𝑋1 ~𝑝ℬ (𝑋0 ) d𝑋𝑡 = 𝑓𝑡 + 𝛽𝑡 𝛁 log Ψ 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d𝑊𝑡 SDEによる定式化 Ψ ෢ とΨを求めるためのPDE (Schrödinger System) d𝑋𝑡 = 𝑓𝑡 + 𝛽𝑡 𝛁 log Ψ 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d𝑊𝑡 d𝑋𝑡 = 𝑓𝑡 − 𝛽𝑡 𝛁 log ෡ Ψ 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d ഥ 𝑊𝑡 d𝑋𝑡 = 𝑓𝑡 − 𝛽𝑡 𝛁 log ෡ Ψ 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d ഥ 𝑊𝑡 ドリフト𝛁 log Ψ 𝑋𝑡 , 𝑡 と𝛁 log ෡ Ψ 𝑋𝑡 , 𝑡 を求める
7. • エントロピー正則化された最適輸送モデル Schrödinger Bridge (SB) 6 𝑋0 ~𝑝𝒜 (𝑋0 ) 𝑋1 ~𝑝ℬ (𝑋0 ) d𝑋𝑡 = 𝑓𝑡 + 𝛽𝑡 𝛁 log Ψ 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d𝑊𝑡 SDEによる定式化 d𝑋𝑡 = 𝑓𝑡 + 𝛽𝑡 𝛁 log Ψ 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d𝑊𝑡 d𝑋𝑡 = 𝑓𝑡 − 𝛽𝑡 𝛁 log ෡ Ψ 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d ഥ 𝑊𝑡 d𝑋𝑡 = 𝑓𝑡 − 𝛽𝑡 𝛁 log ෡ Ψ 𝑋𝑡 , 𝑡 d𝑡 + 𝛽𝑡 d ഥ 𝑊𝑡 𝚿と෡ 𝚿を求めるのは計算コストが高い （そもそも高次元データだと解けない) ドリフト𝛁 log Ψ 𝑋𝑡 , 𝑡 と𝛁 log ෡ Ψ 𝑋𝑡 , 𝑡 を求める Ψ ෢ とΨを求めるためのPDE (Schrödinger System)
8. • Schrödinger Bridgeを拡散モデルの枠組みで学習可能 SBを拡散モデルとして解きたい！ 7 式(9)に示すDMのSDEと境界条件を持たすスコア関数𝛁 log 𝑝 = SBのSDEを満たすドリフト𝛁 log Ψ 𝑋𝑡 , 𝑡 or 𝛁 log ෡ Ψ 𝑋𝑡 , 𝑡 何を意味するのか？ 1. スコア関数𝛁 log 𝑝を近似するための損失関数 𝜖𝜃 𝑋𝑡 , 𝑡 − 𝑋0−𝑋𝑡 𝜎𝑡 2. 定理3.1におけるスコア関数とドリフトの関係 スコア関数𝛁 log 𝑝＝ドリフト𝛁 log ෡ Ψ 𝑋𝑡 , 𝑡 𝝐𝜽 𝑿𝒕 , 𝒕 はドリフト𝛁 𝐥𝐨𝐠 ෡ 𝚿 𝑿𝒕 , 𝒕 を学習
9. I2SBの学習方法 8 変換前 𝑋1 変換後 𝑋0 𝜖𝜃 時刻𝑡のデータ 𝑋𝑡 𝜖𝜃 𝑋𝑡 , 𝑡 − 𝑋0−𝑋𝑡 𝜎𝑡 損失関数 → 拡散モデルの損失関数に相当 式(11) 𝑋0 と𝑋1 の加重和 (+ランダムノイズ) ※ 実装では削除可 ※ 𝜎𝑡 , ത 𝜎𝑡 は既知のパラメータ → 拡散モデルの拡散過程に相当
10. I2SBの推論方法 9 𝜖𝜃 時刻𝑡のデータ 𝑋𝑡 ※ 開始は変換前データ 𝑋1 𝜖𝜃 𝑋𝑡 , 𝑡 𝑋0 𝜖 = 𝑋𝑡 − 𝜎𝑡 𝜖𝜃 𝑋𝑡 , 𝑡 ※ 任意の推定回数 (NFE)を使用可 ※ 推定回数が多いほど精度も改善 推定データ𝑋0 𝜖 𝑋0 𝜖と𝑋1 の加重和 (+ランダムノイズ) 𝑋0 𝜖 𝑋0 𝜖 𝑋0 𝜖
11. • 画像復元タスク：劣化画像→高品質画像への変換 画像復元タスクにおける定量評価 10 Super-resolution JPEG restoration 従来手法より優れた精度を達成！ ※1 比較手法は主に拡散モデルベースの手法 ※2 の手法は追加情報が必要 ※3 他にもInpainting, Deblurringタスクでも検証 FID↓: 画像品質 CA↑: 画像分類の正解率
12. 効率的な画像変換が可能 11 Palette I2SB I2SB Palette 推定回数 (NFE)↓ FID↓ CA↑ 推定回数 (NFE)↓ 少ない推定回数で高精度 (=計算効率が良い)
13. 一般的な画像変換タスクへの応用も可能 12 定量的にも従来手法 (pix2pix)より優れた精度達成
14. • まとめ • 目的：画像変換タスクのための新しい条件付き拡散モデル • 手法：SBと拡散モデルを組み合わせたI2SBの提案 • 結果：従来の拡散モデルを凌駕する精度を達成 • 感想 • 理論的な難しさに比べて実装は簡単で画像変換タスクのベースラインとなりそう • Unpaired-I2Iへの拡張や高精度化 (主に理論的)で研究は進んでいる様子 • 精度も高く，生成過程に解釈性があるモデルなので応用研究に向いてそう • SBが最適輸送を解いているので解釈可能な変換になっているのかにも注目したい まとめと今後の展望 13