拡散過程の時間反転

概要

拡散過程の時間反転を考えると，Hyvärinen スコアがドリフト項に現れる．デノイジングスコアマッチングの目的関数は，これを利用してデータ分布のスコアが推定されるように設計されている．Tweedie の式がこれを正当化するが，この式を用いたサンプリング手法には確率的局所化というものもある．

1 命題

(命題 Haussmann and Pardoux, 1986)

Brown 運動 \(\{B_t\}\subset\mathcal{L}(\Omega;\mathbb{R}^d)\) と可測関数 \(b:[0,1]\times\mathbb{R}^d\to\mathbb{R}^d,\sigma:[0,1]\times\mathbb{R}^d\to M_d(\mathbb{R})\) に関して， \[ dX_t=b_t(X_t)\,dt+\sigma_t(X_t)\,dB_t,\qquad t\in[0,1], \] を Markov 過程とする．さらに次の３条件を仮定する：

\(b_t,\sigma_t\) は \(\mathbb{R}^d\) 上局所 Lipschitz 連続で，線型増大条件を満たす： \[ \lvert b_t(x)\rvert+\lvert\sigma_t(x)\rvert\le K(1+\lvert x\rvert),\qquad x\in\mathbb{R}^d,K>0. \]
\(X_0\) は密度 \(p_0\) をもち，ある \(\lambda<0\) について次を満たす： \[ p_0\in L^2((1+\lvert x\rvert^2)^\lambda\,dx). \]
\(a:=\sigma\sigma^\top\) は一様に正定値である \[ a_t(x)\ge\alpha I_d \] であるか，\(\alpha^{ij}_{x_ix_j}\in L^\infty((0,1)\times\mathbb{R}^d)\) である．

このとき，時間反転 \(\overline{X}_t:=X_{1-t}\) の分布は次の SDE の解である： \[ d\overline{X}_t=\overline{b}_t(\overline{X}_t)\,dt+\overline{\sigma}_t(\overline{X}_t)\,d\overline{B}_t,\qquad t\in[0,1]. \] ただし，\((B_t)\) も Brown 運動で， \[ \overline{b}^i_t(x)=-b_{1-t}^i(x)+\sum_{j=1}^d\frac{\biggr(a^{ij}_{1-t}(x)p_{1-t}(x)\biggl)_{x_j}}{p_{1-t}(x)}, \] \[ \overline{a}^{ij}_t(x)=a^{ij}_{1-t}(x),\qquad\overline{\sigma}^{ij}_t(x)=\sigma^{ij}_{1-t}(x). \]

系

前の命題の３条件を満たす，ドリフト係数 \(\sigma\) が \(x\in\mathbb{R}^d\) に依らない SDE \[ dX_t=b_t(X_t)\,dt+\sigma_t\,dB_t,\qquad t\in[0,1], \] を考える．この \((X)_{t\in[0,1]}\) の時間反転は，\(a_t:=\sigma_t\sigma^\top_t\) に関して \[ d\overline{X}_t=\biggr(-b_{1-t}(\overline{X}_t)+a_{1-t}\nabla_x\log p_{1-t}(\overline{X}_t)\biggl)\,dt+\sigma_{1-t}\,d\overline{B}_t,\qquad\overline{X}_0=X_1, \] と分布同等になる．

SGM (Song and Ermon, 2019), (Song et al., 2021) は， \[ b_t(x)=-x,\qquad\sigma_t=\sqrt{2}, \tag{1}\] と設定し，\((X_t)\) を OU 過程とした．これは \(\mathrm{N}_d(0,I_d)\) に全変動距離・Wasserstein 距離・\(\chi^2\)-距離で指数収束する．

従って，この時間反転を \(\mathrm{N}_d(0,I_d)\) からスタートさせることで，データ分布 \(p_0\) からのサンプリングが可能になる．

しかしこのアイデアを実行するためには，スコア \(\nabla_x\log p_{1-t}(X_t)\) の項を何らかの方法で推定する方法が必要である．

2 スコアマッチングへの応用

2.1 はじめに

Denoising Score Matching (Vincent, 2011) を初めとして，Generalized Flow Matching (Isobe et al., 2024) や Functional Flow Matching (Kerrigan et al., 2024) は，次のような目的関数を持っている： \[ \mathcal{L}(\theta)=\operatorname{E}\biggl[\biggl|s_\theta(\widetilde{X})-\frac{X-\widetilde{X}}{\sigma^2}\biggr|^2\biggr],\qquad X\sim p_0,\widetilde{X}\sim p_0*\mathrm{N}(0,\sigma^2I_d). \tag{2}\]

これはデータ \(X\sim p_0\) と，それに独立な Gauss ノイズを印加したもの \(\widetilde{X}\) との差分を目標としてスコアネットワーク \(s_\theta\) を学習している．

2.2 デノイジング過程としての見方

データにノイズを印加する過程は，\(b_t=0,\sigma_t=I_d\) とした SDE \[ dX_t=dB_t,\qquad t\in[0,1],X_0\sim p_0, \] で \(t=0\) から \(t=\sigma^2\) まで輸送することにあたる： \[ X_\sigma^2=X_0+(B_{\sigma^2}-B_0)\overset{\text{d}}{=}\widetilde{X}. \] この時間反転は \[ d\overline{X}_t=\nabla_x\log p_{1-t}(\overline{X}_t)\,dt+d\overline{B}_t,\qquad\overline{X}_0=X_1, \] と分布同等になる．

この時間反転過程 \((\overline{X}_t)\) は \(\overline{X}_{1-\sigma^2}=\widetilde{X}\) を \(\overline{X}_1=X\) まで運ぶが，この際に \(\sigma\ll1\) ならば次の関係を示唆する： \[\begin{align*} X&\overset{\text{d}}{=}\widetilde{X}+\int^1_{1-\sigma^2}\nabla_x\log p_{1-t}(\overline{X}_t)\,dt+\epsilon\\ &\approx\widetilde{X}+\sigma^2\nabla_x\log p_0(X)+\epsilon,\qquad\epsilon\sim\mathrm{N}(0,\sigma^2I_d). \end{align*}\] \(\sigma\to0\) の極限で次の等号が成り立つ： \[ \lim_{\sigma\to0}\frac{X-\widetilde{X}}{\sigma^2}\overset{\text{d}}{=}\nabla_x\log p_0(X). \]

2.3 Tweedie の式

実は同様の式は，\(\sigma\to0\) の極限で漸近的にではなく正の \(\sigma^2>0\) に関しても，次の意味で成り立つ：

命題

\(X\sim p_0,\widetilde{X}\sim p_0*\mathrm{N}(0,\sigma^2I_d)=:\widetilde{p}_0\) とする．このとき，次が成り立つ： \[ \operatorname{E}[X|\widetilde{X}=\widetilde{x}]=\widetilde{x}+\sigma^2\nabla_{\widetilde{x}}\log \widetilde{p}_0(\widetilde{x}). \]

証明

一般に，\(\phi_\sigma\) を \(\mathrm{N}_d(0,\sigma^2I_d)\) の密度関数とすると， \[ \operatorname{E}[X|\widetilde{X}=\widetilde{x}]=\int_{\mathbb{R}^d}xp_0(x)\phi_\sigma(\widetilde{x}-x)\,dx \] より， \[ \operatorname{E}[X|\widetilde{X}=\widetilde{x}]-\widetilde{x}=\int_{\mathbb{R}^d}(x-\widetilde{x})p_0(x)\phi_\sigma(\widetilde{x}-x)\,dx. \]

一方で， \[\begin{align*} &\qquad\nabla_{\widetilde{x}}\log\left(\int_{\mathbb{R}^d}p_0(x)\phi_\sigma(\widetilde{x}-x)\,dx\right)\\ &=\frac{\int_{\mathbb{R}^d}p_0(x)\nabla_{\widetilde{x}}\phi_\sigma(\widetilde{x}-x)\,dx}{\int_{\mathbb{R}^d}p_0(x)\phi_\sigma(\widetilde{x}-x)\,dx}\\ &=\int_{\mathbb{R}^d}p_0(x)\phi_\sigma(\widetilde{x}-x)\frac{x-\widetilde{x}}{\sigma^2}\,dx\\ &=\frac{\operatorname{E}[X|\widetilde{X}=\widetilde{x}]-\widetilde{x}}{\sigma^2}. \end{align*}\]

すなわち，\(\widetilde{X}\) から \(X\) を不偏推定しようとすることで，スコア \(\nabla_{\widetilde{x}}\log\widetilde{p}(\widetilde{x})\) を学習することができるのである．

ただし，学習されるスコアは，データ分布 \(p_0\) のものではなく，ノイズ分布 \(\widetilde{p}_0\) のものであることに注意．

これが，デノイジングスコアマッチングの目的関数 (2) の背後にある動機付けである．

3 確率的局所化

3.1 OU 過程による SGM

OU 過程の例 (1) に戻ろう．OU 過程 \[ dX_t=-X_t\,dt+\sqrt{2}\,dB_t \] の時間反転は次と分布同等である： \[ d\overline{X}_t=\biggr(\overline{X}_t+2\nabla_x\log p_{1-t}(\overline{X}_t)\biggl)\,dt+\sqrt{2}\,d\overline{B}_t,\qquad\overline{X}_0=X_1. \tag{3}\]

OU 過程 \((X_t)\) は \[ X_t\overset{\text{d}}{=}e^{-t}X_0+\sqrt{1-e^{-2t}}\epsilon,\qquad X_0\sim p_0,\epsilon\sim\mathrm{N}_d(0,I_d) \] という遷移半群を持っているため，\(\mathcal{L}[X_t]=\mathcal{L}[e^{-t}X_0]*\mathrm{N}_d(0,1-e^{-2t})\) であることから，Tweedie の式 2.3 より \[ \nabla_x\log p_t(x_t)=\frac{\operatorname{E}[e^{-t}X_0|X_t=x_t]-x_t}{1-e^{-2t}} \] を得る．従ってこのスコアを式 (3) に代入し， \[ m_t(x_t):=\operatorname{E}[X_0|tX_0+\sqrt{t}\epsilon=x_t],\qquad X_0\sim p_0,\epsilon\sim\mathrm{N}_d(0,I_d), \] とおき， \[ \tau(t)=\frac{1}{e^{2t}-1} \] の変数変換を施すと OU 過程の時間反転 (3) は次のように書き直せる： \[ d\overline{Y}_\tau=\biggr(-\frac{1+\tau}{\tau(1+\tau)}\overline{Y}_\tau+\frac{1}{\sqrt{\tau(1+\tau)}}m_\tau\left(\sqrt{\tau(1+\tau)}\overline{Y}_\tau\right)\biggl)\,d\tau+\frac{1}{\sqrt{\tau(1+\tau)}}\,d\overline{B}_\tau. \tag{4}\]

これが denoising diffusion である．

3.2 もう一つのサンプリング法

確率的局所化 (stochastic localization) は初め，(Eldan, 2013) が高次元空間内の等方的な凸体上での等周不等式を示すために構成した半マルチンゲールが基になっている．

確率的局所化においては，\(p_0\) からのあるサンプル \(x_0\) に対して，その 観測過程 \((Y_t)\) と呼ばれる \(x_0\) のノイズ付きの観測を考える．ただし，\(Y_t\) は時間が進むごとに \(x_0\) に関する情報量が増えるとする．¹

例えば \[ Y_t=tx_0+B_t,\qquad t\in\mathbb{R}_+, \] という場合である．\(B_t\) というノイズは印加されているが，\(x_0\) というメッセージの内容がどんどん大きくなるため，Signal-to-noise 比は増大していく．

この場合については，\(p_0\) が有限な二次の積率を持つならば， \[ dY_\tau=m_\tau(Y_\tau)\,d\tau+dB'_\tau \] という SDE の解と分布同等である (Liptser and Shiryaev, 2001)．

これは式 (4) で与えた OU 過程の時間反転 \((\overline{Y}_\tau)\) に関して，\(Y_\tau=\sqrt{\tau(1+\tau)}\overline{Y}_\tau\) の関係を持つ．

3.3 確率的局所化

実は \((Y_t)\) は， \[ \mu_t:=\mathcal{L}[X_0|Y_t] \] として定まる \(\operatorname{P}(\mathbb{R}^d)\)-値の確率過程 \(\{\mu_t\}\subset\mathcal{L}(\Omega;\mathcal{P}(\mathbb{R}^d))\) について，次の性質を持つ： \[ \mu_0=\mathcal{L}[X_0]=p_0\,d\ell_d,\qquad\mu_t\Rightarrow\delta_{x_0}\qquad t\to\infty, \]

実は上述のサンプリング法は，このような \(p_0\,d\ell_d\) から \(\delta_{X_0}\;(X_0\sim p_0)\) への確率過程が与えられるごとに構成できる．

実際，最も簡単には，重心 \[ M_t:=\int_{\mathbb{R}^d}x\,\mu_t(dx) \] を計算すれば，\(M_t\) は \(t\to\infty\) の極限で \(\mathcal{L}[X_0]\) に収束する．

これが 確率的局所化 である．

確率的局所化に基づいたサンプラーは (Alaoui et al., 2022) により Sherrington-Kirkpatrick 模型の Gibbs 分布からのサンプリングに適用され，(Montanari and Wu, 2023) でさらにベイズ統計への応用のために拡張されている．

また，最良の雑音除去拡散モデルの収束証明は確率的局所化に基づいた証明によって与えられている (Benton et al., 2024)．

4 文献紹介

(Anderson, 1982) では Fokker-Planck 方程式の解に対する条件の言葉で時間反転命題を与えている．また，時間反転も，元の Brown 運動 \(B_t\) と独立な Brown 運動 \(\overline{B}_t\) に関する SDE ではなく，その時間反転 \(\widetilde{B}_t:=B_{1-t}\) に関する SDE で与えている．

Aleandre Thiéry のブログ記事や鈴木大慈氏のスライド，Montanari の講義資料も参照．

Tweedie の式は (Robbins, 1956) によって命名されている．(Efron, 2011) では選択バイアスが存在する状況における経験ベイズ法に応用している．

確率的局所化については (Montanari, 2023) を参考にした．

References

Alaoui, A. E., Montanari, A., and Sellke, M. (2022). Sampling from the sherrington-kirkpatrick gibbs measure via algorithmic stochastic localization. In 2022 IEEE 63rd annual symposium on foundations of computer science (FOCS), pages 323–334.

Anderson, B. D. O. (1982). Reverse-time diffusion equation models. Stochastic Processes and Their Applications, 12(3), 313–326.

Benton, J., Bortoli, V. D., Doucet, A., and Deligiannidis, G. (2024). Nearly \(d\)-linear convergence bounds for diffusion models via stochastic localization.

Efron, B. (2011). Tweedie’s formula and selection bias. Journal of the American Statistical Association, 106(496), 1602–1614.

Eldan, R. (2013). Thin Shell Implies Spectral Gap Up to Polylog via a Stochastic Localization Scheme. Geometric and Functional Analysis, 23(2), 532–569.

Haussmann, U. G., and Pardoux, E. (1986). Time Reversal of Diffusions. The Annals of Probability, 14(4), 1188–1205.

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

Kerrigan, G., Migliorini, G., and Smyth, P. (2024). Functional flow matching. In S. Dasgupta, S. Mandt, and Y. Li, editors, Proceedings of the 27th international conference on artificial intelligence and statistics,Vol. 238, pages 3934–3942. PMLR.

Liptser, R. S., and Shiryaev, A. N. (2001). Statistics of random processes i: General theory. Original Russian edition published by Nauka, Moscow, 1974; Springer Berlin, Heidelberg.

Montanari, A. (2023). Sampling, diffusions, and stochastic localization.

Montanari, A., and Wu, Y. (2023). Posterior sampling from the spiked models via diffusion processes.

Robbins, H. (1956). An Empirical Bayes Approach to Statistics. In Proceedings of the third berkeley symposium on mathematical statistics and probability,Vol. 1, pages 157–163.

Song, Y., and Ermon, S. (2019). Generative Modeling by Estimating Gradients of the Data Distribution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. dAlché-Buc, E. Fox, and R. Garnett, editors, Advances in neural information processing systems,Vol. 32. Curran Associates, Inc.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021). Score-Based Generative Modeling through Stochastic Differential Equations. In International conference on learning representations.

Vincent, P. (2011). A connection between score matching and denoising autoencoders. Neural Computation, 23(7), 1661–1674.

Footnotes

例えば，\(x_*,Y_{t_2},Y_{t_1}\) が長さ３の Markov 連鎖をなす，などの意味で．↩︎