dissuion modelの勉強してたらGauss分布 $\mathcal{N}(x; x_0, \sigma)$ が出てきていたので、そのための式をメモ代わりに置いておきます。$$
\mathcal{N}(x; x_0, \sigma^2) = \frac{1}{\sigma\sqrt{2\pi}}e^{-\frac{(x-x_0)^2}{2\sigma^2}}
$$
なんでGauss分布？
diffusion modelではガウス分布を用いた拡散過程が出てきます。そこでまずはなぜガウス分布が拡散過程を表しているのかを簡単に見てみましょう。
AI のみに興味がある方は飛ばしても問題ないです。
拡散方程式のお話
物理学における有名な方程式の一つに、拡散方程式と呼ばれるものがあります。その導入として、系が粒子数、電荷、エネルギーといった保存量 $Q$ を持つ場合を考えましょう。保存量とはその総量が時間によって変化しない量を指し、$dQ/dt = 0$で特徴付けられます。このような保存量と系の対称性は密接に関係しているのですが、そういうのが気になる方は物理学の本を読むなりNoetherの定理でググるなりして下さい。
保存量に付随する位置 $x$ 時刻 $t$ での密度 $q(x, t), \int dx\,q(x, t) = Q$ を考えると、以下の式を満たす場合があります。$$\frac{\partial q(x, t)}{\partial t} = \kappa\frac{\partial^2 q(x, t)}{\partial x^2}$$ただし、$\kappa>0$ は拡散のし易さを司るパラメータで、拡散係数と呼ばれます。また、多次元に拡張したい場合は $x$ の2階偏微分をラプラシアンに置き換えてください。これから空間1次元の場合の拡散方程式を解いていきますが、全く同じ手順で解けます。
保存量が存在する場合、連続の式が成り立ちます。 $$\frac{\partial q}{\partial t} + \frac{\partial j}{\partial x} = 0.$$ ここで $j(x, t)$ は保存量の流れ（カレント）の密度です。$Q$が電荷の場合、$j$ は電流密度ですね。このカレント密度ですが、仮に $\partial q/\partial x$ に比例したとしましょう。すると拡散方程式が得られることがすぐにわかります。
残念なことに、保存量が存在すれば必ず拡散方程式が成り立つなんてことはありません。成り立つ場合の典型的例としては、熱の拡散や水中のインクの拡散です。この場合の保存量は熱エネルギーとインクの分子の数です。しかしながらこれらは厳密には保存量ではありません。前者は他の形態のエネルギーとの変換を考慮していないですし、後者はインク分子の分解などを考慮していません。ただし、特定の系や時間スケールなどに限定すれば近似的に保存量と見なせるはずです。というわけで、拡散方程式は厳密な意味での保存量以外にも適応可能なので、様々な場所で顔を出します。
また、拡散方程式はマスター方程式にいくつかの近似を施すことで導出することもできます。興味のあるかたはFokker-Plank方程式やKramers-Moyal展開のあたりでググってください。
拡散方程式を解いてみる
一般に変微分方程式を解く場合は、境界条件が必要になります。そこで時刻 $t=t_0$で $q(x, t_0) = q_0(x)$ であったとしましょう。拡散方程式はFourier変換を利用すると簡単に解けます。$$q(x, t) = \int_{-\infty}^{\infty} \frac{dk}{2\pi}\tilde{q}(k, t)e^{ikx}.$$このようにして導入された $\tilde{q}(k, t)$ は $q(x, t)$ と一対一の関係にあり、逆となる変換$$\tilde{q}(k, t) = \int_{-\infty}^{\infty} dx\,q(x, t)e^{-ikx}$$を施すことにより、（性質の良い関数に関しては）一切の損失なく情報を復元できます。
拡散方程式に代入してまとめると、$$\int_{-\infty}^{\infty} \frac{dk}{2\pi}\left(\frac{\partial}{\partial t} + \kappa k^2 \right)\tilde{q}(k, t)e^{ikx} = 0$$
となります。両辺を逆Fourier変換することで微分方程式$$\left( \frac{\partial}{\partial t} + \kappa k^2\right)\tilde{q}(k, t) = 0$$が得られるので、解いてあげます。$$\tilde{q}(k, t) = c(k)e^{-\kappa kt}.$$
ここで初期条件を使うと、$$\tilde{q}(k, t_0) =
\int_{-\infty}^{\infty} dx\,q_0(x)e^{-ikx}=
c(k)e^{-\kappa k^2t_0}$$となり、$c(k)$が得られました。
最後に $\tilde{q}(k, t)$ を逆Fourier変換してあげます。$$q(x, t) =
\int_{-\infty}^{\infty} \frac{dk}{2\pi}
\int_{-\infty}^{\infty} dx’\,q_0(x’)e^{-\kappa k^2(t-t_0)}e^{-ikx’}e^{ikx}.
$$ $k$ 積分を実行するため、平方完成 $$e^{-\kappa k^2(t-t_0)}e^{-ikx’}e^{ikx}
=
e^{-\kappa (t-t_0)\left(k – \frac{i(x-x’)}{2 \kappa (t – t_0)}
\right)^2 – \frac{(x-x’)^2}{4 \kappa (t – t_0)}
}$$を行ってあげると $k$ に関するGauss積分が出てきましたね。というわけで積分を実行し、最終的に時刻 $t > t_0$ での密度分布$$q(x, t) =
\frac{1}{\sqrt{4 \kappa (t – t_0)\pi}}
\int_{-\infty}^{\infty} dx’\,e^{- \frac{(x-x’)^2}{4 \kappa (t – t_0)}}q_0(x’)
$$が得られました。
$t > t_0$ としたのは、$t \geq t_0$では積分が収束しないからです。また $e^{-\kappa(t-t_0)\left(k – \frac{i(x-x’)}{2\kappa (t – t_0)}
\right)^2}$ の $k$ 積分に関して、カッコの中に虚数があるけどどうすんの？と思うかもしれません。これに関しては積分範囲を複素空間に拡張し、Cauchyの積分定理を用いることで積分経路を実軸 $k\in \mathbb{R}$ から $k = k’ + i(x-x’)/2\kappa (t – t_0), k’\in \mathbb{R}$ に変更することで対処できます。その際に余分な経路が出るかと思いますが、それらは消えるのでご安心ください。
Gauss分布と拡散方程式
拡散方程式が解けました。得られた解ですが、Gauss分布を用いて
$$q(x, t) =
\int_{-\infty}^{\infty} dx’\,\mathcal{N}(x; x’, 2\kappa (t – t_0))q_0(x’)
$$と書けますね。$q$ が非負（非正）である場合 $q/Q$ 確率分布として扱えます。そうしてみると、時刻 $t_0$ で確率変数 $x’$ が観測されたという条件のもとで、時刻 $t$ で確率変数 $x$ が観測される場合の条件付確率は、Gauss分布であることがわかりますね。というわけでGauss分布での確率分布間の遷移は拡散方程式によるものであると見做せるのでした。
ついでに初期時刻 $t_0$ から $t_1$ まで時間発展させて、さらに $t_1$ を初期条件として $t_2$ まで時間発展させることで、$$
q(x_2, t_2) =
\int_{-\infty}^{\infty} dx_1\,\mathcal{N}(x_2; x_1, 2\kappa (t_2 – t_1))
\int_{-\infty}^{\infty} dx_0\,\mathcal{N}(x_1; x_0, 2\kappa (t_1 – t_0))q(x_0, t_0)
$$を得ます。もちろんこれは初期時刻 $t_0$ から $t_2$ まで時間発展させた結果と等価なので比較し、ついでにGauss分布の標準偏差のところが任意の正の数に置き換えられるので、$$
\mathcal{N}(x_2; x_0, \sigma_1^2+\sigma_0^2)
=
\int_{-\infty}^{\infty} dx_1\,\mathcal{N}(x_2; x_1, \sigma_1^2)
\mathcal{N}(x_1; x_0, \sigma_0^2)
$$が成り立ちます。
この式を利用すると、以下の「diffusion processを繰り返したときの分布」で証明する式を簡単に示すことができます。
diffusion modelにおけるGauss分布
それでは本題に入ります。
どんな問題だっけ
まずは問題設定を思い出しておきましょう。diffusion modelでは、入力されたデータ $z_{t-1}$ に対して次のような形式でGauss型のノイズを加えていくのでした。$$
z_t = \sqrt{1-\beta_t}z_{t-1} + \sqrt{\beta_t}\epsilon, \quad \beta \in (0, 1).
$$ ただし、$\epsilon$ は確率分布 $\mathcal{N}(x; 0, 1)$ に従いサンプリングされた数です。ここでは $t$ は整数値をとる離散的なインデックスなので、先ほどの連続的な時間という意味から変わっていることに注意してください。
問題意識としては以下の2点です。
- 生成された $z_t$ はどんな分布に従うのか
- この操作を繰り返し行うと分布はどうなるのか
慣れている方には自明なことかもしれませんが、せっかくなので数式を使ってきっちり説明しておきたいと思います。
また、diffusion modelでは $z$ は多次元のベクトル量でした。しかしながらdiffusion processでは異なるベクトルの成分は独立な確率過程になるようにノイズを加えていたので、1次元の場合をやっておけば多次元への拡張は自明です。なので以下では $z$ はすべて1次元の確率変数です。
$z_t$ の分布
答えはGauss分布です。正直やっていることはもとのGauss分布をスケール変換して、並進移動しただけなので、Gauss分布のグラフを思い浮かべれば自明だと思います。$\epsilon$ が他の分布から生成された場合でも同様の考え方で対処できます。しかし（neural networkでは需要はないと思いますが）$z_{t}$ が $\epsilon$ により複雑に依存している場合にも対応できるように、数式を使ってきっちり確認しましょう。
タイムステップ $t-1$ で事象 $z_{t-1}$ であったという条件のもと、$t$ で事象 $z_{t}$ が生じる条件付確率は以下で定義することができます。$$
q(z_{t}|z_{t-1}) = E_{\epsilon}[\delta(z_{t} – f(z_{t-1}; \epsilon)].
$$今の場合は $f(z_{t-1}; \epsilon) = \sqrt{1-\beta_t}z_{t-1} + \sqrt{\beta_t}\epsilon$ です。ただし $\delta$ はDiracのデルタ関数で、期待値 $E_{\epsilon}[\cdot]$ は $\epsilon$ に関して取られていることにご注意ください。
$\epsilon$ が確率分布 $p(\epsilon)$ に従うとすれば$$
q(z_{t}|z_{t-1})
=
\int d\epsilon \,\delta(z_{t} – f(z_{t-1}; \epsilon))p(x)
$$となります。ここでデルタ関数に関しては以下の公式が成り立ちます。$$
\delta(g(x))
=
\sum_{i}\frac{1}{|g'(x_i)|}\delta(x – x_i).
$$ただし、$x_i$ は $g(x_i) = 0$ となる点であり、一般には複数存在します。以上より $q(z_{t}|z_{t-1})$ は$$
q(z_{t}|z_{t-1})
=
\sum_{i}\frac{p(\epsilon_i(z_{t-1}, z_{t}))}{|f'(z_{t-1}; \epsilon_i(z_{t-1}, z_{t}))|}
$$となることがわかりました。ただし、$\epsilon_i(z_{t-1}, z_{t})$ は $z_{t} = f(z_{t-1}; \epsilon)$ となる点です。
一般の場合には面倒なことになりましたが、$f(z_{t-1}; \epsilon) = \sqrt{1-\beta_t}z_{t-1} + \sqrt{\beta_t}\epsilon$ の場合にはとっても簡単です。なぜなら $\epsilon_i$ は $(z_{t} – \sqrt{1-\beta_t}z_{t-1})/\sqrt{\beta_t}$ しかない上に、$|f’| = \sqrt{\beta_t}$ だからです。というわけで代入してみると、$$
q(z_{t}|z_{t-1})
=
\frac{1}{\sqrt{\beta_t}}p(z_{t} – \sqrt{1-\beta_t}z_{t-1})/\sqrt{\beta_t})
$$となります。
以上より、$\epsilon$ がGauss分布に従う場合は
$$
q(z_{t}|z_{t-1})
=
\frac{1}{\sqrt{2\pi\beta_t}}e^{-\frac{(z_{t} – \sqrt{1-\beta_t}z_{t-1})^2}{2\beta_t}}
=
\mathcal{N}(z_{t}; \sqrt{1-\beta_t}z_{t-1}, \beta_t)
$$となることがわかりました。
diffusion processを繰り返したときの分布
この問題は、
$$
q(z_t|z_0)
=
\int dz_{t-1}\cdots\int dz_{1}\, q(z_{t}|z_{t-1})\cdots q(z_{1}|z_{0})
$$
がどのような分布であるかを答えるという問題であるといえます。まずは $q(z_2|z_0)$ の場合を調べます。
$\epsilon$ を使う方法
まずは 。$z_2$ は$$
z_2 = \sqrt{\alpha_2}(\sqrt{\alpha_1}z_{0} + \sqrt{\beta_1}\epsilon_1) + \sqrt{\beta_2}\epsilon_2
$$に従い生成されます。ただし、見やすいように $1-\beta_t = \alpha_t$ としました。すると、先程と同様に$$
q(z_2|z_0)
=
E_{\epsilon_1, \epsilon_2}[\delta(z_2 – \sqrt{\alpha_2}(\sqrt{\alpha_1}z_{0} + \sqrt{\beta_1}\epsilon_1) + \sqrt{\beta_2}\epsilon_2)]
$$となります。先に $\epsilon_2$ の期待値を実行すると、先ほどの計算の一部を置き換えるだけなので簡単です。$$
q(z_2|z_0)
=
\int d\epsilon_1
\frac{1}{\sqrt{2\pi\beta_2}}
e^{-\frac{(z_{2} – \sqrt{\alpha_2\alpha_1}z_{0} – \sqrt{\alpha_2\beta_1}\epsilon_1)^2}{2\beta_2}}
\frac{1}{\sqrt{2\pi}}e^{-\frac{\epsilon_1^2}{2}}
$$これを気合でまとめると以下のようになります。$$
q(z_2|z_0)
=
\int d\epsilon_1
\frac{1}{2\pi\sqrt{\beta_2}}
e^{
-\frac{1-\alpha_1\alpha_2}{2\beta_2}\left(
\epsilon_1 – \frac{\sqrt{\alpha_2\beta_1}}{1-\alpha_1\alpha_2}
(z_2 – \sqrt{\alpha_1\alpha_2}z_0)
\right)^2
}e^{-
\frac{(z_2 – \sqrt{\alpha_1\alpha_2}z_0)^2}{2(1-\alpha_1\alpha_2)}
}$$式を追う人のために少しだけ補足すると、 $1-\alpha_1\alpha_2 = \alpha_2\beta_1 + \beta_2$ という関係式を使います。以上より、Gauss積分を実行することで$$
q(z_2|z_0)
=\frac{1}{\sqrt{2\pi(1-\alpha_1\alpha_2)}}
e^{-
\frac{(z_2 – \sqrt{\alpha_1\alpha_2}z_0)^2}{2(1-\alpha_1\alpha_2)}
}
=
\mathcal{N}(z_1; \sqrt{\alpha_1\alpha_2}z_0, 1 – \alpha_1\alpha_2)
$$が得られました。
そのまま積分する方法
もちろんそのまま代入することでも同じ答えを得ることができます。しかし面倒なので、「Gauss分布と拡散方程式」で得られた公式を用います。$$q(z_2|z_0)
=
\int dz_1 \mathcal{N}(z_2; \sqrt{\alpha_2}z_1, \beta_2) \mathcal{N}(z_1; \sqrt{\alpha_1}z_0, \beta_1)
$$に対して、まずは変数変換 $\sqrt{\alpha_2}z_1 = z$ を行い、$$
q(z_2|z_0)
=
\int \frac{dz}{\sqrt{\alpha_2}} \mathcal{N}(z_2; z, \beta_2) \mathcal{N}(z/\sqrt{\alpha_2}; \sqrt{\alpha_1}z_0, \beta_1)
$$を得ます。次にGauss分布の形を見てあげると$$
\mathcal{N}(z/\sqrt{\alpha_2}; \sqrt{\alpha_1}z_0, \beta_1)
=
\sqrt{\alpha_1}\mathcal{N}(z; \sqrt{\alpha_1\alpha_2}z_0, \alpha_2\beta_1)
$$が成り立つことがわかると思います。最後に、公式を用いることで$$
q(z_2|z_0)
= \mathcal{N}(z_1; \sqrt{\alpha_1\alpha_2}z_0, 1 – \alpha_1\alpha_2)
$$が得られました。ただし、標準偏差の部分は $\alpha_2\beta_1 + \beta_2 = 1 – \alpha_1\alpha_2$ を用いてきれいにまとめています。
$t$ 回繰り返す場合
得られた $q(z_2|z_0)$ の結果を利用すると、$\alpha_1\alpha_2 = \tilde{\alpha}_2$ と置き換えることで、再帰的に $q(z_3|z_0)$ が構成可能になります。つまり $q(z_2|z_0)=\mathcal{N}(z_2; \sqrt{\tilde{\alpha_2}}z_0, 1 – \tilde{\alpha}_2)$ と書けるので、先ほどと全く同じ議論をすることにより、$$
q(z_3|z_0)
= \mathcal{N}(z_1; \sqrt{\tilde{\alpha_2}\alpha_3}z_0, 1 – \tilde{\alpha}_2\alpha_3)
$$が得られます。従って、$$
\tilde{\alpha_t}
=
\prod_{i=1}^{t}\alpha_t
$$を用いることで、$$
q(z_t|z_0)
=
\mathcal{N}(z_1; \sqrt{\tilde{\alpha_t}}z_0, 1 – \tilde{\alpha}_t)
$$であることが分かりました。
以上よりdiffusion processを何度繰り返そうと、中間部分に関して積分をとってしまえば、一回の diffusion process と等価であるという結論が得られました。$\epsilon$ を使って書いておくと$$
z_t = \sqrt{\tilde{\alpha}_t}z_{t-1} + \sqrt{1 – \tilde{\alpha}_t}\epsilon
$$です。
コメント