拡散モデルからシュレディンガー橋へ

関連記事一覧

A Blog Entry on Bayesian Computation by an Applied Mathematician

$$

$$

1 Schrödinger 橋としての条件付き生成

1.1 はじめに

雑音除去過程（OU 過程の時間反転）を事後分布 \(p(x|y)\) からの正確なサンプリングに用いるには，OU 過程の遷移密度 \((p_t(x_t|y))_{t\in[0,T]}\) の終端分布に関する近似的な関係 \[ p_T(x_T|y)\approx\mathrm{N}_d(0,I_d) \] を正確に成り立たせる必要がある．

これは OU 過程の代わりに Schrödinger 橋を用いることに相当する．

Schrödinger 橋を用いたサンプラーは (Bernton et al., 2019), (De Bortoli et al., 2021), (Shi et al., 2022) で提案された．

1.2 定義

前述の OU 過程の分布を \(\mathbb{P}_y\in\mathcal{P}(C([0,T];\mathcal{X}))\) と表そう．OU 過程がエルゴード性を持とうとも，このままでは \((\mathbb{P}_y)_t=\mathrm{N}_d(0,I_d)\) は近似的にしか成り立たない．

Schrödinger 橋 (SB) とは，\(\mathbb{P}:=\mathbb{P}_{y}\otimes\delta_{p(y)}\) から最も KL 乖離度が小さい分布 \[ \Pi^*:=\operatorname*{argmin}_{\Pi\in\mathcal{A}}\operatorname{KL}(\Pi,\mathbb{P}), \tag{1}\] \[ \mathcal{A}:=\left\{\Pi\in\mathcal{P}(C([0,T];\mathcal{X}\times\mathcal{Y}))\,\middle|\,\begin{array}{l}\Pi_0\sim p(x,y)dxdy,\\\Pi_T=\mathrm{N}_d(0,I_d)\otimes p(y)dy\end{array}\right\}, \] に従う確率過程をいう．

ただし，\(\delta_{p(y)}:=p(y)dy\otimes\mathrm{U}([0,T])\) は \(X_0\sim p(y)dy\) を決めたのちに動かない過程 \(X_t\equiv X_0\) が \(C([0,T];\mathcal{Y})\) 上に定める確率分布とした．

1.3 SB を使ったサンプリング

SB 問題 (1) の解は表示 \[ \Pi^*=\mathbb{P}^*_{y}\otimes\delta_{p(y)} \] を持つ．

そのため，仮に \(\mathbb{P}^*_y\) に従う過程の時間反転 \((Z_t^y)\) がシミュレーションできたならば，ランダム性の種 \(Z_T^y\sim\mathrm{N}_d(0,I_d)\) とラベル \(y\) を初期条件として \((Z_t^y)\) で流すことで， \[ Z_0^y\sim p(x|y)dx,\qquad p(y)\text{-a.s.} \] を得ることになる．

問題は \(\Pi^*,\mathbb{P}^*_y\) を計算し，これそシミュレーションすることである．

2 シミュレーション法

SB 問題 (1) の解 \(\Pi^*\) は 逐次的比例フィッティング (IPF: Iterative Proportional Fitting) により得られる．

2.1 IPF とは

IPF の歴史

IPF アルゴリズムと method of iterative propotions の名前は離散的な形で (Deming and Stephan, 1940) が分割表データ解析の研究で提案している．

その手続きを (Ireland and Kullback, 1968) が距離の最小化として特徴付け，(Kullback, 1968) が確率密度に対しても一般化した．

ただし，この確率密度に対するアルゴリズムは (Fortet, 1940) が Schrödinger 方程式の研究ですでに提案しているものである．

IPF は元々，指定した２つの確率ベクトル \(r\in(0,\infty)^{d_r},c\in(0,\infty)^{d_c}\) を周辺分布に持つ結合分布（カップリング）のうち，指定の行列 \(W\in M_{d_rd_c}(\mathbb{R}_+)\) に最も近い KL 乖離度を持つカップリングを見つけるための逐次アルゴリズムである (Kurras, 2015)．

種々の分野で再発見され，複数の名前を持っているようである．少なくとも，Sheleikhovskii 法，Kruithof アルゴリズム，Furness 法，Sinkhorn-Knopp アルゴリズム，RAS 法など (Kurras, 2015)．¹

\(W\) の成分が正である場合は，(Sinkhorn, 1967) がアルゴリズムの収束と解の一意性を示している．²

しかし，\(W\) の成分が零を含む場合，零成分の位置に依存してアルゴリズムは収束しないことがあり得ることを (Sinkhorn and Knopp, 1967) が \(d_r=d_c=1\) の場合について示している．

2.2 アルゴリズム

IPF アルゴリズムは観念的には，OU 過程の分布 \(\mathbb{P}=\mathbb{P}_y\otimes\delta_{p(y)}\) からはじめ，２つの周辺分布制約 \[ \Pi_0(x_0,y_0)=p(x_0,y_0)=:\Pi_{\text{data}}(x_0,y_0) \] \[ \Pi_T=\mathrm{N}_d(0,I_d)\otimes p(y_T)dy_T=:\Pi_{\text{prior}} \] のうち片方のみを満たすもののうち，KL 距離を最小にする解への射影を返していく： \[ \Pi^{2n+1}:=\operatorname*{argmin}_{\Pi\in\mathcal{P}(C([0,T];\mathcal{X}\times\mathcal{Y}))}\biggl\{\operatorname{KL}(\Pi,\Pi^{2n})\,\bigg|\,\Pi_T=\Pi_{\text{prior}}\biggr\}, \] \[ \Pi^{2n+2}:=\operatorname*{argmin}_{\Pi\in\mathcal{P}(C([0,T];\mathcal{X}\times\mathcal{Y}))}\biggl\{\operatorname{KL}(\Pi,\Pi^{2n+1})\,\bigg|\,\Pi_0=\Pi_{\text{data}}\biggr\}. \]

今回の場合， \[ \Pi^{2n+1}=\mathbb{P}_{y_T}^{2n+1}\otimes\delta_{p(y)},\qquad\Pi^{2n+2}=\mathbb{P}^{2n+2}_{y_0}\otimes\delta_{p(y)}, \] と分解される．

ただし奇数回実施後の \(\mathbb{P}_{y_T}^{2n+1}\) は次で定まるノイズをデータにする過程 \((Z_t)\) の時間反転とした： \[ dZ_t=f_{T-t}^{2n+1}(Z_t,y_T)\,dt+dW_t,\qquad Z_0\sim\mathrm{N}_d(0,I_d),f_t^{2n+1}(x_t,y):=-f_t^{2n}(x_t,y)+\nabla_{x}\log\Pi^{2n}_t(x_t|y), \] 一方で偶数回実施後の \(\mathbb{P}_{y_0}^{2n+2}\) は次で定まるデータをノイズにする過程 \((X_t)\) の経路測度となる： \[ dX_t=f^{2n+2}_t(X_t,y_0)\,dt+dB_t,\qquad X_0\sim p(x|y_0)\,dx,f_t^{2n+2}(x_t,y):=-f_t^{2n+1}(x_t,y)+\nabla_x\log\Pi_t^{2n+1}(x_t|y). \]

ただし初期分布は OU 過程 \(\mathbb{P}_y\) とするから，\(f^0_t(x_t)=-x_t/2\)．

つまり IPF は OU 過程が時刻 \(T\) で正確に不変分布に到達しているようにするために，ドリフト項 \(f_t\) を逐次的に修正する過程と見れる．

2.3 IPF の近似方法

最初のイテレーション \(n=0\) における \(\mathbb{P}^1_y\) が雑音除去拡散 DD に対応する．

DD をスコアマッチングによって学習した ように，\(\mathbb{P}^2_y,\mathbb{P}^3_y,\cdots\) と逐次的にスコアマッチングを実行することが考えられる．

この際には mean-matching (De Bortoli et al., 2021), (Shi et al., 2022) に基づく効率的なアルゴリズムが与えられている．

文献紹介

SB の応用

I²SB (Liu et al., 2023) のプロジェクトページは こちら．

Text-to-Speech にも応用されている (Chen et al., 2023)．

エントロピー最適輸送との関係

1. Langevin 動力学は，平衡分布との KL 乖離度を最小化する Wasserstein 勾配流になっている (Jordan et al., 1998), (4.4節 Figalli and Glaudo, 2023)．
2. その 時間反転過程 は OU 過程からの KL 乖離度とエネルギーの和を，\(\overline{Y}_T\sim p_0\) の境界条件の下最小化する過程になっている．
3. さらに境界条件を加えたものが SB である．

鈴木大慈氏のスライド も参照．

機械学習の観点からは，エントロピー正則化項を，帰納バイアスを入れる方法として再解釈することができる (Koshizuka and Sato, 2023), (Isobe et al., 2024)．

References

Bedford, T., Daneshkhah, A., and Wilson, K. J. (2016). Approximate uncertainty modeling in risk analysis with vine copulas. Risk Analysis, 36(4), 792–815.

Bernton, E., Heng, J., Doucet, A., and Jacob, P. E. (2019). Schrödinger bridge samplers.

Chen, Z., He, G., Zheng, K., Tan, X., and Zhu, J. (2023). Schrodinger bridges beat diffusion models on text-to-speech synthesis.

De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. (2021). Diffusion schrödinger bridge with applications to score-based generative modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. S. Liang, and J. W. Vaughan, editors, Advances in neural information processing systems,Vol. 34, pages 17695–17709. Curran Associates, Inc.

Deming, W. E., and Stephan, F. F. (1940). On a least squares adjustment of a sampled frequency table when the expected marginal totals are known. The Annals of Mathematical Statistics, 11(4), 427–444.

Figalli, A., and Glaudo, F. (2023). An Invitation to Optimal Transport, Wasserstein Distances, and Gradient Flows. European Mathematical Society.

Fortet, R. (1940). Résolution d’un système d’équations de M. Schrödinger. Journal de Mathématiques Pures Et Appliquées, Series 9, 19(1-4), 83–105.

Ireland, C. T., and Kullback, S. (1968). Contingency tables with given marginals. Biometrika, 55(1), 179–188.

Isobe, N., Koyama, M., Zhang, J., Hayashi, K., and Fukumizu, K. (2024). Extended flow matching: A method of conditional generation with generalized continuity equation.

Jordan, R., Kinderlehrer, D., and Otto, F. (1998). The variational formulation of the fokker–planck equation. SIAM Journal on Mathematical Analysis, 29(1), 1–17.

Koshizuka, T., and Sato, I. (2023). Neural lagrangian schr\”{o}dinger bridge: Diffusion modeling for population dynamics. In The eleventh international conference on learning representations.

Kullback, S. (1968). Probability densities with given marginals. The Annals of Mathematical Statistics, 39(4), 1236–1243.

Kurras, S. (2015). Symmetric Iterative Proportional Fitting. In G. Lebanon and S. V. N. Vishwanathan, editors, Proceedings of the eighteenth international conference on artificial intelligence and statistics,Vol. 38, pages 526–534. San Diego, California, USA: PMLR.

Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, E., Nie, W., and Anandkumar, A. (2023). I\(^2\)SB: Image-to-image schrödinger bridge. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, Proceedings of the 40th international conference on machine learning,Vol. 202, pages 22042–22062. PMLR.

Shi, Y., De Bortoli, V., Deligiannidis, G., and Doucet, A. (2022). Conditional simulation using diffusion Schrödinger bridges. In J. Cussens and K. Zhang, editors, Proceedings of the thirty-eighth conference on uncertainty in artificial intelligence,Vol. 180, pages 1792–1802. PMLR.

Sinkhorn, R. (1967). Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly, 74(4), 402–405.

Sinkhorn, R., and Knopp, P. (1967). Concerning Nonnegative Matrices and Doubly Stochastic Matrices. Pacific Journal of Mathematics, 21(2), 343–348.

清智也. (2021). 最小情報コピュラとその周辺. 日本統計学会誌, 51(1), 75–99.

Footnotes

1. また，行列スケーリングを通じた最小情報コピュラとの関連を (Bedford et al., 2016), (清智也, 2021) が指摘している．

2. ただし (Sinkhorn, 1967) は (Deming and Stephan, 1940) にも (Fortet, 1940) にも言及しておらず，Markov 連鎖の遷移確率の推定という文脈で研究している．