从零实现扩散模型：数学原理与PyTorch实战图像生成

1. 项目概述与核心价值最近几年AI图像生成领域最让人兴奋的突破莫过于扩散模型Diffusion Models的崛起。从DALL·E 2、Midjourney到Stable Diffusion这些能根据一句话就生成惊艳图片的工具其核心引擎都是扩散模型。但你是否和我一样最初看到那些关于“前向过程”、“反向过程”、“变分下界”的论文时感觉像在读天书这个项目就是要把这层神秘的面纱彻底揭开。“扩散模型从数学原理到图像生成的去噪扩散概率模型”这个标题精准地概括了我们要做的两件事一是理解其数学内核弄明白它到底是如何工作的二是动手实践构建一个能真正生成图像的扩散模型。这不仅仅是复现一个算法更是从最底层的数学公式出发一步步推导、编码最终见证噪声如何被“雕刻”成一张图片的完整过程。对于开发者、研究者甚至是好奇心强的爱好者来说亲手实现一遍扩散模型是理解其强大能力和局限性的最佳方式。它能让你不再只是调用API的“用户”而是成为能调整、改进甚至创造新模型的“建造者”。2. 核心思路与数学框架拆解2.1 灵感来源物理世界的扩散现象扩散模型的灵感源于热力学中的扩散过程。想象一滴墨水落入清水墨水分子会从高浓度区域清晰图像逐渐、随机地扩散到整个水中最终变成一片均匀的浑浊液体纯噪声。这个过程是前向过程Forward Process它被建模为一个马尔可夫链Markov Chain每一步都向当前图像添加一点点高斯噪声。关键在于如果我们能学会这个扩散过程的逆过程就能从一片均匀的噪声中逐步“回溯”出最初的墨水图案也就是从噪声中生成图像。这就是扩散模型最核心、也最反直觉的思想通过系统地破坏数据来学习如何构造数据。2.2 数学建模前向与反向过程我们用数学语言精确描述这个过程。给定一张真实图像x₀前向过程在T个时间步内逐步生成一系列噪声越来越大的图像x₁, x₂, ..., x_T。在t步我们根据x_{t-1}生成x_tx_t √(1 - β_t) * x_{t-1} √β_t * ε其中ε是标准高斯噪声β_t是一个预先定义好的、很小的噪声调度表Noise Schedule它随着t增大而增大控制着每一步添加的噪声量。一个重要的性质是由于每一步都是高斯噪声的叠加我们可以直接从x₀计算出任意t步的x_tx_t √(ᾱ_t) * x₀ √(1 - ᾱ_t) * ε这里α_t 1 - β_t,ᾱ_t Π_{s1}^{t} α_s。这个闭式解closed-form极大地简化了训练过程因为我们不需要真的迭代t步可以随机采样时间步t并直接计算加噪后的图像。反向过程Reverse Process则是我们的学习目标。我们需要训练一个神经网络通常是U-Net来预测给定x_t和t时前向过程中添加的噪声ε或者等价地去噪后的图像x₀或x_{t-1}的分布。这个网络被训练来最大化数据似然的变分下界ELBO经过推导这等价于一个简单的噪声预测均方误差损失L(θ) E_{x₀, ε, t} [ || ε - ε_θ(x_t, t) ||² ]其中ε_θ就是我们的神经网络。这个损失函数直观得惊人我们只是让网络学会预测我们亲手加进去的噪声。注意理解这个损失函数是理解扩散模型的关键。它不像GAN那样需要对抗训练也不像VAE需要复杂的后验分布近似。它就是一个朴素的回归任务预测噪声。这种训练的稳定性是扩散模型后来居上的重要原因。2.3 为何是U-Net网络架构的必然选择为什么扩散模型普遍使用U-Net这源于图像生成任务的内在需求。U-Net的编码器-解码器结构配合跳跃连接Skip Connections使其具备两大优势多尺度特征提取编码器逐步下采样捕获图像的全局语义信息如“这是一只猫”解码器逐步上采样结合跳跃连接传递的细节信息恢复出清晰的局部纹理如猫的胡须和毛发。条件注入时间步t的信息需要通过嵌入层如正弦位置编码或MLP注入到U-Net的每一层甚至是通过自适应组归一化AdaGN来调制特征。U-Net的模块化结构便于这种条件的灵活融入。此外为了处理更高分辨率的图像现代扩散模型如Stable Diffusion引入了交叉注意力机制Cross-Attention将文本提示prompt的语义信息注入到U-Net中实现了文生图功能。这本质上是让U-Net在去噪的每一步都“瞥一眼”文本描述确保生成的内容与之对齐。3. 从零构建一个简易扩散模型实战3.1 环境准备与数据加载我们使用PyTorch框架并选择经典的CIFAR-10数据集32x32分辨率作为起点以控制计算成本。# 环境依赖 pip install torch torchvision matplotlib numpy tqdmimport torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms from torchvision.utils import make_grid import matplotlib.pyplot as plt from tqdm import tqdm import numpy as np # 数据预处理与加载 def get_dataloader(batch_size128): transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 将像素值归一化到[-1, 1] ]) dataset datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) dataloader DataLoader(dataset, batch_sizebatch_size, shuffleTrue, num_workers2, pin_memoryTrue) return dataloader3.2 噪声调度与采样工具函数噪声调度表β_t的设计至关重要。我们采用余弦调度cosine schedule它在开始和结束时变化平缓中间变化较快经验上比线性调度能产生更好的样本质量。def cosine_beta_schedule(timesteps, s0.008): 余弦噪声调度表。 参考Improved Denoising Diffusion Probabilistic Models steps timesteps 1 x torch.linspace(0, timesteps, steps) alphas_cumprod torch.cos(((x / timesteps) s) / (1 s) * torch.pi * 0.5) ** 2 alphas_cumprod alphas_cumprod / alphas_cumprod[0] betas 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0.0001, 0.9999) def extract(a, t, x_shape): 从数组a中根据索引t获取对应的值并reshape到目标形状x_shape。 用于批量获取不同时间步的ᾱ_t, β_t等参数。 batch_size t.shape[0] out a.gather(-1, t.cpu()) return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device) # 定义前向过程加噪函数利用重参数化技巧 def q_sample(x_start, t, noiseNone): 给定x_start和t直接计算x_t。 x_t sqrt(ᾱ_t) * x_start sqrt(1 - ᾱ_t) * ε if noise is None: noise torch.randn_like(x_start) sqrt_alphas_cumprod_t extract(sqrt_alphas_cumprod, t, x_start.shape) sqrt_one_minus_alphas_cumprod_t extract(sqrt_one_minus_alphas_cumprod, t, x_start.shape) return sqrt_alphas_cumprod_t * x_start sqrt_one_minus_alphas_cumprod_t * noise3.3 U-Net噪声预测模型实现这里实现一个适用于CIFAR-10的简化版U-Net。关键点在于时间步嵌入Time Embedding的注入。class SinusoidalPositionEmbeddings(nn.Module): 将时间步t转换为固定维度的向量 def __init__(self, dim): super().__init__() self.dim dim def forward(self, time): device time.device half_dim self.dim // 2 embeddings torch.log(torch.tensor(10000.0)) / (half_dim - 1) embeddings torch.exp(torch.arange(half_dim, devicedevice) * -embeddings) embeddings time[:, None] * embeddings[None, :] embeddings torch.cat((embeddings.sin(), embeddings.cos()), dim-1) return embeddings class Block(nn.Module): 基础卷积块组归一化SiLU激活卷积 def __init__(self, in_ch, out_ch, time_emb_dimNone): super().__init__() self.mlp nn.Sequential( nn.SiLU(), nn.Linear(time_emb_dim, out_ch * 2) ) if time_emb_dim is not None else None self.conv1 nn.Conv2d(in_ch, out_ch, 3, padding1) self.norm1 nn.GroupNorm(8, out_ch) self.conv2 nn.Conv2d(out_ch, out_ch, 3, padding1) self.norm2 nn.GroupNorm(8, out_ch) self.residual_conv nn.Conv2d(in_ch, out_ch, 1) if in_ch ! out_ch else nn.Identity() def forward(self, x, tNone): h self.norm1(F.silu(self.conv1(x))) if self.mlp is not None and t is not None: time_emb self.mlp(t) time_emb time_emb[(..., ) (None, ) * 2] # reshape to [B, C, 1, 1] scale, shift time_emb.chunk(2, dim1) h h * (scale 1) shift # AdaGN操作 h self.norm2(F.silu(self.conv2(h))) return h self.residual_conv(x) class SimpleUNet(nn.Module): 简化版U-Net包含下采样和上采样 def __init__(self, in_channels3, out_channels3, base_channels64): super().__init__() time_emb_dim base_channels * 4 self.time_mlp nn.Sequential( SinusoidalPositionEmbeddings(base_channels), nn.Linear(base_channels, time_emb_dim), nn.SiLU(), nn.Linear(time_emb_dim, time_emb_dim) ) # 下采样路径 self.down1 Block(in_channels, base_channels, time_emb_dim) self.down2 Block(base_channels, base_channels*2, time_emb_dim) self.down3 Block(base_channels*2, base_channels*4, time_emb_dim) self.bottleneck Block(base_channels*4, base_channels*8, time_emb_dim) # 上采样路径 self.up3 Block(base_channels*8 base_channels*4, base_channels*4, time_emb_dim) # 跳跃连接 self.up2 Block(base_channels*4 base_channels*2, base_channels*2, time_emb_dim) self.up1 Block(base_channels*2 base_channels, base_channels, time_emb_dim) self.final_conv nn.Conv2d(base_channels, out_channels, 1) self.pool nn.MaxPool2d(2) self.upsample nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) def forward(self, x, timestep): t self.time_mlp(timestep) # 编码器 d1 self.down1(x, t) # [B, 64, 32, 32] d2 self.down2(self.pool(d1), t) # [B, 128, 16, 16] d3 self.down3(self.pool(d2), t) # [B, 256, 8, 8] b self.bottleneck(self.pool(d3), t) # [B, 512, 4, 4] # 解码器带跳跃连接 u3 self.up3(torch.cat([self.upsample(b), d3], dim1), t) # [B, 256, 8, 8] u2 self.up2(torch.cat([self.upsample(u3), d2], dim1), t) # [B, 128, 16, 16] u1 self.up1(torch.cat([self.upsample(u2), d1], dim1), t) # [B, 64, 32, 32] return self.final_conv(u1) # 预测的噪声 [B, 3, 32, 32]3.4 训练循环最小化噪声预测误差训练过程的核心就是反复执行加噪 - 预测噪声 - 计算损失 - 反向传播。def train_epoch(model, dataloader, optimizer, device, timesteps1000): model.train() total_loss 0 pbar tqdm(dataloader, descTraining) for batch_idx, (images, _) in enumerate(pbar): images images.to(device) batch_size images.shape[0] # 1. 随机采样时间步t t torch.randint(0, timesteps, (batch_size,), devicedevice).long() # 2. 采样随机噪声并计算加噪后的图像x_t noise torch.randn_like(images) x_t q_sample(images, t, noise) # 3. 神经网络预测噪声 predicted_noise model(x_t, t) # 4. 计算简单的均方误差损失 loss F.mse_loss(predicted_noise, noise) # 5. 反向传播与优化 optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() pbar.set_postfix({Loss: loss.item()}) return total_loss / len(dataloader) # 初始化与训练准备 device torch.device(cuda if torch.cuda.is_available() else cpu) timesteps 1000 # 预计算噪声调度表相关参数 betas cosine_beta_schedule(timesteps).to(device) alphas 1. - betas alphas_cumprod torch.cumprod(alphas, dim0) sqrt_alphas_cumprod torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod torch.sqrt(1. - alphas_cumprod) model SimpleUNet().to(device) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) dataloader get_dataloader() # 开始训练示例实际需要更多轮次 num_epochs 50 for epoch in range(num_epochs): avg_loss train_epoch(model, dataloader, optimizer, device, timesteps) print(fEpoch {epoch1}/{num_epochs}, Average Loss: {avg_loss:.4f}) # 可在此处添加模型保存和采样生成代码实操心得训练扩散模型是个“慢工出细活”的过程。在CIFAR-10上你可能需要训练50-100个epoch才能看到比较清晰的图像。损失值不会像分类任务那样快速降到零它会稳定在一个较低的水平。耐心是关键。另外学习率不宜过大1e-4或3e-4是常见的起点使用学习率预热Warmup和余弦衰减Cosine Decay策略通常效果更好。3.5 采样生成从噪声到图像的魔法训练完成后我们就可以运行反向过程从纯高斯噪声开始一步步去噪生成新图像。这里使用DDPM论文中的简化采样算法。torch.no_grad() def p_sample(model, x, t, t_index): 反向过程的一步采样从x_t预测x_{t-1}。 使用DDPM的简化公式。 betas_t extract(betas, t, x.shape) sqrt_one_minus_alphas_cumprod_t extract(sqrt_one_minus_alphas_cumprod, t, x.shape) sqrt_recip_alphas_t extract(torch.sqrt(1.0 / alphas), t, x.shape) # 1. 用模型预测噪声 pred_noise model(x, t) # 2. 计算x_0的估计值去噪后的图像 pred_x_start sqrt_recip_alphas_t * (x - sqrt_one_minus_alphas_cumprod_t * pred_noise) pred_x_start torch.clamp(pred_x_start, -1., 1.) # 3. 计算均值根据公式推导 model_mean sqrt_recip_alphas_t * (x - betas_t * pred_noise / sqrt_one_minus_alphas_cumprod_t) if t_index 0: return model_mean # 最后一步不加噪声 else: posterior_variance_t extract(betas, t, x.shape) # 这里简化使用β_t作为方差 noise torch.randn_like(x) return model_mean torch.sqrt(posterior_variance_t) * noise torch.no_grad() def p_sample_loop(model, shape, timesteps1000): 完整的反向采样循环从x_T ~ N(0, I) 开始逐步生成x_0。 device next(model.parameters()).device b shape[0] # 从纯噪声开始 img torch.randn(shape, devicedevice) imgs [] for i in tqdm(reversed(range(0, timesteps)), descSampling, totaltimesteps): t torch.full((b,), i, devicedevice, dtypetorch.long) img p_sample(model, img, t, i) # 可选保存中间过程观察去噪进展 if i % (timesteps // 10) 0 or i 0: imgs.append(img.cpu()) return imgs # 生成图像示例 model.eval() sample_shape (16, 3, 32, 32) # 生成16张32x32的RGB图像 generated_imgs p_sample_loop(model, sample_shape, timesteps1000) # 可视化最后生成的图像 final_imgs generated_imgs[-1] grid make_grid(final_imgs, nrow4, normalizeTrue, value_range(-1, 1)) plt.figure(figsize(10,10)) plt.imshow(grid.permute(1, 2, 0).numpy()) plt.axis(off) plt.show()当你第一次看到模糊的色块逐渐凝聚成可辨识的物体如汽车、小鸟时那种感觉是无与伦比的。这不仅仅是代码在运行更是你构建的数学模型在“创造”。4. 关键参数解析与调优经验4.1 噪声调度表Noise Schedule生成质量的节拍器β_t序列控制着噪声添加的节奏。它的设计直接影响训练稳定性和生成质量。线性调度β_t从很小的值如0.0001线性增加到较大的值如0.02。简单直观但可能导致早期去噪步信息不足晚期步噪声过大。余弦调度目前的主流选择。它让ᾱ_t信号保留系数随t呈余弦函数下降。这意味着在过程开始和结束时ᾱ_t变化较慢图像信息/噪声变化平缓在中间阶段变化较快。这为模型提供了更均衡的学习信号。平方调度β_t与t的平方成正比在早期添加噪声更快。现已较少使用。调优建议对于新任务优先使用余弦调度。你可以通过参数s如0.008来微调曲线偏移较小的s会让曲线更早地接近零保留更多初始信号。4.2 时间步数Timesteps T精度与效率的权衡更多时间步如1000前向过程更平滑反向过程的每一步需要预测的变化更小理论上训练更稳定生成质量更高。但代价是采样速度极慢生成一张图需要迭代1000次网络前向传播。更少时间步如50, 100采样速度大大加快但要求模型在每一步做出更大的预测跳跃这更困难可能导致质量下降。解决方案使用知识蒸馏或一致性模型Consistency Models技术训练一个“一步”或“少步”模型来模仿多步模型的输出从而在保持质量的同时加速采样。对于研究和初步实现1000步是标准配置便于理解原理。4.3 损失函数与训练技巧虽然基础损失是噪声的MSE但实践中有些变体和技巧v-预测参数化不直接预测噪声ε而是预测一个速度向量v定义为v ᾱ_t * ε - √(1-ᾱ_t) * x_0。论文指出这种参数化有时能带来更稳定的训练和略好的效果。学习率策略使用Warmup如前5000次迭代线性增加学习率可以避免训练初期的不稳定。配合余弦衰减让学习率在训练后期平滑下降至零。梯度裁剪虽然扩散模型训练比GAN稳定但对非常深的U-Net梯度爆炸仍可能发生。设置梯度裁剪如torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)是个好习惯。指数移动平均EMA维护一个模型权重的滑动平均版本用于最终的采样这通常能提高生成样本的稳定性和质量。# 简单的EMA实现示例 class EMA: def __init__(self, model, decay0.995): self.model model self.decay decay self.shadow {name: param.clone().detach() for name, param in model.named_parameters()} def update(self): with torch.no_grad(): for name, param in self.model.named_parameters(): self.shadow[name] self.shadow[name] * self.decay param.data * (1 - self.decay) def apply_shadow(self): with torch.no_grad(): for name, param in self.model.named_parameters(): param.data.copy_(self.shadow[name])5. 进阶探索与常见问题排查5.1 为何我的生成结果全是灰色或模糊的斑点这是新手最常见的问题通常有几个原因训练不充分扩散模型需要很长的训练时间。在CIFAR-10上Loss降到0.02以下才可能看到清晰图像。确保训练轮次足够50 epochs。数据归一化错误确认输入图像是否被正确归一化到[-1, 1]。在可视化前是否将图像反归一化回[0, 1]使用torchvision.utils.make_grid时设置normalizeTrue, value_range(-1, 1)。噪声调度过于激进如果β_t起始值太大或增长太快图像信息过早被破坏模型难以学习。尝试换用余弦调度或减小β_t的最大值。模型容量不足U-Net的通道数base_channels可能太小无法捕捉足够复杂的分布。尝试增加到128或256需考虑显存。采样过程错误检查p_sample函数中的公式是否正确特别是系数提取部分。一个常见的错误是sqrt_alphas_cumprod等张量的形状没有正确广播到输入图像x的维度。5.2 采样速度太慢有什么加速方法加速采样是扩散模型研究的热点主要有以下方向减少采样步数使用DDIMDenoising Diffusion Implicit Models采样器。DDIM允许在保持确定性相同噪声种子产生相同输出的前提下使用远少于训练步数如20-50步进行采样且质量损失很小。其核心是使用不同的反向过程方差将扩散过程重写为非马尔可夫过程。更先进的采样器DPM-Solver、UniPC等是专门设计的高阶求解器能用更少的步数达到高精度通常10-20步就能获得不错的结果。知识蒸馏训练一个“学生”网络直接学习从噪声到图像的映射或者学习“教师”扩散模型少步采样的结果实现一步或几步生成。5.3 如何扩展到更高分辨率如256x256直接训练高分辨率扩散模型对显存和计算要求极高。主流方案是采用级联扩散模型或潜在扩散模型LDM。级联模型训练多个扩散模型。第一个模型生成低分辨率如64x64图像后续模型依次对图像进行超分辨率上采样。这分而治之降低了难度。潜在扩散模型Stable Diffusion的核心这是革命性的方法。它不在像素空间操作而是在一个预训练的自编码器VAE的潜在空间中进行扩散。因为潜在空间维度远低于像素空间例如256x256x3的图像被压缩到32x32x4的潜在表示极大地减少了计算量。U-Net也主要在潜在空间中运行。生成后再用VAE的解码器转换回像素图像。5.4 如何实现文生图Text-to-Image这需要引入条件扩散模型。核心是修改U-Net使其能够接受文本描述作为额外条件。文本编码使用一个预训练的文本编码器如CLIP的文本编码器或T5将文本提示词转换为一个嵌入向量序列。条件注入最有效的方式是通过交叉注意力机制。在U-Net的瓶颈层或每层添加Cross-Attention层其中Query来自U-Net的特征Key和Value来自文本嵌入。这样去噪过程就能“关注”文本描述。Classifier-Free Guidance一种强大的技巧在采样时通过一个指导尺度guidance scale来放大文本条件的影响。它同时训练一个有条件模型和一个无条件模型采样时按ε_cond guidance_scale * (ε_cond - ε_uncond)的方向推进从而生成更贴合文本、质量更高的图像。实现文生图是一个系统工程通常基于现有大型模型如Stable Diffusion进行微调而非完全从零开始。从数学原理的推导到第一张由你代码生成的图像这个过程充满了挑战与惊喜。扩散模型的美在于它将一个复杂的生成问题优雅地转化为了一个可学习的去噪问题。虽然当前最前沿的模型包含了大量工程优化和技巧但其基石始终是本文所探讨的这些核心概念。亲手实现它不仅能让你透彻理解其工作原理更能让你具备定制和优化模型以满足特定需求的能力。在后续的探索中你可以尝试更换数据集、调整U-Net架构、实现DDIM采样甚至向潜在扩散模型迈进每一步都是对这项强大技术更深层次的驾驭。