拡散モデル

拡散モデルとは

拡散モデルは、データに徐々にノイズを加えていく過程（拡散過程）と、 ノイズから元のデータを復元する過程（逆拡散過程）を学習する生成モデルです。

数学的な定式化

PyTorch実装（DDPM）

import torch
import torch.nn as nn
import torch.nn.functional as F

class DiffusionModel:
    def __init__(self, n_steps=1000, beta_start=1e-4, beta_end=0.02):
        # ノイズスケジュール
        self.betas = torch.linspace(beta_start, beta_end, n_steps)
        self.alphas = 1 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)

    def add_noise(self, x0, t, noise=None):
        """拡散過程: x0にノイズを加えてxtを生成"""
        if noise is None:
            noise = torch.randn_like(x0)

        alpha_bar_t = self.alpha_bars[t].view(-1, 1, 1, 1)
        # xt = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * noise
        xt = torch.sqrt(alpha_bar_t) * x0 + torch.sqrt(1 - alpha_bar_t) * noise
        return xt, noise

    def sample(self, model, shape, device):
        """逆拡散過程: ノイズから画像を生成"""
        # 純粋なノイズから開始
        x = torch.randn(shape, device=device)

        for t in reversed(range(len(self.betas))):
            t_tensor = torch.full((shape[0],), t, device=device, dtype=torch.long)

            # ノイズを予測
            predicted_noise = model(x, t_tensor)

            # 1ステップ分ノイズを除去
            alpha_t = self.alphas[t]
            alpha_bar_t = self.alpha_bars[t]
            beta_t = self.betas[t]

            # 平均を計算
            mean = (1 / torch.sqrt(alpha_t)) * (
                x - (beta_t / torch.sqrt(1 - alpha_bar_t)) * predicted_noise
            )

            if t > 0:
                noise = torch.randn_like(x)
                x = mean + torch.sqrt(beta_t) * noise
            else:
                x = mean

        return x


# ノイズ予測ネットワーク（U-Net構造）
class NoisePredictor(nn.Module):
    def __init__(self, channels=1, time_embed_dim=256):
        super().__init__()

        # 時間埋め込み
        self.time_mlp = nn.Sequential(
            nn.Linear(1, time_embed_dim),
            nn.ReLU(),
            nn.Linear(time_embed_dim, time_embed_dim)
        )

        # 簡略化したU-Net（実際はもっと複雑）
        self.down1 = nn.Conv2d(channels, 64, 3, padding=1)
        self.down2 = nn.Conv2d(64, 128, 3, stride=2, padding=1)
        self.mid = nn.Conv2d(128, 128, 3, padding=1)
        self.up1 = nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1)
        self.out = nn.Conv2d(64, channels, 3, padding=1)

    def forward(self, x, t):
        # 時間情報を埋め込み
        t_emb = self.time_mlp(t.float().unsqueeze(-1) / 1000)

        # ダウンサンプリング
        h1 = F.relu(self.down1(x))
        h2 = F.relu(self.down2(h1))

        # ミドル（時間情報を追加）
        h = F.relu(self.mid(h2))

        # アップサンプリング
        h = F.relu(self.up1(h))
        return self.out(h + h1)

学習プロセス

# 学習ループ
model = NoisePredictor()
diffusion = DiffusionModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(epochs):
    for x0 in dataloader:
        # ランダムな時刻tを選択
        t = torch.randint(0, 1000, (x0.size(0),))

        # ノイズを加える
        xt, noise = diffusion.add_noise(x0, t)

        # ノイズを予測
        predicted_noise = model(xt, t)

        # 損失: 予測ノイズと実際のノイズの差
        loss = F.mse_loss(predicted_noise, noise)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

主要な拡散モデル

理解度チェック