★ FEATURED ARTICLE 1. 从零手写DDPM为什么我不建议你直接调库很多人第一次接触扩散模型都是从调用现成的管线开始的——几行代码加载预训练权重输入提示词等几十秒出图。这当然爽但爽完之后呢采样器里那些参数到底在干什么时间步长为什么是1000而不是500噪声调度器换个名字结果就天差地别这些问题一个都答不上来。我自己就经历过这个阶段用别人的轮子跑了小半年遇到出图质量崩了、训练loss不收敛完全不知道从哪下手排查。所以后来我下定决心抛开所有封装好的高层接口从最原始的DDPMDenoising Diffusion Probabilistic Models开始一行一行把前向加噪、反向去噪、训练循环、采样过程全部手写一遍。这篇文章就是那次踩坑的完整记录。我会把整个实现拆成几个核心模块讲清楚每个公式背后的直觉把我在实现过程中遇到的坑一个个标出来包括那些让我debug到凌晨三点的维度不匹配和数值溢出问题。这篇文章适合谁看如果你已经会用扩散模型出图但想搞清楚底层到底发生了什么或者你正在学深度学习想找一个足够复杂但又不会复杂到劝退的项目来练手再或者你准备做扩散模型相关的改进研究需要一个干净的baseline代码——那这篇内容应该能帮到你。我会尽量用大白话解释数学用可运行的代码片段展示实现同时把那些文档里不会写的实操细节全部抖出来。2. DDPM核心原理拆解加噪与去噪的博弈2.1 前向过程把一张图慢慢“溶解”成噪声DDPM的前向过程本质上是一个固定的马尔可夫链从真实图像x₀出发每一步都往里面加一点点高斯噪声经过T步之后图像就变成了纯噪声。这个过程可以用一个公式概括q(x_t | x_{t-1}) N(x_t; √(1-β_t) x_{t-1}, β_t I)其中β_t是一个预先定义好的方差调度序列通常从0.0001线性增长到0.02T取1000。这个公式的直觉是每一步都把上一时刻的图像乘以一个略小于1的系数让信号衰减再加上一个很小的高斯噪声。但实际实现的时候我们不会真的循环1000次去加噪——那太慢了。利用高斯分布的性质可以直接从x₀跳到任意时刻x_tq(x_t | x₀) N(x_t; √(ᾱ_t) x₀, (1-ᾱ_t) I)其中ᾱ_t ∏(1-β_i)是从0到t所有(1-β)的累乘。这个重参数化技巧是整个DDPM训练效率的关键它让我们可以在一次前向计算中就得到任意时刻的噪声图像。注意ᾱ_t在t1000时会非常接近0这意味着x_t几乎就是纯噪声了。如果你发现ᾱ_t在某个时刻变成了负数或者NaN那一定是β的调度出了问题——β不能太大否则累乘会数值下溢。我在实现的时候犯过一个低级错误把ᾱ_t的累乘写成了累加。结果就是ᾱ_t在t1000时变成了一个很大的数整个训练直接爆炸。这个bug花了我一个下午才定位到因为loss看起来只是“有点大”而不是明显的NaN。所以建议你在实现完累乘之后先打印几个关键时间步的ᾱ_t值检查一下t0时应该接近1t500时应该在0.5左右t1000时应该接近0。2.2 反向过程让神经网络学会“去噪”反向过程才是DDPM真正需要学习的地方。我们希望从纯噪声x_T出发一步步去噪最终得到一张清晰的图像。这个反向过程被建模为p_θ(x_{t-1} | x_t) N(x_{t-1}; μ_θ(x_t, t), Σ_θ(x_t, t))其中μ_θ是需要神经网络预测的均值Σ_θ是方差。DDPM的巧妙之处在于它不直接预测x_{t-1}而是预测添加到x_t中的噪声ε。具体来说给定x_t和预测的噪声ε_θ(x_t, t)我们可以计算出x_{t-1}的均值μ_θ(x_t, t) (1/√(α_t)) (x_t - (β_t/√(1-ᾱ_t)) ε_θ(x_t, t))这个公式看起来复杂但拆开看就很好理解先把x_t除以√α_t做缩放然后减去一个与预测噪声成比例的项。方差Σ_θ在原始DDPM论文中被固定为β_t或β̃_t不需要学习。训练目标也非常简洁——就是让预测噪声和真实噪声之间的均方误差最小L E[||ε - ε_θ(√ᾱ_t x₀ √(1-ᾱ_t) ε, t)||²]这个loss的美妙之处在于它只需要一次前向传播就能计算不需要像GAN那样做对抗训练也不需要像VAE那样推导复杂的变分下界。你只需要随机采样一个时间步t随机采样一个噪声ε构造噪声图像x_t让网络预测ε然后算MSE。2.3 时间步嵌入让网络知道“现在有多吵”网络需要知道当前处理的是哪个时间步的图像因为不同时间步的去噪难度完全不同。t100时图像还很清晰网络只需要做微调t900时图像几乎全是噪声网络需要做大幅度的重建。DDPM使用正弦位置编码来嵌入时间步和Transformer里的位置编码类似PE(t, 2i) sin(t / 10000^(2i/d)) PE(t, 2i1) cos(t / 10000^(2i/d))其中d是嵌入维度。这个编码的好处是不同时间步的嵌入向量之间有平滑的过渡网络可以泛化到训练时没见过的具体时间步。我试过用可学习的位置嵌入代替正弦编码效果差不多但正弦编码不需要额外的参数而且理论上可以外推到任意长度。所以如果你不是要做时间步相关的特殊分析直接用正弦编码就行。3. 网络架构选型U-Net为什么是标配3.1 U-Net的整体结构设计DDPM的噪声预测网络几乎清一色使用U-Net架构这不是没有道理的。U-Net的编码器-解码器结构配合跳跃连接天然适合这种“输入和输出尺寸相同”的任务。编码器逐层下采样提取多尺度特征解码器逐层上采样恢复空间分辨率跳跃连接把编码器的细节信息直接传到解码器避免了下采样过程中的信息丢失。具体到DDPMU-Net的输入是噪声图像x_t通道数3输出是预测噪声ε通道数也是3。时间步嵌入通过在每个残差块中注入的方式融入网络。我实现的时候用了4层下采样和4层上采样基础通道数设为64每层翻倍最大到512。提示通道数不要设得太大DDPM的训练本身就比较慢通道数翻倍会让参数量和显存占用呈平方级增长。我一开始用了128的基础通道结果单卡24G显存只能跑batch size 4训练一个epoch要几个小时。后来降到64batch size能到16训练速度直接起飞。3.2 残差块与注意力机制的配合每个U-Net层里我放了两个残差块每个残差块包含GroupNorm → SiLU激活 → 卷积 → 时间步嵌入注入 → GroupNorm → SiLU → 卷积 → 残差连接。时间步嵌入的注入方式是把嵌入向量通过一个线性层映射到和特征图通道数相同的维度然后直接加到特征图上。注意力机制我放在了16×16分辨率的层。这个分辨率下特征图的空间尺寸已经比较小了自注意力的计算量可以接受同时又能捕捉全局依赖关系。如果你在更高分辨率加注意力计算量会爆炸如果在更低分辨率加全局信息的捕捉效果又不够。我踩过一个坑一开始在所有分辨率都加了注意力结果训练速度慢了3倍显存直接爆了。后来改成只在16×16加效果几乎没降速度回来了。所以注意力不是越多越好放对位置才是关键。3.3 上采样与下采样的实现细节下采样我用的是步长为2的卷积而不是池化。卷积下采样可以学习到如何压缩信息比固定的池化操作更灵活。上采样用的是最近邻插值加卷积而不是转置卷积。转置卷积容易产生棋盘格伪影最近邻插值加卷积的组合更稳定。跳跃连接的处理也需要注意编码器的特征图和解码器的特征图在通道数上可能不一致需要用一个1×1卷积做通道对齐。我一开始忘了做这个对齐结果拼接的时候维度对不上报错信息还特别隐晦找了半天才发现是通道数的问题。4. 训练循环实现从数据加载到loss收敛4.1 数据预处理与加载训练数据我用的是一个标准图像数据集分辨率统一缩放到32×32。为什么不直接用256×256因为DDPM在32×32上训练快容易看到效果适合作为第一个手写实现。等32×32跑通了再迁移到高分辨率就是改几个参数的事。数据预处理包括随机水平翻转、归一化到[-1, 1]。归一化这一步很关键因为DDPM的前向过程假设数据分布在[-1, 1]附近。如果你用[0, 1]的输入ᾱ_t的累乘会让信号衰减得太快训练会变得不稳定。# 数据加载的核心逻辑 transform transforms.Compose([ transforms.Resize(32), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])4.2 噪声调度的实现与验证噪声调度我用了线性调度β从0.0001到0.02T1000。实现的时候需要预先计算好所有时间步的α、ᾱ、β̃等系数存成张量训练时直接索引。betas torch.linspace(0.0001, 0.02, 1000) alphas 1.0 - betas alphas_cumprod torch.cumprod(alphas, dim0) sqrt_alphas_cumprod torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod torch.sqrt(1.0 - alphas_cumprod)验证调度是否正确的方法打印t0, 250, 500, 750, 1000时的ᾱ_t值。正确的值应该分别是接近1、约0.8、约0.5、约0.2、接近0。如果ᾱ_t在t500时就已经接近0了说明β太大了需要调小。4.3 训练循环的完整实现训练循环的核心逻辑非常简洁for epoch in range(num_epochs): for x0, _ in dataloader: t torch.randint(0, 1000, (x0.shape[0],)) noise torch.randn_like(x0) x_t sqrt_alphas_cumprod[t] * x0 sqrt_one_minus_alphas_cumprod[t] * noise predicted_noise model(x_t, t) loss F.mse_loss(predicted_noise, noise) optimizer.zero_grad() loss.backward() optimizer.step()就这么简单。没有对抗训练没有复杂的loss组合就是一个MSE。但简单不代表容易收敛。我遇到的最大问题是loss在0.1左右就降不下去了生成的图像全是模糊的色块。排查过程先检查了数据归一化没问题再检查了噪声调度ᾱ_t的值也正常最后发现是学习率太大了。我一开始用了1e-3的学习率对于DDPM来说太大了网络在噪声预测上震荡。改成2e-4之后loss顺利降到了0.02以下生成的图像开始有清晰的轮廓。注意DDPM的训练对学习率非常敏感。建议从2e-4开始试如果loss震荡就降到1e-4如果收敛太慢就升到3e-4。不要用1e-3以上的学习率除非你有特殊的调度策略。4.4 训练过程中的监控指标除了loss我还监控了两个指标梯度范数和ᾱ_t的数值稳定性。梯度范数如果突然变得很大说明可能有梯度爆炸需要加梯度裁剪。ᾱ_t如果出现NaN说明数值下溢了需要检查β的调度。梯度裁剪我设的是1.0实测下来对训练稳定性有帮助。特别是在训练初期梯度范数偶尔会冲到5以上裁剪之后训练曲线平滑了很多。5. 采样过程实现从纯噪声生成图像5.1 采样算法的逐步推导采样过程是反向过程的实现。从x_T ~ N(0, I)出发对t从T-1到0逐步去噪x_t torch.randn(1, 3, 32, 32) for t in reversed(range(1000)): t_tensor torch.tensor([t]) predicted_noise model(x_t, t_tensor) alpha_t alphas[t] alpha_cumprod_t alphas_cumprod[t] beta_t betas[t] mean (1 / torch.sqrt(alpha_t)) * (x_t - (beta_t / torch.sqrt(1 - alpha_cumprod_t)) * predicted_noise) if t 0: noise torch.randn_like(x_t) variance torch.sqrt(beta_t) * noise x_t mean variance else: x_t mean这个循环跑1000步每一步都要过一次网络所以采样速度比训练慢很多。我实测下来32×32的图像在单卡上采样一张大概需要15秒。如果你要生成大量图像建议用DDIM采样可以把步数降到50-100步速度提升10倍以上质量损失很小。5.2 采样中的常见问题与修复采样阶段我遇到了两个典型问题。第一个是生成的图像有严重的网格状伪影看起来像马赛克。排查后发现是上采样用了转置卷积导致的改成最近邻插值加卷积之后就消失了。第二个问题是生成的图像颜色偏暗整体发灰。这个问题的根源是训练时的数据归一化用了[-1, 1]但采样最后没有做反归一化。加上反归一化之后颜色就正常了。还有一个比较隐蔽的问题采样时如果t0也加了噪声图像会变得很模糊。因为t0时理论上不应该再加噪声了方差应该为0。我在代码里加了判断t0时只取均值不加噪声图像清晰度立刻上了一个台阶。5.3 采样步数与生成质量的权衡完整的1000步采样虽然质量最好但速度太慢。我试过用DDIM采样50步就能达到接近1000步的效果。DDIM的核心思想是让采样过程变成确定性的不再每步都加噪声而是直接沿着概率流的方向走。DDIM的采样公式和DDPM类似但去掉了随机噪声项同时引入了一个η参数控制随机性。η0时完全确定性η1时退化为DDPM。我一般用η050步采样生成一张32×32的图像只需要0.3秒左右。提示如果你要做图像编辑或者插值DDIM的确定性采样非常有用因为相同的初始噪声会生成相同的图像方便做可控实验。6. 踩坑实录与排查技巧6.1 维度不匹配最常见的报错来源维度问题是手写DDPM时最高频的报错。我遇到过的包括时间步嵌入的维度和特征图通道数不匹配、跳跃连接拼接时通道数不一致、噪声调度的张量形状和输入不匹配。排查这类问题的技巧在关键位置打印张量的shape。比如在U-Net的每个残差块前后打印特征图的shape在时间步嵌入注入前后打印嵌入的shape。一旦发现某个位置的shape和预期不符就能快速定位到问题所在。还有一个容易忽略的点时间步t的形状。如果你用DataLoader加载数据t的batch维度要和x_t的batch维度一致。我一开始把t设成了标量结果广播的时候出了各种奇怪的问题。后来统一用t.expand(x_t.shape[0])问题就消失了。6.2 数值稳定性NaN和Inf的预防DDPM涉及大量的平方根和除法运算数值稳定性需要特别注意。ᾱ_t在t接近1000时会非常小sqrt(1-ᾱ_t)接近1这没问题。但如果你不小心把ᾱ_t算成了0sqrt(1-ᾱ_t)就是1除法的时候分母是1也没问题。真正危险的是ᾱ_t变成负数sqrt会直接产生NaN。预防措施在计算ᾱ_t的累乘时用torch.cumprod而不是手动循环cumprod的数值稳定性更好。另外可以在β的调度上加一个下限比如最小1e-5避免β0导致ᾱ_t不衰减。还有一个数值陷阱预测噪声的方差。如果网络输出的噪声值特别大mean的计算可能会溢出。我加了一个clamp操作把预测噪声限制在[-10, 10]之间实测下来对生成质量没有影响但稳定性提升了很多。6.3 训练不收敛从loss曲线看问题训练不收敛的表现有很多种loss震荡、loss下降太慢、loss降到某个值就不动了。每种表现对应的原因不同。loss震荡通常是学习率太大降到1e-4或者2e-4试试。loss下降太慢可能是网络容量不够增加通道数或者加深网络。loss卡住不动可能是数据归一化有问题检查输入是否在[-1, 1]范围内。我遇到过一个比较特殊的情况loss在前1000步下降很快然后突然跳到一个很高的值之后再也降不下来。排查后发现是某个batch的数据有问题图像全是纯黑或者纯白导致ᾱ_t的计算出现了极端值。解决方案是在数据加载时加一个过滤把方差过小的图像剔除掉。6.4 生成质量差从模糊到清晰的调优路径生成质量差的表现也分几种图像模糊、颜色失真、结构混乱。图像模糊通常是采样步数不够或者网络欠拟合增加步数或者延长训练时间。颜色失真是归一化的问题检查训练和采样的归一化是否一致。结构混乱可能是注意力机制的位置不对或者时间步嵌入没有正确注入。我调优的路径是先保证训练loss降到0.02以下然后检查采样步数是否足够最后调整网络架构。大部分生成质量问题在前两步就能解决不需要动网络结构。7. 完整代码结构与关键模块速查7.1 项目文件组织我的项目结构是这样的ddpm/ ├── model.py # U-Net模型定义 ├── diffusion.py # 噪声调度和前向/反向过程 ├── train.py # 训练循环 ├── sample.py # 采样脚本 ├── dataset.py # 数据加载 └── config.py # 超参数配置这种模块化的组织方式方便调试和复用。比如你想换一个数据集只需要改dataset.py想换一个网络架构只需要改model.py。7.2 关键超参数速查表参数名推荐值说明T1000总时间步数β_start0.0001初始噪声方差β_end0.02最终噪声方差基础通道数64U-Net第一层通道数学习率2e-4Adam优化器学习率batch_size16根据显存调整注意力分辨率16×16自注意力所在层梯度裁剪1.0防止梯度爆炸采样步数50 (DDIM)平衡速度和质量7.3 快速验证实现是否正确的方法写完代码之后不要急着跑完整训练。先用一个极小的数据集比如10张图跑100步看看loss能不能降到接近0。如果过拟合都做不到说明实现有问题。然后检查采样用训练了1000步的模型采样生成的图像应该有一些模糊的结构而不是纯噪声。如果采样出来还是纯噪声说明反向过程的公式写错了。最后检查ᾱ_t的数值打印t0, 500, 1000时的ᾱ_t确认在合理范围内。这一步能排除大部分数值问题。8. 从DDPM到Stable Diffusion的扩展思路DDPM是扩散模型的基础但直接用它生成高分辨率图像是不现实的——256×256的图像训练需要大量算力采样速度也慢得让人抓狂。Stable Diffusion的核心改进是在潜在空间做扩散而不是在像素空间。具体来说先用一个VAE把图像压缩到潜在空间比如64×64×4然后在潜在空间上训练DDPM最后用VAE解码器把潜在表示还原成图像。这个改进带来的好处是巨大的潜在空间的分辨率只有像素空间的1/4到1/8计算量直接降了一个数量级。而且VAE的压缩过程去掉了高频细节扩散模型只需要学习语义层面的信息训练也更容易收敛。如果你想从手写DDPM过渡到Stable Diffusion我建议的路径是先把DDPM在32×32上跑通然后加一个VAE把图像压缩到16×16×4在潜在空间上训练DDPM最后把VAE解码器接上。这个过程会让你对Stable Diffusion的整个管线有非常清晰的理解。还有一个值得尝试的扩展是条件生成。DDPM本身是无条件的但你可以通过交叉注意力把文本嵌入注入到U-Net中实现文本到图像的生成。这就是Stable Diffusion的核心机制之一。我试过在32×32的DDPM上加一个简单的文本编码器用CLIP的嵌入做条件生成的图像虽然分辨率低但确实能反映出文本的语义信息。最后分享一个我在实现过程中体会最深的点扩散模型的数学看起来很复杂但代码实现其实非常简洁。核心就是两个公式——前向的加噪公式和反向的去噪公式。把这两个公式理解透了剩下的都是工程细节。我一开始被那些概率符号吓到了后来把公式拆开用代码一行行实现发现每一步的直觉都很清晰。所以如果你也在手写DDPM不要被数学吓退先把代码跑起来再回头理解公式效果会好很多。
阅读完成 · 觉得有帮助？
