【実装】拡散モデルを自作する — MNISTからはじめる

28×28の手書き数字で拡散モデルを最小構成から組み上げます。ノイズスケジュールが満たすべき2条件、UNetにステップ番号を注入する方法、サンプリングの最後にノイズを足し直す理由まで、自分で書いて初めて詰まる場所を順に潰していきます。

Denoising Diffusion Probabilistic Models

砂に埋めた彫刻を、1掃きずつ掘り出す

拡散モデルの解説を読むと、たいてい「ノイズを足して、引く」で腑に落ちた気になります。ところが自分で書こうとすると手が止まる。$\sqrt{\bar{\alpha}_t}$ の配列はいつ作るのか、ステップ番号 $t$ はどうやってネットワークに教えるのか、学習が終わったあと最初の1枚はどこから湧いてくるのか。分かった気になっている部分と本当に分かっている部分の境目は、コードを書くと一瞬で露出します。

この記事では MNIST（28×28 のグレースケール手書き数字）を題材に、拡散モデルを最小構成から組みます。MNIST を選ぶのは弱気だからではありません。28×28×1 は、手元のGPUでも数十分で「数字に見えるもの」が出る最小の題材で、しかも失敗が目で読めます。灰色の靄しか出なければ正規化かスケジュール、同じ字ばかり出れば多様性、全体がざらつくならサンプリングの最終ステップ——と原因が絞れる。512×512のカラーで同じ失敗をすると、切り分ける前に一晩溶けます。

比喩を1つだけ置きます。砂場に彫刻を埋め、上から砂をひと掬いずつ、何百回もかけて完全に見えなくする。拡散モデルの学習とは、その埋める作業を録画しておいて、「いまの砂の状態を見て、直前にかけられたひと掬いを言い当てる」練習を延々と繰り返すことです。言い当てられるようになれば、砂だけの状態から始めて、ひと掬いずつ取り除いていける。取り除き終わったとき、そこには埋めた覚えのない彫刻が現れます。

壊し方を自分で決める、という発想の転換

生成モデルが難しいのは、データの分布そのものが複雑すぎて直接は扱えないからです。GANは本物らしさの判定を別のネットワークに任せ、VAEは潜在変数を挟んで近似しました（VAEを1から）。拡散モデルの答えはどちらとも違います。難しい問題を一発で解くのをやめ、簡単な問題を何百個も並べたものに置き換える。

鍵は、壊す側を自分で設計してしまうことです。元画像にほんの少しノイズを足す、を何百回繰り返せば最後は完全なノイズになる。この過程は人間が決めたルールなので、どの時点の状態も正確に計算できます。そして1回あたりの変化が十分小さければ、その逆向きの1歩もまた同じ形の（ガウス分布の）操作で近似できることが知られています。だから学習対象は「1ステップ分の小さなズレ」に固定される。ネットワークは毎回、同じ形の易しい回帰問題だけを解けばよくなります。

前向き過程: T段の階段を、1回のジャンプで降りる

壊す側の定義から書きます。$\beta_t$ は「ステップ $t$ でどれだけ壊すか」を決める小さな正の数です。

つまり「1つ前の画像をほんの少し縮めて、そこに分散 $\beta_t$ のノイズを足す」だけです。足す前に縮めるのは、何百回繰り返しても値の大きさが暴れないようにするための調整で、これがあるおかげで最後の状態がちょうど標準ガウス分布に落ち着きます。

このまま実装すると、$x_t$ を得るのに $t$ 回ループが要ります。学習ではサンプルごとにランダムな $t$ を引くので、これは致命的に遅い。ところが $\alpha_t = 1-\beta_t$、$\bar{\alpha}_t = \prod_{s \le t}\alpha_s$ と置くと、途中を全部飛ばした形が手に入ります。

要するに「元画像を $\sqrt{\bar{\alpha}_t}$ 倍に薄め、残りをノイズで埋める」。$\bar{\alpha}_t$ は1から0へ落ちていく数で、その落ち具合がそのまま画像とノイズの混合比になります。実装上はこの1本がすべてです。ループは消え、任意の $t$ の状態が1行で作れます。

ノイズスケジュール — 実質ここだけが設計

T = 1000
betas = torch.linspace(1e-4, 0.02, T)          # DDPM原論文の線形スケジュール
alphas = 1.0 - betas
alphas_bar = torch.cumprod(alphas, dim=0)      # ᾱ_t。事前に1本作って使い回す

def q_sample(x0, t, eps):                      # t はバッチ内でバラバラでよい
    ab = alphas_bar[t].view(-1, 1, 1, 1)
    return ab.sqrt() * x0 + (1 - ab).sqrt() * eps

$\beta$ の決め方が、拡散モデルで実質的に唯一の設計です。DDPM原論文は $T=1000$ とし、$\beta$ を $10^{-4}$ から $0.02$ まで線形に増やしました。見るべき条件は2つあります。

1つ目は、終点で $\bar{\alpha}_T$ が十分0に近いこと。ここが破れると事故になります。生成は純粋なノイズから始まるのに、学習中のモデルは元画像がうっすら残った入力しか見ていない。初手がいきなり分布外になり、出てくるのは灰色の靄です。MNISTは簡単だからと $T$ を200へ減らし、$\beta$ の範囲を据え置くと、まさにこれが起きます。$T$ を減らすなら $\beta$ の終点も上げる。

2つ目は、壊れ方が途中に偏らないこと。線形スケジュールは低解像度の画像を前半で壊しすぎると後続研究で指摘され、序盤をゆるやかにするコサイン型が提案されました。28×28 はまさにその低解像度側です。