注釈付き拡散モデル – DDPM実装の詳細なウォークスルー

TL;DR

Hugging Face は、元の Denoising Diffusion Probabilistic Model (DDPM) を実装した完全な注釈付き PyTorch ノートブックを公開しました。これは、純粋なガウスノイズから画像をデノイズすることを学習するニューラルネットワークで、低解像度データセットでの訓練と高品質画像のサンプリングを実演しています。このリソースは、数学、ネットワークアーキテクチャ、訓練ループ、推論手順を分かりやすく解説し、拡散モデルを実務家にとって利用しやすくします。

拡散モデルが何をするか（重要なポイント）

拡散モデルは、ランダムなガウスノイズを現実的なデータへと変換する逆デノイズプロセスを学習します。前方プロセスは事前に定められたスケジュールでノイズを付加し、U‑Net スタイルのニューラルネットワークが各タイムステップで付加されたノイズを予測することで、固定回数のステップ後に画像生成を可能にします。

前方拡散：分散スケジュールによるノイズ付加

結論: 前方プロセスは閉形式のガウス遷移であり、任意のノイズ付きタイムステップを直接サンプリングできるため、反復的なノイズ付加が不要になります。

前方拡散分布は次のようになります。

q(x_t | x_{t-1}) = N(x_t ; sqrt(1-β_t)·x_{t-1}, β_t·I)

ここで分散スケジュール (β_1,…,β_T) は単調（例：線形、余弦、二次、シグモイド）です。この遷移を繰り返し適用することで、周辺分布

q(x_t | x_0) = N(x_t ; sqrt(\bar α_t)·x_0, (1-\bar α_t)·I)

は、事前に計算された (\bar α_t = \prod_{s=1}^t (1-β_s)) を用いて直接サンプリングできます。この「便利な性質」により、訓練アルゴリズムはランダムなタイムステップ (t) を選び、実画像 (x_0) を一ステップで汚染できます。

逆拡散：ノイズ予測による平均の学習

結論: モデルはタイムステップ (t) で付加された正確なガウスノイズ (ε) を予測するように訓練され、予測されたノイズは逆ガウス分布の平均を計算するために使用されます。

逆条件分布はガウスと仮定されます：

p_θ(x_{t-1} | x_t) = N(x_{t-1} ; μ_θ(x_t, t), σ_t^2·I)

DDPM は既知の前方分散 (σ_t^2) を固定し、(μ_θ) のみを学習します。平均をノイズ予測子 (ε_θ) で再パラメータ化すると：

μ_θ(x_t, t) = (1/√α_t)·(x_t - (β_t/√(1-\bar α_t))·ε_θ(x_t, t))

訓練損失は、真のノイズ (ε) とネットワーク出力 (ε_θ) との単純な平均二乗誤差に縮小されます：

L_t = || ε - ε_θ(x_t, t) ||^2

各バッチでランダムに (t) をサンプリングすることで、変分下界の不偏推定量が得られます。

ネットワークアーキテクチャ：条件付き U‑Net

結論: DDPM は、時間条件付き U‑Net を使用し、サイン波位置埋め込み、残差ブロック、アテンション、グループ正規化を組み合わせて、すべての空間解像度でノイズを予測します。

主要コンポーネント：

  • サイン波位置埋め込み はタイムステップ (t) をエンコードし、小さな MLP を通して各 ResNet ブロックに注入されます。
  • 重み標準化畳み込み はグループ正規化と組み合わせることで訓練の安定性を向上させます。
  • ResNet ブロック（時間埋め込みからのオプションのスケールシフトを伴う 2 つの convolution‑norm‑SiLU 層）は、主な特徴変換を提供します。
  • アテンションモジュール（フルマルチヘッドまたは線形アテンションのいずれか）は、長距離依存関係を捉えます。
  • グループ正規化 はアテンションの前に適用されます（PreNorm）。
  • ダウンサンプリング/アップサンプリングパス は空間解像度を半分または倍にしながらチャネル深さを保持し、古典的な U‑Net 設計を鏡像します。
  • 最終ヘッド は結合されたボトルネック特徴を画像形状に戻します。

完全な PyTorch クラス Unet はこれらの部品を組み立て、ノイズ画像テンソルとタイムステップテンソルを受け取り、同形状のノイズテンソルを返します。

訓練ループ：確率的タイムステップサンプリングと Huber 損失

結論: 訓練はバッチごとにランダムなタイムステップをサンプリングし、閉形式の前方プロセスで入力を汚染し、真のノイズと予測ノイズ間の Huber 損失を最小化することで進行します。

疑似コード（簡略版）：

for epoch in range(num_epochs):
    for batch in dataloader:
        t = torch.randint(0, T, (batch_size,)).to(device)          # uniform timestep
        loss = p_losses(model, batch, t, loss_type='huber')        # MSE/Huber on noise
        loss.backward()
        optimizer.step()

ヘルパー関数 p_losses は q_sample を呼び出して (x_t) を取得し、ネットワーク出力 model(x_t, t) に対する損失を計算します。定期的なサンプリング（逆プロセスを使用）で訓練の進捗を可視化します。

サンプリング（推論）：拡散チェーンの逆転

結論: 生成は純粋なガウスノイズから開始し、学習されたデノイズステップを反復的に適用します。各ステップは予測されたノイズを用いて事後平均を計算し、既知の分散スケジュールに従って校正されたガウスノイズを加えます。

アルゴリズム（DDPM 論文のアルゴリズム 2）：

x_T = torch.randn(shape)                     # start from noise
for t in reversed(range(T)):
    pred_noise = model(x_t, t)
    mean = (1/√α_t) * (x_t - β_t/√(1-\bar α_t) * pred_noise)
    if t > 0:
        x_{t-1} = mean + √posterior_variance_t * torch.randn_like(x_t)
    else:
        x_0 = mean

提供されたノートブックはこのループを p_sample_loop で実装し、デノイズの軌跡を GIF として可視化します。

Fashion‑MNIST のエンドツーエンド例

結論: このチュートリアルは 28×28 の Fashion‑MNIST データセットで DDPM を訓練し、数千ステップの訓練後に認識可能な衣類を生成します。

  • データロードは 🤗 datasets ライブラリを使用し、オンザフライ変換（ランダム水平フリップ、([-1,1]) へのスケーリング）を行います。
  • モデルハイパーパラメータ: グレースケール画像の場合、dim=image_size、dim_mults=(1,2,4)、channels=1。
  • 最適化手法: 学習率 1e‑3 の Adam。
  • 訓練は 6 エポック実行され、損失はすぐに 0.05 未満に低下します。
  • 訓練後のサンプリングは明瞭な T‑シャツ形状を生成し、実装が低解像度データでも機能することを確認します。

拡散文献における位置付け

結論: この注釈付き実装はオリジナルの DDPM（Ho et al., 2020）を再現し、分散学習（Nichol et al., 2021）、カスケード拡散（Ho et al., 2021）、classifier‑free guidance、そして大規模テキスト‑画像モデル（DALL‑E 2、Imagen）など、後続の多くの進展の基盤となります。

ブログ記事で挙げられた主要なフォローアップ研究：

  • Improved DDPM – 平均と分散の両方を学習し、サンプル品質を向上させます。
  • Cascaded Diffusion Models – 複数の拡散モデルを積み重ねて高解像度合成を実現します。
  • Diffusion Models Beat GANs – アーキテクチャの調整と classifier guidance により、優れた FID スコアを示します。
  • Classifier‑Free Guidance – 条件付き生成時に外部 classifier を不要にします。
  • DALL‑E 2 & ImageGen – 拡散と CLIP 埋め込みや大規模言語モデルを組み合わせ、テキスト条件付き画像合成を行います。

主な欠点は多くのデノイズステップが必要なこと（通常は 1000 ステップ）ですが、最近の研究（例：高次ソルバー）によりこれが 10 ステップ程度に削減されています。

開発者への実践的なポイント

  • ノートブックはすぐに実行できるリファレンス実装を提供します。データセットを置き換えたり、T を増やしたり、U‑Net をより大きなバックボーンに置き換えることでスケールアップできます。
  • 分散スケジュール（linear_beta_schedule、cosine_beta_schedule など）を調整して、サンプルの忠実度と速度のトレードオフが可能です。
  • ロバスト性のために Huber 損失を使用します。コードは L1/L2 もサポートしています。
  • 定期的なサンプリング（save_and_sample_every）は、モード崩壊や訓練不安定性の監視に不可欠です。
  • より高速な推論のために、DDIM や確率的サンプラーのバリエーションを実装するか、最近の「少ステップ」ソルバーを採用してください。

Sources