完全なノイズから、三日月形が本当に浮かび上がってくるのか。拡散モデル(DDPM)を自分で実装して確かめてみた
約10分で読めます

完全なノイズから、三日月形が本当に浮かび上がってくるのか。拡散モデル(DDPM)を自分で実装して確かめてみた

画像生成AIの説明でよく見る「完全なノイズから、少しずつノイズを取り除いていくと画像が生まれる」という話を、これまでは仕組みとしてなんとなく理解した気になっていた。実際に手を動かして中身を確かめたことはなかったので、今回は拡散モデル(Denoising Diffusion Probabilistic Models、DDPM)を、2次元の点群データに対してnumpyだけで一から実装してみた。画像だと計算が重いので、代わりに「two moons」と呼ばれる三日月が2つ組み合わさった形のおもちゃデータを、ノイズから本当に復元できるか試している。

前向き過程: データに少しずつノイズを混ぜる

拡散モデルの学習は、まず「データを徐々に壊してノイズにする」前向き過程から始まる。時刻 t=0t=0 が元のデータ、t=Tt=T が完全なノイズになるようにスケジュールを組む。

xt=αˉt x0+1−αˉt ϵ,ϵ∼N(0,I)x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \epsilon, \qquad \epsilon \sim \mathcal{N}(0, I) 
betas = np.linspace(1e-4, 0.02, T)          # ノイズを混ぜる強さのスケジュール(T=100ステップ)
alphas = 1 - betas
alpha_bars = np.cumprod(alphas)

def q_sample(x0, t_idx, noise):
    ab = alpha_bars[t_idx][:, None]
    return np.sqrt(ab) * x0 + np.sqrt(1 - ab) * noise

実際にtwo moonsのデータに、ttを0から99まで動かしながらノイズを混ぜていくとこうなる。

two moonsの形をしたデータに、tを0から99まで増やしながらノイズを混ぜていく過程を示す6枚の散布図。t=0では2つのはっきりした三日月形だが、t=25あたりから輪郭がぼやけ始め、t=50を超えると形はほぼ判別できなくなり、t=99では完全にランダムな点の塊になっている

t=25t=25あたりまでは形がまだ見えているが、t=50t=50を超えるとほぼ判別できなくなり、t=99t=99では完全にランダムな点の塊になる。ここまでは単にノイズを足しているだけなので、当然の結果だ。

ノイズ予測ネットワークを学習する

拡散モデルの本体は、この前向き過程を逆向きに辿るために、「時刻ttのノイズ入り点xtx_tを見て、そこに混ざっているノイズϵ\epsilonを推定する」ネットワーク ϵθ(xt,t)\epsilon_\theta(x_t, t) を学習することにある。今回は隠れ層1つ・256ユニットの小さなMLPを、numpyだけで(順伝播・逆伝播とも手書きで)実装した。時刻ttの情報は、前回の位置エンコーディングの記事で使ったのと同じsin/cosのベクトルに変換して、座標と連結して入力している。

def time_embedding(t_idx, dim=8):
    angle_rates = 1.0 / np.power(1000, (2 * (np.arange(dim) // 2)) / dim)
    angles = t_idx[:, None] * angle_rates[None, :]
    emb = np.zeros((len(t_idx), dim))
    emb[:, 0::2] = np.sin(angles[:, 0::2])
    emb[:, 1::2] = np.cos(angles[:, 1::2])
    return emb

# 学習ループ: ランダムなt・ランダムなノイズでx_tを作り、混ぜたノイズそのものを予測させる
for it in range(1, n_iters + 1):
    x0 = X0[rng.integers(0, len(X0), size=batch_size)]
    t_idx = rng.integers(0, T, size=batch_size)
    noise = rng.normal(0, 1, size=x0.shape)
    xt = q_sample(x0, t_idx, noise)
    pred_noise, cache = forward(xt, time_embedding(t_idx))
    loss = np.mean((pred_noise - noise) ** 2)
    # Adam(この記事シリーズ3本連続で同じ実装を使っている)で更新
    ...

20000イテレーション学習させたところ、損失は0.93から0.68まで下がった(ノイズは標準正規分布なので、何も学習していない状態の損失はちょうど1.0になる)。完璧な復元ではないが、多少ノイズの向きを言い当てられるようにはなった、というくらいの学習具合だ。

逆向き過程: ノイズから形を復元する

学習したネットワークを使って、完全なランダムノイズから出発し、DDPMの逆向き更新式を100ステップ繰り返して点群を少しずつ「denoise」していく。

for t_idx in reversed(range(T)):
    eps_theta = predict_noise(x, t_idx)
    mean = (1 / np.sqrt(alphas[t_idx])) * (x - (betas[t_idx] / np.sqrt(1 - alpha_bars[t_idx])) * eps_theta)
    x = mean + np.sqrt(betas[t_idx]) * rng.normal(0, 1, size=x.shape) if t_idx > 0 else mean

この過程をアニメーションにした。

完全なランダムノイズの点群から出発し、拡散モデルの逆向き過程を100ステップかけて実行していくアニメーション。最初は円形に広がったノイズの塊だが、ステップが進むにつれて徐々に2つの弧を描く形へと収束していき、最終的にtwo moons(2つの三日月)の形がはっきりと浮かび上がる

最初は完全にランダムな円形の点の塊だったのが、ステップを重ねるごとに少しずつ形が現れてくる。最終的な生成結果を、本物のデータと並べてみるとこうなる。

本物のtwo moonsデータ(左、青)と、拡散モデルが生成したサンプル(右、赤)を並べた散布図。どちらも上下2つの弧を描く三日月形をしているが、生成サンプルの方はやや輪郭がぼやけてノイズが多い

本物ほどくっきりしてはいないものの、上の弧と下の弧、2つの三日月という構造そのものはちゃんと再現できている。 ネットワークが「時刻ttでどれだけノイズが混ざっているはずか」を学習しただけなのに、そこから実際に構造のある形を組み立てられるのは、実装して自分の目で見るまでは半信半疑だった。

個々の点がどう動いたかを3Dで追う

最後に、1500点のうち180点をサンプルして、各点がt=100t=100(ノイズ)からt=0t=0(生成結果)までの100ステップでどう動いたかを、時刻を高さ方向にした3Dの軌跡として可視化した。色は最終的にどちらの三日月に着地したかを表している。

回転させて上から見ると、上部(t=100付近)では軌跡の出発点がほぼ円形にランダムに散らばっているのに対し、下部(t=0付近)では青と赤の2つの弧にはっきり分かれて収束しているのが分かる。横から見ると、序盤は各点がほとんど動かず、ttが0に近づくにつれて動きが速く・大きくなっていく様子も見える。これは前向き過程のノイズスケジュール(betaが線形に増える)を逆から辿っているためで、終盤ほど「まだ削るべきノイズの量」に対して踏み込む一歩が大きくなる。

手を動かして意外だったこと

一番意外だったのは、ネットワークが学習しているのは「今この瞬間、どんなノイズが混ざっているか」という、かなり局所的で近視眼的な問いへの答えだけだという点だ。ネットワークは「two moonsとは何か」という全体的な形の知識を直接教わっているわけではない。それにもかかわらず、その場その場のノイズ除去を100回繰り返すだけで、大域的に一貫した形が組み上がっていく。

これは、大きな目標を一気に達成しようとするのではなく、「今の状態から見て、少しだけ良い方向」を繰り返し選び続けることの強さを表している気がする。拡散モデルは、最初から完成形を思い描いて点を配置しているわけではなく、ただ「このノイズはここが少しおかしい」という一歩ずつの修正を100回積み重ねているだけだ。それでも、各ステップでの修正の向きさえ正しければ、最終的にはちゃんと意味のある形にたどり着く。3D軌跡を見ても、終盤に近づくほど動きが大胆になっていくのも面白い。最初は間違った方向に大きく踏み出すリスクを避けて小さく探り、状況がはっきりしてきた終盤ほど思い切って動く、というのは、物事を段階的に詰めていくときの感覚にも近い。

まとめ

  • DDPMの前向き過程を実装し、two moonsデータにt=0t=0からt=99t=99までノイズを混ぜていくと、t=50t=50を超えたあたりで形がほぼ判別不能になることを確認した
  • 時刻ttを位置エンコーディングと同じsin/cos埋め込みで表現し、隠れ層1つのMLPをnumpyだけで実装・学習させ、20000イテレーションで損失を1.0(無学習相当)から0.68まで下げた
  • 学習したノイズ予測ネットワークを使い、完全なランダムノイズから100ステップの逆向き更新を行うと、本物ほど鮮明ではないものの、two moonsの2つの弧という構造を実際に再現できた
  • 180点の3D軌跡を可視化すると、序盤(ノイズが多い段階)は各点の動きが小さく、終盤(ノイズが少ない段階)に近づくほど動きが大きく確信を持ったものになっていく様子が見えた
  • ネットワークは「今のノイズ量に対する局所的な修正」しか学習していないのに、それを100回積み重ねるだけで大域的に一貫した形が組み上がるという、拡散モデルの本質的な面白さを実際に確認できた

完全版(有料教材)を公開しました

この実験の「裏側」を全部見せる完全版をnoteで公開した。数式の導出(なぜ任意の時刻へ一気に飛べるのか、逆向きの更新式はどこから来るのか)を1本ずつ、逆伝播まで手書きの全コードを行単位の解説付きで、そして本記事には書いていない**「壊す実験」4本**——ノイズスケジュールを変える・ステップ数を10まで減らす・時刻埋め込みを外す——の実測結果まで。実は本記事の実装には理論的な欠陥が1つあった(のに、なぜか動いていた)という顛末も、正直に検証している。

▼ ノイズから形が生まれる瞬間を、全部自分の手で作る — numpyだけで学ぶ拡散モデル(DDPM)完全版(note)

関連記事