みんなのためのAI 第6回 — 111万パラメータの拡散モデルで単語から数字を描く

• はじめに — 第5回が残した宿題
• 核心のアイデア — 壊すのは簡単、戻すのは難しい
  • Forward — ノイズを混ぜる過程
  • 学習目標 — 画像ではなくノイズを当てる
• 条件を付ける — 単語で絵を指示する
• Reverse — 400回逆にたどる
• 学習ログ
• 結果
• まとめ
• 🧠 理解度チェッククイズ
• 参考資料

はじめに — 第5回が残した宿題

第5回のカラー化は形をよく保ちましたが色が褪せました。原因はモデルではなく損失でした。白黒の自動車は赤かもしれず青かもしれないのに、L1 は「絶対誤差を最小化せよ」と指示するので、 複数の正解の中間値である灰色を打つのが最も安全な戦略 になります。

今回の問題も正解が複数あります。"three" という単語に対応する手書きの3は数千通りです。もし第5回のように回帰損失で「3の画像を直接予測せよ」と指示すれば、すべての3の平均であるぼやけた染みが出てくるでしょう。

拡散モデルはこの問題を正面から回避します。 一度に正解を作らず、ノイズから少しずつ取り除いていく過程に変えます。

核心のアイデア — 壊すのは簡単、戻すのは難しい

きれいな画像にノイズを少しずつ混ぜて完全な雑音にするのは簡単です。数式1行で済みます。難しいのはその逆方向です。

拡散モデルの発想はこうです。 難しい問題を簡単な問題400個に分割する。 雑音から画像を一度に作るのは難しいですが、「ほんの少しだけきれいにする」なら学べます。それを400回繰り返します。

Forward — ノイズを混ぜる過程

T = 400
beta = torch.linspace(1e-4, 0.02, T)   # 各段階で混ぜるノイズの量
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)         # 累積積

abar[t] は「時刻 t まで元画像がどれだけ残っているか」を表します。t=0 ではほぼ1 (元のまま)、t=399 ではほぼ0 (完全な雑音) です。

ここで拡散モデルの第一の仕掛けが登場します。 400段階を順にたどらず、任意の t へ一度にジャンプできます。

t = torch.randint(0, T, (256,), device=dev)
eps = torch.randn_like(x0)                          # 標準正規分布のノイズ
a = abar[t][:, None, None, None]
xt = a.sqrt() * x0 + (1 - a).sqrt() * eps           # 1行で t 段階分のノイズ適用

この1行のおかげで学習が効率的になります。バッチごとに異なる t を引いて、すべての時刻を並列に学習させられます。

学習目標 — 画像ではなくノイズを当てる

loss = F.mse_loss(model(xt, t.float(), w), eps)

モデルが予測するのは元画像 x0 ではなく 混ざり込んだノイズ eps です。

数学的には両者は等価です。xt と t が分かれば eps から x0 を計算でき、逆も成り立ちます。しかし実際にはノイズ予測のほうがはるかにうまく動きます。

理由は目標の性質にあります。x0 を予測させると、t が大きいとき (ほぼ雑音のとき) モデルは事実上情報なしで画像をでっち上げねばならず、ここで第5回の平均回帰が再発します。一方 eps は どの t でも常に標準正規分布 です。目標のスケールと分布が一定なので学習が安定します。

条件を付ける — 単語で絵を指示する

時刻 t と単語を同じベクトル空間で足し、U-Net の中央に注入します。

class CondUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.temb = nn.Sequential(nn.Linear(64, 128), nn.GELU(), nn.Linear(128, 128))
        self.wemb = nn.Embedding(10, 128)          # zero~nine
        self.e1 = block(1, 48); self.e2 = block(48, 96); self.e3 = block(96, 192)
        self.cond3 = nn.Linear(128, 192)
        self.d2 = block(192 + 96, 96); self.d1 = block(96 + 48, 48)
        self.out = nn.Conv2d(48, 1, 1)

    def forward(self, x, t, w):
        half = torch.exp(-np.log(10000) * torch.arange(32, device=x.device) / 32)
        te = torch.cat([torch.sin(t[:, None] * half), torch.cos(t[:, None] * half)], 1)
        c = self.temb(te) + self.wemb(w)           # 時刻 + 単語
        s1 = self.e1(x)
        s2 = self.e2(F.max_pool2d(s1, 2))
        h = self.e3(F.max_pool2d(s2, 2)) + self.cond3(c)[:, :, None, None]
        ...

時刻埋め込みにサイン・コサインを使うのは、トランスフォーマーの位置エンコーディングと同じ発想です。整数 t をそのまま入れると神経網が扱いにくいスケールになりますが、異なる周波数のサイン波に広げれば隣接する時刻同士が似たベクトルを持つようになります。

注入位置が 最も深い層 (e3 の出力) であるのも意図的です。解像度が8×8で最も小さくチャンネルあたりの情報密度が高いので、ここで条件を掛ければ以後のデコーダ全体が影響を受けます。

self.cond3(c)[:, :, None, None] の None 2つは、(バッチ, チャンネル) のベクトルを (バッチ, チャンネル, 1, 1) にして空間次元全体へブロードキャストするためのものです。つまり 条件がすべての位置に等しく加算されます。

Reverse — 400回逆にたどる

@torch.no_grad()
def sample(words):
    x = torch.randn(len(words), 1, 28, 28, device=dev)      # 純粋な雑音から出発
    w = torch.tensor([WORDS.index(s) for s in words], device=dev)
    for ti in reversed(range(T)):                            # 399 -> 0
        t = torch.full((len(words),), ti, device=dev, dtype=torch.float32)
        eps = model(x, t, w)                                 # ノイズを予測
        a, ab = alpha[ti], abar[ti]
        x = (x - (1 - a) / (1 - ab).sqrt() * eps) / a.sqrt() # 1段階分を除去
        if ti > 0:
            x = x + beta[ti].sqrt() * torch.randn_like(x)    # 新しいノイズを再注入
    return ((x.clamp(-1, 1) + 1) / 2).squeeze(1).cpu().numpy()

最後の行のノイズ再注入は奇妙に見えるかもしれません。せっかく取り除いたノイズをなぜまた入れるのでしょうか。

それが多様性の源だからです。 毎段階で新しい乱数を入れるので、同じ単語を入れても毎回違う3が出てきます。第5回で回帰が一つの平均値に囚われたのとは正反対です。拡散モデルは 複数の正解のうち一つを確率的に選び取る 構造です。最後の段階 (ti == 0) でのみ再注入を省き、きれいな結果を得ます。

学習ログ

[    9.2s] params {"params": 1111777, "T": 400, "n": 60000}
[   10.1s] step      0 mse=1.1502
[   12.2s] step    100 mse=0.0807
[  959.7s] step  44200 mse=0.0324
[  960.6s] SUMMARY {"steps": 44215, "final_mse": 0.0320}

100ステップ、2秒で MSE が1.15から0.081に落ちました。予測対象が標準正規分布なので初期損失が1付近から始まるのも自然です。以後なだらかに0.032まで下がりました。

結果

各行は "zero" から "nine" までの10単語で、2つの行は同じ単語に対する異なるサンプルです。

1行目は0から9まですべて正確です。 画の太さと傾きも手書きらしくなっています。

2行目で2つ間違えました。 "five" の位置に9に近い形が、"nine" の位置に7に近い形が出ました。10個中8個正解、80%です。

2つの行を比べると拡散モデルの性質が現れます。 同じ単語なのに字体が違います。 1行目の2と2行目の2は曲線の曲がり具合が異なり、1行目の7と2行目の7は横画の有無が違います。平均でつぶれることなく毎回一つの具体的な手書きを作り出しているということです。第5回の褪せた色と対比される点です。

間違えた5と9は形の上で互いに似ている数字です。111万パラメータ、16分の学習ではこの程度の混同が残ります。

まとめ

要点を3行にまとめるとこうなります。

拡散モデルは 難しい生成問題を簡単なノイズ除去400個に分割します。 forward は数式1行で、学習は任意の t へジャンプして並列化されます。

モデルは画像ではなく ノイズを予測します。 目標がどの t でも標準正規分布で一定なので学習が安定し、x0 を直接予測するときに起きる平均回帰を避けられます。

reverse 過程の ノイズ再注入が多様性を作ります。 同じ条件でも毎回違う結果が出るのが回帰モデルとの決定的な違いです。

ただし今回の実装は基本形です。400段階すべてをたどる必要があるので遅く、条件をどれだけ強く掛けるか調節する手段もありません。シリーズ第11回でこの2つを解決する Classifier-Free Guidance と DDIM を扱いながら、2つの技法が互いに独立ではないことも実測で見ることになります。

🧠 理解度チェッククイズ

参考資料

• MNIST データセット
• Denoising Diffusion Probabilistic Models (DDPM)
• Diffusion Models Beat GANs on Image Synthesis
• Attention Is All You Need — サイン・コサイン位置エンコーディング