一个S型曲线小demo，彻底搞懂扩散模型的核心机制

简介这是一份面向初学者的扩散模型演示项目内置 Jupyter 笔记本源码与演示动画通过生成 S 型曲线来展示扩散过程从缓慢起步、快速攀升到逐渐饱和的完整三阶段可帮助读者直观理解新技术或信息在群体中传播的规律。压缩包共包含 8 个文件大小约 9.74MB其中核心是一个 notebook 文件适合在 Jupyter 环境中直接运行、修改和反复调试另附一张 GIF 动图用于查看生成效果其余 xml、iml、gitignore 文件属于工程配置与版本控制辅助内容。目前已有 1615 人学习下载。跟着这个小型示例读者可以动手调节模型参数观察 S 型曲线形态的变化进而掌握扩散速度、饱和状态等关键概念同时还能将同样的建模思路迁移到市场渗透、病毒传播等真实问题中。整体体量轻巧、配置完整是连接数学公式与代码实现的理想入门材料有助于初学者快速建立从理论到实践的完整认知。1. 一条 S 型曲线讲透扩散模型这个小 demo 值不值得跑扩散模型diffusion model的入门门槛多半是被图片 demo 抬高的。MNIST、CIFAR 的生成效果确实漂亮但高维像素把加噪、去噪、噪声预测这几个核心动作全揉进了黑匣子新手跑完往往只记住网络很强大回头面对 DDPM 论文里的公式仍然一头雾水。这个小 demo 刻意反着来数据只有一条 S 型曲线上的二维点模型是一个几十 KB 的 MLPCPU 上几十秒就能完成训练每一步中间态都可以画出来看。它解决的核心问题是让你在最低计算成本下亲眼确认扩散模型前向加噪、逆向去噪到底怎么发生网络到底在学什么。适合两类人刚读完理论、想用代码验证一遍理解的初学者以及要给别人讲扩散模型、需要一份可复现教学示例的工程师。2. 扩散模型的核心机制S 型曲线为什么是理解 DDPM 的最佳教材2.1 前向过程确定性数据如何被逐步腐蚀成噪声扩散模型的前向过程forward process没有任何可学习参数它是一条固定的马尔可夫链从干净数据 x_0 出发每一步叠加一点点高斯噪声经过 T 步之后数据分布完全退化成标准正态分布。用公式说第 t 步的结果不需要真的从 x_1 一步步迭代过来DDPM 给出了一个闭式解x_t sqrt(ᾱ_t) * x_0 sqrt(1 - ᾱ_t) * ε其中 ε 是标准高斯噪声ᾱ_t 是 1 到 t 时刻 (1 - β_t) 的累积乘积β_t 是预先设计好的噪声表。这个公式是训练能高效跑起来的关键每个 batch 可以直接从 x_0 跳跃到任意第 t 步的加噪结果而不必模拟完整的前向链。import numpy as np def linear_beta_schedule(T, beta_start1e-4, beta_end0.02): 线性噪声表beta 从 beta_start 线性增长到 beta_end return np.linspace(beta_start, beta_end, T) def compute_alpha_bar(beta): alpha_bar[t] cumprod(1 - beta)[t]一次算完所有累积值 alpha 1.0 - beta return np.cumprod(alpha) def q_sample(x_0, t, alpha_bar, noiseNone): 前向加噪从干净数据 x_0 直接跳到第 t 步的加噪结果 t 可以是单个整数也可以是批量整数数组 if noise is None: noise np.random.randn(*x_0.shape) sqrt_alpha_bar np.sqrt(alpha_bar[t]) sqrt_one_minus np.sqrt(1.0 - alpha_bar[t]) x_t sqrt_alpha_bar * x_0 sqrt_one_minus * noise return x_t, noise这段代码里我按原论文惯例把 alpha_bar 预先算好q_sample 只做一次乘加操作训练开销可以忽略。注意 t 在这里是数组索引alpha_bar[t] 必须保证 t 是整数实际训练里 t 通常由 torch.randint 采样转成 long 再去索引张量这个细节后面避坑章节会重点讲。关键趋势是当 t 很小时sqrt(ᾱ_t) 接近 1数据基本保持原样当 t 接近 T 时sqrt(ᾱ_t) 接近 0x_t 几乎完全由噪声决定。这条从有到无的渐变就是前向过程可视化要展示的核心也是 2.3 里我坚持选一维数据的原因之一。2.2 逆向过程网络学的不是分布而是猜噪声逆向过程reverse process才是扩散模型真正需要训练的部分。理论上学习一个神经网络 p_θ(x_{t-1} | x_t)让它一步步把纯噪声拽回数据分布。但 DDPM 论文做了一个重要简化不直接学 x_{t-1} 的分布而是让网络去预测前向加噪时注入的噪声 ε。训练目标是让预测噪声和真实噪声尽可能接近L E[ || ε - ε_θ(x_t, t) ||² ]这里有个新手最容易误解的点网络输出的是噪声不是带噪数据本身也不是干净数据 x_0。采样时之所以能还原出曲线靠的是减去预测出的噪声这个逆向操作而不是网络直接画了一条曲线出来。import torch import torch.nn as nn class NoisePredictor(nn.Module): 极小的噪声预测网络输入 (x, t)输出与 x 同维度的噪声估计 时间步先embedding再与输入拼接是低维扩散任务最简单的结构 def __init__(self, in_dim2, hidden_dim64, T200): super().__init__() self.T T self.time_embed nn.Linear(1, hidden_dim) self.net nn.Sequential( nn.Linear(in_dim hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, in_dim), ) def forward(self, x, t): # t 归一化到 [0,1] 再进 embedding避免数值幅度差异 t_embed self.time_embed(t / float(self.T)) h torch.cat([x, t_embed], dim-1) return self.net(h)我习惯做两个处理。第一t 归一化再进网络否则 t 从 0 到 200 的大数值会让 embedding 层的梯度不稳定。第二time_embed 输出维度和 hidden_dim 保持一致concat 后不会被某一维特征主导。这个网络三层全连接参数量几千CPU 训练毫无压力。需要强调的是这里的 self.T 必须和噪声表的 T 一致。一旦改了扩散步数却忘记同步这个属性时间步条件就会失效模型退化成不分时间步的普通回归网络生成结果会明显变差这是最隐蔽的坑之一。2.3 选型理由一维数据的可视化优势与结构优势选 S 型曲线而不是随机点云有两个实际理由。第一个是它自带非线性结构x 与 y 之间是 sigmoid 映射数据点密集落成一条细带。模型真学到了分布采样点会重新落回细带附近没学会采样点会散成一团。这种学没学会一眼能看出来的特性是纯高斯点云做不到的。第二个理由是可视化。图片数据在加噪过程中只能看到整体变花但很难确认每个像素的噪声比例而 S 型曲线的二维散点图可以同时画出 x_0、x_50、x_150、x_T 四个状态肉眼直接看出曲线先变粗再溃散成圆盘的完整过程。这个 demo 把抽象的马尔可夫链变成了一帧帧能看的动画价值比代码量本身大得多。工程上还有一个附带优势S 型曲线的噪声尺度很明确。x 分布范围在 [-5, 5]y 在 [0, 1] 附近标准化到零均值、单位方差后前向噪声的幅度设计非常可控不会出现图片数据需要先归一化到 [-1, 1] 才能稳定训练的问题。这些都是这个教材级载体值得先跑一遍的原因。3. 把 demo 跑起来环境配置、训练循环与采样代码走读3.1 环境准备三个依赖CPU 就能跑这个 demo 依赖极简。PyTorch 负责模型和训练NumPy 负责噪声表和数据处理Matplotlib 负责出图。不需要 CUDACPU 完全够用我在一台 8GB 内存的办公笔记本上完整跑完训练和采样全程不到一分钟。torch 版本 1.13 以上即可因为代码只用到了 nn.Linear、ReLU、TensorDataset 这些基础组件。python -m venv .venv source .venv/bin/activate pip install torch --index-url https://download.pytorch.org/whl/cpu pip install numpy matplotlib我的习惯是把 torch 的 CPU 版本单独装避免在开发机上不小心拉下来几个 GB 的 CUDA 依赖。后续想在 GPU 上跑更大版本再把 torch 换成对应 CUDA 版本即可业务代码一行不用改。3.2 数据生成从 sigmoid 采样出 S 型点云数据生成是整份代码里最直白的部分。核心做法是在 x 轴上均匀采样经过 sigmoid 映射得到 y再叠加一点高斯噪声模拟真实数据的不确定性。噪声幅度不能设太大否则细带糊成一片S 型的形状特征就没了。def generate_s_curve_data(n5000, noise_scale0.05): 生成 S 型曲线数据x 均匀分布y sigmoid(x) 小噪声 noise_scale 控制细带的厚度建议保持 0.02~0.1 之间 x np.random.uniform(-5.0, 5.0, n) y 1.0 / (1.0 np.exp(-x)) y y np.random.randn(n) * noise_scale data np.stack([x, y], axis1).astype(np.float32) # 标准化到零均值、单位方差方便噪声表设计 mean data.mean(axis0) std data.std(axis0) data (data - mean) / std return data, mean, std标准化这步建议保留。虽然 demo 规模小、不标准化也能训但保留 mean 和 std 可以在采样完成后把结果还原回原始坐标方便和真实曲线对比。采样脚本要用这两个值做逆变换否则画出来的图坐标对不上。3.3 训练主循环时间步采样与 noise prediction loss训练循环的核心逻辑只有四步取一个 batch 的 x_0随机采样每个样本的时间步 t根据噪声表对 x_0 做前向加噪得到 x_t把 x_t 和 t 喂给网络计算预测噪声和真实噪声的 MSE。def train(model, data, T, alpha_bar, epochs300, batch_size256, lr1e-3): optimizer torch.optim.Adam(model.parameters(), lrlr) loss_fn nn.MSELoss() dataset torch.utils.data.TensorDataset(torch.from_numpy(data)) loader torch.utils.data.DataLoader(dataset, batch_sizebatch_size, shuffleTrue) for epoch in range(epochs): epoch_loss 0.0 for batch in loader: x_0 batch[0] # 每个样本独立随机采样时间步保证所有 t 都被均匀训练 t torch.randint(0, T, (x_0.shape[0], 1), dtypetorch.float32) noise torch.randn_like(x_0) alpha_bar_t torch.tensor(alpha_bar, dtypetorch.float32)[t.long().squeeze()] sqrt_alpha_bar torch.sqrt(alpha_bar_t).unsqueeze(1) sqrt_one_minus torch.sqrt(1.0 - alpha_bar_t).unsqueeze(1) x_t sqrt_alpha_bar * x_0 sqrt_one_minus * noise pred_noise model(x_t, t) loss loss_fn(pred_noise, noise) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss loss.item() * x_0.shape[0] if (epoch 1) % 50 0: print(fepoch {epoch 1:4d} loss {epoch_loss / len(dataset):.4f})两个要点。第一t 是 batch 内逐样本随机采样的不是整个 batch 共用一个 t这样每个 epoch 网络都能见到全范围时间步。第二loss 用的是原论文的简化版没有加权项低维小数据下完全够用。如果你实验发现早期 t 的训练明显滞后可以考虑给 loss 加权重但在 S 型曲线这个场景下没必要。训练轮数 300 是我的默认值。可以观察 loss 曲线提前停当 loss 跌到 0.02 以下并稳定时生成质量基本收敛再训只是浪费时间。3.4 采样过程从纯噪声一步步逆推回曲线采样是训练的反向操作。从标准正态分布初始化 x_T按照 DDPM 的更新公式逐 t 步迭代先用网络预测噪声再从 x_t 中减去这个噪声的贡献并补回一步高斯扰动。只有在 t0 的最后一步不加扰动否则生成结果会有残余噪声。def sample(model, T, alpha_bar, alpha, beta, n_samples1000): model.eval() x torch.randn(n_samples, 2) with torch.no_grad(): for t in range(T - 1, -1, -1): t_tensor torch.full((n_samples, 1), float(t)) pred_noise model(x, t_tensor) sqrt_one_minus torch.sqrt(1.0 - alpha_bar[t]) x 1.0 / torch.sqrt(alpha[t]) * ( x - beta[t] / sqrt_one_minus * pred_noise ) if t 0: x x torch.sqrt(beta[t]) * torch.randn_like(x) return x更新公式每一项都有明确来源1/sqrt(α_t) 是对当前估计的缩放β_t / sqrt(1 - ᾱ_t) 是噪声系数最后的 sqrt(β_t) * z 是这一时间步的随机性补偿其中 alpha 1.0 - beta在生成噪声表时顺手算好即可。采样时切到 eval 模式并包在 no_grad 里是必须的否则中间变量的图计算会平白多占内存。跑完 sample 之后把输出 x 用前面保存的 mean、std 还原回原始坐标散点图画出来一条 S 型曲线就在眼前了。我第一次跑通这个流程的时候盯着散点愣了好几秒——一堆随机噪声迭代 200 步之后居然真的落回了那条线上。4. 训练与采样参数噪声表、扩散步数与采样密度怎么配4.1 噪声表选型线性表稳余弦表精噪声表beta schedule是前向过程的节奏控制。线性表从 1e-4 线性涨到 0.02是 DDPM 原论文的默认配置早 t 段噪声增量小、数据结构保持完整适合大多数任务。余弦表的噪声增量中间段更平缓、两端更陡在图片任务上能改善早期信噪比下降过快的问题。def cosine_beta_schedule(T, s0.008): 余弦噪声表噪声总量不变的前提下让中间段加噪更平缓 s 是原论文的平滑系数一般保持默认 steps np.linspace(0, T, T 1) alpha_bar np.cos((steps / T s) / (1 s) * np.pi / 2) ** 2 alpha_bar alpha_bar / alpha_bar[0] beta 1.0 - alpha_bar[1:] / alpha_bar[:-1] return np.clip(beta, 0.0, 0.999)S 型曲线这个任务上线性表和余弦表的生成质量差别很小因为低维流形对信噪比衰减不敏感。我的选择是教学演示用线性表参数直观好解释如果拿这个 demo 当基底改成其他任务再换余弦表也不迟。4.2 扩散步数 T100 够演示200 是平衡点T 决定前向过程拆得有多细。T 太小比如 32每个时间步的噪声增量过大逆向每一步都在跨大步网络拟合难度明显上升采样结果偏散。T 太大比如 1000训练不会崩但每个 epoch 的时间步覆盖更稀疏同样轮数下每个 t 被训练到的次数减少小 demo 收敛速度会被拖慢。T训练开销生成质量适用场景32最低偏散形状模糊只演示前向过程100低可接受有轻微噪点快速跑通流程200中细带清晰推荐教学与参数实验1000高与 200 差别不大理论对照实验我在这个 demo 里默认 T200。改了 T 之后凡是出现 T 的地方都要同步噪声表长度、NoisePredictor 里的 self.T、采样循环的遍历次数。注意改扩散步数 T 后最容易漏的是 NoisePredictor.init里的 self.T漏改会让时间步条件失效表现就是训练正常但生成偏散。4.3 采样加速从 200 步压到 20 步的 DDIMDDPM 采样需要完整遍历 T 步T200 时已经很快但现场演示要反复采样找手感时可以用 DDIM 压缩到 20~50 步。DDIM 的思路是跳过中间节奏、每次用大步长迭代同时去掉每步的高斯扰动项让采样变成确定性过程。def sample_ddim(model, alpha_bar, n_steps20, n_samples1000, T200): DDIM 采样大步长跳跃去掉随机扰动项结果更稳定 model.eval() # 等间隔取 n_steps 个时间点去重后降序排列最后一步落到 0 times np.unique(np.linspace(0, T - 1, n_steps).astype(int))[::-1] if times[-1] ! 0: times np.append(times, 0) x torch.randn(n_samples, 2) with torch.no_grad(): for i, t in enumerate(times): t_tensor torch.full((n_samples, 1), float(t)) pred_noise model(x, t_tensor) # DDIM 公式里的 prev 是采样序列中的下一步更小的 t t_prev times[i 1] if i 1 len(times) else 0 alpha_bar_t alpha_bar[t] alpha_bar_prev alpha_bar[t_prev] # 确定性 DDIMsigma0去噪估计与噪声估计加权合成 x torch.sqrt(alpha_bar_prev) * ( (x - torch.sqrt(1 - alpha_bar_t) * pred_noise) / torch.sqrt(alpha_bar_t) ) torch.sqrt(1 - alpha_bar_prev) * pred_noise return x这段是 DDIM 论文公式的直接实现去掉了每步重采样生成过程确定性强。注意这里传给网络的 t 必须是真实的整数时间步不能传归一化后的值因为 alpha_bar 的索引依赖真实 t。DDIM 在 20 步时的结果和 200 步 DDPM 肉眼看不出明显差别现场演示用这个就对了。4.4 学习率与 batch_size哪些参数真的值得调学习率是这个 demo 里最不需要调参的参数。Adam 配 1e-3 默认学习率在小模型上收敛稳定。如果你把 hidden_dim 改大到 256 以上或换了更复杂的时间步编码可以降到 5e-4但一般没必要。真正影响收敛的是 batch_size。batch_size 太小比如 32每个 batch 的 t 覆盖不全loss 曲线抖得厉害太大比如 2048每个 epoch 更新次数太少收敛偏慢。256 是我试下来最稳的。另外不要额外加梯度裁剪小模型加上 clip 反而可能让早期时间步的噪声预测学得更慢。5. 避坑指南loss 不降、采样全噪声、曲线翻车怎么排查这一章列的是我在不同机器上复现这个 demo 时踩过的坑每条都按现象、原因、解决的顺序写。多数问题不是模型理论难而是参数和细节没对齐属于典型的低成本翻车点。5.1 loss 纹丝不动时间步传参与形状问题现象训练 50 轮后 loss 还在 1.0 附近几乎不下降网络像没学过一样。原因最常见是时间步 t 的传参形状不对。t 被 squeeze 成 [batch] 而不是 [batch, 1] 时time_embed 的输入维度对不上PyTorch 有时不报错、只是隐性广播embedding 实际失效模型分不清 t退化成普通回归网络。解决在 train 和 sample 里统一检查 model(x_t, t) 的 t 是否保持 [batch, 1] 形状。另一个排查点是 alpha_bar[t] 的索引类型t 必须转 longfloat 索引会直接报错或产生类型转换的隐性 bug。5.2 采样全是噪声逆向更新公式写反现象loss 明明降到 0.01 以下采样结果却还是一团高斯噪声完全看不出 S 型。原因八成在采样更新公式。把 1/sqrt(α_t) 写成 sqrt(α_t)或把 beta_t 和 sqrt(1 - ᾱ_t) 的位置放反都会让逆向过程越迭代越噪。另一个高频错误是最后一步没做特殊处理t0 时仍加随机扰动图形上表现为整条曲线蒙着一层白噪声。解决对照 3.4 节的 sample 函数逐行核对重点看 beta[t] 除以的是 sqrt(1 - alpha_bar[t]) 而不是 alpha[t]t0 时的随机扰动项必须跳过。5.3 曲线比真实数据胖噪声尺度与 T 不匹配现象采样点确实落在 S 型细带附近但比训练数据的细带宽很多像额外叠了一层噪声。原因训练数据 noise_scale 设得偏大前向过程最终状态与标准正态的偏差变大而采样从标准正态初始化起点就偏了。另一个原因是 T 太小比如 32每一步逆向步长过大累积误差让生成分布变散。解决把 noise_scale 调回 0.02~0.05或把 T 提到 100 以上。可以做个对比验证分别用 T32 和 T200 采样如果后者明显更紧致问题就是步长过大。5.4 多次采样结果漂移随机种子和标准化没固化现象同一个模型两次采样差异巨大一次是完整曲线一次断成两截。原因小模型对初始化敏感度比大模型高很多。不固定 seed每次训练落在不同局部最优生成形状自然漂移。此外标准化用的 mean/std 如果每次都重新计算数据本身也在漂移。解决训练和采样脚本最前面固定 seednp.random.seed(42)和torch.manual_seed(42)都写上。数据标准化用的 mean/std 保存到文件训练、采样统一读同一份。这些看着像玄学实际是小规模实验里最影响复现性的两个因素。5.5 loss 锯齿状波动batch_size 太小时间步覆盖不均现象loss 曲线像锯齿后期没有明显下降趋势收敛后仍不稳定。原因batch_size 过小导致时间步覆盖不均匀。每个 batch 只有 32 个样本时随机 t 的分布方差很大这个 batch 恰好全是小 t容易学下个 batch 全是大 t难学loss 自然上下跳。解决batch_size 提到 256。还可以固定 t 的采样方式每个 batch 先保证 t 在 [0, T) 上大致均匀取一遍再 shuffle 使用后者在大模型里不常见但在这个几百轮的小 demo 里非常有效算是我试出来的血泪经验。6. 进阶验证训练完先画三张图再决定要不要接真实项目在把这个 demo 的代码改造成自己的东西之前值得花十分钟做三件验证。第一张是前向过程中间态散点图把 x_0、x_50、x_100、x_199 四个状态并排画出确认加噪节奏符合预期。第二张是 loss 曲线确认没有 5.1 和 5.5 里的异常模式。第三张是采样结果与真实数据的叠加对比肉眼确认生成分布与真实分布的重合程度。def visualize_forward(data, alpha_bar, T, steps(0, 50, 100, 199)): 画出指定时间步的加噪结果验证前向过程节奏 import matplotlib.pyplot as plt fig, axes plt.subplots(1, len(steps), figsize(12, 3)) for ax, t in zip(axes, steps): if t 0: x_vis data[:300] else: x_t, _ q_sample(data[:300], t, alpha_bar) x_vis x_t ax.scatter(x_vis[:, 0], x_vis[:, 1], s1, alpha0.5) ax.set_title(ft{t}) plt.tight_layout() plt.savefig(forward_states.png, dpi150)这三张图比任何 loss 数值都更能说明模型是否真的学会了。如果前向过程正确你会看到散点从细带逐渐膨胀成圆盘如果采样结果落回细带逆向过程也正确。两者都对这个模型的每一步就是可解释的之后你才放心拿它当基线去替换成自己的数据和维度。我自己跑 diffusion 相关实验的习惯是无论任务是一维曲线还是后续要上图片训练完第一件事就是把可视化脚本原样贴过去先画前向中间态再画采样对比。从那以后每次训练结束我都强制走一遍三张图的验证流程再决定要不要继续调参。数据换成其他分布时这套流程同样成立先看加噪是否把结构融化干净再看采样是否把结构重建回来。希望帮到你。本文还有配套的精品资源点击获取