从零实现Diffusion模型：200行代码理解AI绘画核心原理

1. 项目概述：为什么从零实现Diffusion模型值得尝试

2022年至今，Stable Diffusion等AI绘画模型已经改变了数字内容创作的方式。但大多数使用者仅停留在调用API或使用现成工具的层面，对底层原理知之甚少。实际上，亲手实现一个基础版Diffusion模型是理解现代生成式AI的最佳实践路径。

我在计算机视觉领域工作多年，第一次完整实现Diffusion模型时，仅200行的PyTorch代码就让我对噪声预测、时间步调度等核心概念有了颠覆性认知。这远比阅读十篇论文来得深刻。本文将带你用约300行Python代码，构建一个能生成28x28手写数字的简易Diffusion模型。

提示：完整代码已开源，建议配合Jupyter Notebook边读边实践。即使没有GPU，也可以在Colab上免费运行。

2. 核心原理拆解：Diffusion模型如何"去噪"

2.1 前向扩散：数据破坏的艺术

前向扩散过程本质是逐步向图像添加高斯噪声。设原始图像为x₀，经过T步扩散后得到纯噪声x_T。每步噪声添加遵循：

python复制def forward_diffuse(x0, t):
    """计算第t步的噪声图像"""
    sqrt_alpha = torch.sqrt(alphas[t])  # 噪声调度系数
    sqrt_one_minus_alpha = torch.sqrt(1 - alphas[t])
    noise = torch.randn_like(x0)  # 标准高斯噪声
    return sqrt_alpha * x0 + sqrt_one_minus_alpha * noise

其中alphas是预设的噪声调度数组，控制噪声添加的节奏。我常用cosine调度，它在开始和结束时的变化更平缓。

2.2 逆向生成：噪声预测是关键

模型需要学会预测给定噪声图像x_t和时间步t时的噪声ε。训练目标函数为：

math复制L = ||ε - ε_θ(x_t, t)||^2

实现时，UNet的输入是噪声图像和嵌入后的时间步，输出是与输入同尺寸的噪声图。这个设计让模型专注于学习噪声模式而非直接生成图像。

2.3 采样过程：从噪声中迭代重建

采样时，我们从纯噪声x_T出发，逐步去噪：

python复制for t in reversed(range(T)):
    predicted_noise = model(xt, t)
    xt = 1/torch.sqrt(alpha[t]) * (xt - (1-alpha[t])/torch.sqrt(1-alpha_bar[t]) * predicted_noise)
    if t > 0:
        xt += torch.sqrt(beta[t]) * torch.randn_like(xt)  # 添加随机性

其中alpha_bar是alphas的累积乘积。这个过程如同将模糊的照片逐渐对焦。

3. 代码实现详解：200行构建可运行模型

3.1 数据准备：MNIST的预处理技巧

使用MNIST数据集时，我推荐以下处理：

python复制transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: (x - 0.5) * 2)  # 归一化到[-1,1]
])

注意：保持输入输出在同一数值范围(-1到1)对模型收敛至关重要。许多训练失败案例源于错误的归一化。

3.2 网络设计：轻量级UNet实现

我们的微型UNet包含：

python复制class TimeEmbedding(nn.Module):
    def forward(self, t):
        # 将时间步转换为128维向量
        return sinusoidal_embedding(t, dim=128)

class Block(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.GroupNorm(8, out_ch),  # 比BN更适合小batch
            nn.SiLU()
        )

完整UNet约50万参数，在Colab上训练1小时即可得到不错效果。

3.3 训练技巧：三个关键超参数

1. 学习率：从3e-4开始，配合OneCycle调度
2. Batch Size：至少32，太小会导致噪声预测不稳定
3. 时间步T：400-1000步为宜，太少降质，太多浪费算力

我的最佳实践配置：

python复制optimizer = AdamW(model.parameters(), lr=3e-4)
scheduler = OneCycleLR(optimizer, max_lr=3e-4, total_steps=5000)

4. 效果优化与问题排查

4.1 生成质量诊断表

4.2 高级技巧：Classifier-Free Guidance

虽然基础版不需要条件输入，但通过添加10%的随机dropout可以提升效果：

python复制# 训练时随机丢弃标签
if random() < 0.1:
    class_label = None

采样时通过guidance_scale控制条件强度：

python复制noise_pred = unconditional_pred + guidance_scale * (conditional_pred - unconditional_pred)

5. 从玩具模型到实用工具

完成基础实现后，你可以尝试：

1. 替换为Stable Diffusion的VAE
2. 增加注意力机制
3. 尝试更高分辨率

我在本地用RTX 3090训练256x256模型时，发现以下改进最有效：

• 使用EMA模型权重
• 添加自适应组归一化(AdaGN)
• 采用v-prediction参数化

重要提醒：扩散模型对超参数极其敏感。建议先用小规模实验确定配置，再放大训练。我的第一个256px模型因学习率过高浪费了三天算力。

实现过程中最让我意外的发现是：即使简单的线性噪声调度也能工作，但cosine调度可以节省30%训练时间。这印证了论文《Improved Denoising Diffusion Probabilistic Models》中的结论——噪声调度比想象中更重要。