拡散モデルの数学 — スコアとSDEで見る生成

拡散モデルを「ノイズを足して引く」の一段下、スコア（対数確率の勾配）の言葉で捉え直す。なぜノイズ除去がスコア推定と同じことなのか、前向き・逆向きSDEと確率フローODEが何を言っているのか、そしてその式が推論設定のつまみにどう化けるかまで。

Generative Modeling by Estimating Gradients of the Data Distribution

標高は諦めて、傾きだけを覚える

夜の山中で、地図もGPSもないまま頂上を目指すとします。自分の標高が何メートルかを知る方法はありません。それでも登れます。足元の傾きさえ分かれば、上り坂の方へ一歩ずつ進めばいいからです。

生成モデルが直面しているのは、まさにこの状況です。「それらしい画像」の確率分布 $p(x)$ を直接学ぼうとすると、必ず正規化定数が邪魔をします。モデルをエネルギー関数 $E_\theta(x)$ で書くと、確率はこうなります。

$Z_\theta$ は「全部足したら1になるように」割るための数で、積分はあらゆる画像にわたって取ります。$256\times256$ のカラー画像なら約20万次元の空間全体の積分です。誰にも計算できません。

ところが、両辺の対数を取ってから $x$ で微分すると、$Z_\theta$ は $x$ を含まない定数なので消えてしまいます。

これがスコアです。名前は統計学のスコア関数（対数尤度の勾配）から来ていますが、パラメータではなくデータ $x$ の方で微分するのが拡散モデル流です。$\nabla_x$ は「$x$ の各成分で微分して並べたもの」、つまり入力と同じ形のベクトルです。画像なら、画素と同じ枚数・同じ解像度の矢印の場が出てくると思ってください。意味は「この画素をほんの少しどちらへ動かせば、それらしさが上がるか」。標高そのものではなく、傾きだけを覚える作戦です。

拡散モデル全体の流れ（ノイズを足す前向き過程と、引く逆過程）は拡散モデルを1から理解するで扱いました。この記事はその一段下、なぜあの手続きが正しいのかをスコアの言葉で見ていきます。

傾きだけで、どうやって標本を作るのか

スコアが手に入ったとして、そこからサンプルを取り出すのがランジュバン動力学です。式はこれだけです。

$\varepsilon$ は歩幅、$z_k$ は毎回引き直す標準正規乱数です。言い換えると「濃い方へ半歩進んで、サイコロの分だけよろける」を延々と繰り返すだけ。$\varepsilon$ を十分小さく、回数を十分多くすれば、$x_k$ の分布は $p_\theta$ に近づいていきます。

第2項の揺らぎが本質的です。これを外すと単なる勾配上昇になり、一番確率の高い1点に張り付いて動かなくなります。それは「最頻値を求めること」であって「分布から標本を引くこと」ではありません。ノイズがあるから斜面をうろつき、確率の高い場所には長く、低い場所には短く滞在する——その滞在時間の比が、そのまま確率になります。

具体例で手触りを掴みましょう。$p(x) = \mathcal{N}(0, \sigma^2)$ のスコアは、$\log p(x) = -x^2/(2\sigma^2) + \text{const}$ を微分して

原点に向かって引き戻すバネです。中心から遠いほど強く引かれ、$\sigma$ が小さい（分布が尖っている）ほどバネが硬い。つまりガウス分布に対するランジュバンは、放物線の谷を球が転がるのに、毎回ランダムな小突きを加えたものと完全に同じ運動です。歩幅を上げすぎるとどうなるかも、放物線の話とそっくり同じことが起きます。

「ノイズ除去」と「スコア推定」が同じことである理由

ここで大きな問題があります。スコアの正解ラベルは、どこにもありません。真の $\nabla_x \log p(x)$ を知っているなら、そもそも生成モデルは要らないのです。

素朴なスコアマッチング（Hyvärinen, 2005）は、部分積分を使って未知の $p$ を目的関数から消し去り、$\frac{1}{2}\|s_\theta(x)\|^2 + \operatorname{tr}(\nabla_x s_\theta(x))$ の期待値を最小化すればよい形に持ち込みます。理論は完璧ですが、右の項はヤコビ行列の対角和で、次元の数だけ逆伝播が要ります。20万次元の画像では話になりません。

突破口は、データをわざと汚すことでした。きれいなデータの分布は分かりませんが、そこにガウスノイズを足した条件付き分布なら、こちらが作ったのだから完全に分かっています。$z \sim \mathcal{N}(0,I)$ として $\tilde{x} = x + \sigma z$ と汚せば $p_\sigma(\tilde{x}\mid x) = \mathcal{N}(\tilde{x}; x, \sigma^2 I)$、そのスコアはさっきのバネの式そのものです。