Diffusion 扩散模型入门指南：从 DDPM 到 Latent Diffusion 与 DiT

import torch

def extract(values, t, x_shape):
    """从长度 T 的表中取出每个样本对应的 t，并扩展到可广播形状。"""
    return values.gather(0, t).reshape(t.shape[0], *((1,) * (len(x_shape) - 1)))

def q_sample(x0, t, alpha_bar, noise=None):
    noise = torch.randn_like(x0) if noise is None else noise
    signal = extract(alpha_bar.sqrt(), t, x0.shape)
    noise_scale = extract((1.0 - alpha_bar).sqrt(), t, x0.shape)
    return signal * x0 + noise_scale * noise, noise

T = 1000
beta = torch.linspace(1e-4, 2e-2, T)
alpha = 1.0 - beta
alpha_bar = torch.cumprod(alpha, dim=0)

x0 = torch.randn(8, 3, 32, 32)       # [B,C,H,W]
t = torch.randint(0, T, (x0.size(0),))
xt, epsilon = q_sample(x0, t, alpha_bar)
assert xt.shape == x0.shape == epsilon.shape

from pathlib import Path

import matplotlib.pyplot as plt
import torch


steps = 1000

# Linear beta schedule
beta_linear = torch.linspace(1e-4, 2e-2, steps)
alpha_bar_linear = torch.cumprod(1.0 - beta_linear, dim=0)

# Cosine cumulative schedule，并保证起始累计信号为 1
offset = 0.008
grid = torch.linspace(0, 1, steps + 1)
alpha_bar_cosine = torch.cos(
    ((grid + offset) / (1 + offset)) * torch.pi / 2
).square()
alpha_bar_cosine = alpha_bar_cosine / alpha_bar_cosine[0]
alpha_bar_cosine = alpha_bar_cosine[1:].clamp_min(1e-8)

# SNR(t) = alpha_bar(t) / (1 - alpha_bar(t))
snr_linear = alpha_bar_linear / (1.0 - alpha_bar_linear)
snr_cosine = alpha_bar_cosine / (1.0 - alpha_bar_cosine)
t = torch.arange(steps)

plt.style.use("seaborn-v0_8-whitegrid")
fig, axes = plt.subplots(1, 2, figsize=(12, 4.2))

axes[0].plot(t, alpha_bar_linear, label="Linear beta", linewidth=2.2)
axes[0].plot(t, alpha_bar_cosine, label="Cosine", linewidth=2.2)
axes[0].set(
    title="Cumulative signal retention",
    xlabel="Timestep t",
    ylabel="alpha_bar(t)",
)
axes[0].set_ylim(-0.02, 1.02)
axes[0].legend()

axes[1].semilogy(
    t, snr_linear.clamp_min(1e-8), label="Linear beta", linewidth=2.2
)
axes[1].semilogy(
    t, snr_cosine.clamp_min(1e-8), label="Cosine", linewidth=2.2
)
axes[1].set(
    title="Signal-to-noise ratio",
    xlabel="Timestep t",
    ylabel="SNR(t), log scale",
)
axes[1].legend()

fig.suptitle("Noise schedules describe a nonlinear loss of signal", fontsize=14)
fig.tight_layout()
output = Path("schedule-and-snr.png")
fig.savefig(output, dpi=180, bbox_inches="tight")
print(f"saved: {output.resolve()}")

真实样本 x₀
  → 随机采样时间步 t
  → 随机采样噪声 ε
  → 闭式构造 xₜ
  → 网络预测 εθ(xₜ,t)，用 MSE(εθ, ε) 更新参数

import math
import torch

def timestep_embedding(t, dim, max_period=10_000):
    half = dim // 2
    freq = torch.exp(
        -math.log(max_period)
        * torch.arange(half, device=t.device, dtype=torch.float32)
        / max(half - 1, 1)
    )
    angles = t.float()[:, None] * freq[None, :]
    emb = torch.cat([angles.cos(), angles.sin()], dim=-1)
    if dim % 2:
        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
    return emb

import torch

def classifier_free_guidance(denoiser, xt, t, cond, null_cond, scale):
    # 两个分支合并成一个 2B batch；前半无条件，后半有条件
    model_x = torch.cat([xt, xt], dim=0)
    model_t = torch.cat([t, t], dim=0)
    model_c = torch.cat([null_cond, cond], dim=0)
    eps_uncond, eps_cond = denoiser(model_x, model_t, model_c).chunk(2)
    return eps_uncond + scale * (eps_cond - eps_uncond)

python -m pip install torch matplotlib

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"

# T = Number of Timesteps；D = Time Embedding Dimension
T = 100
D = 64
# 缩短到 100 步时不能直接照搬 1000 步线性 beta 的端点。
# 从 cosine 累计量反推出每步 beta，再重算实际累计量。
grid = torch.linspace(0, 1, T + 1, device=device, dtype=torch.float64)
f = torch.cos((grid + 0.008) / 1.008 * math.pi / 2).square()
f = f / f[0]
beta = (1 - f[1:] / f[:-1]).clamp(1e-8, 0.999).float()
alpha = 1.0 - beta
alpha_bar = torch.cumprod(alpha, dim=0)
assert alpha_bar[-1] < 1e-5  # 末端应接近采样使用的标准高斯先验
alpha_bar_prev = F.pad(alpha_bar[:-1], (1, 0), value=1.0)
posterior_var = beta * (1.0 - alpha_bar_prev) / (1.0 - alpha_bar)

def sample_data(batch_size):
    """八个高斯簇组成的二维训练分布。"""
    index = torch.randint(0, 8, (batch_size,), device=device)
    angle = index.float() * (2 * math.pi / 8)
    centers = 4.0 * torch.stack([angle.cos(), angle.sin()], dim=-1)
    return centers + 0.30 * torch.randn(batch_size, 2, device=device)

def time_embedding(t, dim=D):
    half = dim // 2
    freq = torch.exp(
        -math.log(10_000)
        * torch.arange(half, device=device)
        / (half - 1)
    )
    angles = t.float()[:, None] * freq[None, :]
    return torch.cat([angles.sin(), angles.cos()], dim=-1)

def take(values, t):
    return values.gather(0, t)[:, None]

def add_noise(x0, t, noise):
    return take(alpha_bar.sqrt(), t) * x0 + take(
        (1.0 - alpha_bar).sqrt(), t
    ) * noise

class NoisePredictor(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2 + D, 128), nn.SiLU(),
            nn.Linear(128, 128), nn.SiLU(),
            nn.Linear(128, 2),
        )

    def forward(self, xt, t):
        return self.net(torch.cat([xt, time_embedding(t)], dim=-1))

model = NoisePredictor().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3)

model.train()
for step in range(3000):
    x0 = sample_data(512)
    t = torch.randint(0, T, (x0.size(0),), device=device)
    noise = torch.randn_like(x0)
    xt = add_noise(x0, t, noise)
    predicted_noise = model(xt, t)
    loss = F.mse_loss(predicted_noise, noise)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 500 == 0:
        print(f"step={step:4d} loss={loss.item():.4f}")

@torch.no_grad()
def sample(model, count=4000):
    model.eval()
    x = torch.randn(count, 2, device=device)
    for step in reversed(range(T)):
        t = torch.full((count,), step, device=device, dtype=torch.long)
        eps = model(x, t)
        mean = (x - beta[step] / (1.0 - alpha_bar[step]).sqrt() * eps)
        mean = mean / alpha[step].sqrt()
        if step > 0:
            x = mean + posterior_var[step].sqrt() * torch.randn_like(x)
        else:
            x = mean
    return x.cpu()

real = sample_data(4000).cpu()
generated = sample(model)
fig, axes = plt.subplots(1, 2, figsize=(9, 4))
axes[0].scatter(real[:, 0], real[:, 1], s=3, alpha=0.5)
axes[0].set_title("Training data x0")
axes[1].scatter(generated[:, 0], generated[:, 1], s=3, alpha=0.5)
axes[1].set_title("Generated samples")
for ax in axes:
    ax.set_aspect("equal")
    ax.set_xlim(-6, 6)
    ax.set_ylim(-6, 6)
plt.tight_layout()
plt.savefig("diffusion_2d_result.png", dpi=160)
print("saved: diffusion_2d_result.png")

图像：       x₀, xₜ, ε, εθ  [B,C,H,W]
时间步：     t               [B]
调度表：     β, α, ᾱ        [T]
广播系数：   ᾱ[t]            [B,1,1,1]
文本上下文： c               [B,L,C_text]
潜变量：     z               [B,C_latent,H/f,W/f]