扩散模型连续时间框架：SDE与ODE原理及一步生成实践

扩散模型火到现在大家聊得最多的往往是 U-Net 结构怎么改、采样器选哪个、CFG 调到多少。但如果你翻过几篇原始论文会发现一个绕不开的东西连续时间框架下的 SDE 与 ODE。这两个缩写看起来像数学系的专属黑话实际上它们决定了你训练时的噪声调度、采样时的步数、以及为什么有些模型能一步出图。我最初接触这块的时候也是一头雾水后来把几个开源实现翻了个遍才慢慢把这条线捋清楚。这篇就按我自己的理解路径从连续时间怎么来的、SDE 和 ODE 各自扮演什么角色、一直到一步生成到底是怎么实现的尽量讲透。1. 为什么离散扩散不够用非要搞连续时间1.1 从 DDPM 的固定步数说起最早的 DDPM 把加噪过程定义成 T 步通常是 1000 步每一步加一点点高斯噪声前向过程是一个马尔可夫链。这个设定很直观代码也好写但问题在于T 是一个人为设定的超参数。你想加更多步就得重新定义整个转移核你想换一种噪声调度就得把每一步的系数重新算一遍。更麻烦的是理论分析。离散框架下你要证明某个采样器的收敛性得一步步展开T 一大推导就爆炸。而且不同的论文用不同的 T、不同的 β 调度互相之间很难直接比较——你说你的 FID 好是因为模型强还是因为步数多说不清楚。我自己的体会是离散框架像是一张固定分辨率的网格你只能在格点上做文章。而连续时间框架把这张网格变成了一条连续的曲线T 不再是必须的你可以在这条曲线上任意取点。1.2 连续时间框架的核心抽象连续时间扩散的核心思路是把离散的 t 0, 1, ..., T 换成连续的 t ∈ [0, T]前向过程从一个马尔可夫链变成一个随机微分方程SDEdx f(x, t) dt g(t) dw这里 f(x, t) 叫漂移项driftg(t) 叫扩散项diffusion coefficientdw 是标准维纳过程的增量。直观理解f 决定整体往哪个方向漂g 决定每一步抖多厉害。这个式子的好处是它把加噪过程写成了一个统一的微分方程。不同的噪声调度linear、cosine、VP、VE本质上就是选了不同的 f 和 g。你换调度不用重写整个框架只换这两个函数就行。提示VPVariance Preserving和 VEVariance Exploding是两种最常见的 SDE 形式。VP 保持方差有界VE 让方差随时间爆炸增长。Stable Diffusion 用的是 VP 那一套。1.3 反向过程才是生成的关键前向加噪谁都会真正难的是反向怎么从纯噪声一步步还原出图像。连续时间框架下反向过程也是一个 SDE但方向是时间倒流的dx [f(x, t) - g(t)² ∇_x log p_t(x)] dt g(t) dw̄这里多出来一项∇_x log p_t(x)也就是分数函数score function。它表示在时刻 t 的数据分布上往哪个方向走能让概率密度增大。这一项是未知的所以我们要训练一个神经网络 s_θ(x, t) 去逼近它。这就是扩散模型训练目标的本质不是直接学怎么去噪而是学数据分布的分数函数。理解了这一点后面 SDE 和 ODE 的关系就顺理成章了。2. SDE 与 ODE 的等价性同一个分数函数两条采样路径2.1 概率流 ODE 是怎么冒出来的反向 SDE 能采样但它有个特点每一步都带随机性。即使你从同一个初始噪声出发两次采样结果也不一样。这在某些场景下是好事多样性但在需要可复现、需要快速采样的场景下就是麻烦。于是有人问能不能把随机性去掉构造一个确定性的过程让它产生同样的边缘分布 p_t(x)答案是能。这个确定性过程就是概率流 ODEProbability Flow ODEdx [f(x, t) - ½ g(t)² ∇_x log p_t(x)] dt对比一下反向 SDE你会发现漂移项里 g(t)² 的系数从 1 变成了 ½而且没有 dw̄ 那一项了。神奇的地方在于这两个过程在任意时刻 t 的边缘分布 p_t(x) 是完全相同的。我第一次看到这个结论的时候觉得挺反直觉的——一个带噪声一个不带噪声怎么会分布一样后来想明白了SDE 里的随机性被折叠进了 ODE 的漂移项里。你可以把 SDE 看成 ODE 加上一个精心设计的噪声这个噪声恰好不改变边缘分布只改变样本的具体轨迹。2.2 两条路径的对比特性反向 SDE概率流 ODE随机性有每次采样不同无确定性边缘分布p_t(x)p_t(x)与 SDE 相同采样步数通常较多可用高阶求解器减少步数可复现性差好同输入同输出似然计算困难可用瞬时变量变换计算典型求解器Euler-MaruyamaEuler、Heun、RK45这张表是我自己在选采样方案时的决策依据。如果你要做图像编辑、插值这类需要可复现的任务ODE 明显更合适如果你要的是多样性采样SDE 更自然。2.3 为什么 ODE 能减少采样步数ODE 的一大优势是可以用高阶数值求解器。SDE 因为带随机项求解器的阶数提升很受限Euler-Maruyama 基本就是一阶而 ODE 可以用 Heun二阶、甚至 RK45自适应步长。实际测下来同样的模型SDE 采样可能要 1000 步ODE 用 Heun 求解器 50 步就能出差不多的结果。这就是为什么后来的快速采样器DPM-Solver、UniPC 等大多基于 ODE 框架——它们本质上是在用更聪明的数值方法解这个 ODE。注意ODE 步数少不代表免费。高阶求解器每一步要算多次模型前向Heun 是每步两次。所以50 步 Heun实际的计算量约等于100 步 Euler。真正的加速来自求解器阶数和步数的平衡。3. 一步生成把 ODE 的轨迹拉直3.1 一步生成为什么难理解了 ODE就能理解一步生成的野心如果从噪声到图像的映射是一个 ODE 的解那能不能训练一个网络直接一步把这个 ODE 从 tT 积到 t0难点在于原始的概率流 ODE 轨迹是弯曲的。你从 tT 的一个点出发沿着 ODE 走路径会拐来拐去。数值求解器之所以需要很多步就是因为要跟着这条弯弯曲曲的路径走。如果能让轨迹变直那一步就能到。这就是**一致性模型Consistency Models和流匹配Flow Matching**这类方法的出发点。3.2 一致性模型的核心思想一致性模型的想法很直接训练一个函数 f(x, t)使得同一条 ODE 轨迹上的任意两点经过 f 映射后都落到同一个点。也就是说f(x_t, t) f(x_{t}, t) 对同一条轨迹上的任意 t, t如果这个性质成立那我在 tT 时直接调用 f(x_T, T)就能一步得到 t0 的结果。训练时通过自一致性约束来逼近这个性质推理时就是一次前向。我试过在小型数据集上复现一致性蒸馏最大的感受是它对教师模型的 ODE 轨迹质量非常敏感。如果教师模型的采样轨迹本身就不够平滑蒸馏出来的学生模型一步生成的效果会明显掉档。3.3 流匹配与直线轨迹流匹配Flow Matching走的是另一条路不去蒸馏已有的 ODE而是直接设计一条简单的轨迹比如从噪声到数据的直线然后训练一个速度场 v(x, t) 去拟合这条直线。直线轨迹的好处是理论上一步就能积分完因为速度恒定。实际中因为网络拟合有误差通常还是要几步但比传统扩散少得多。这里有个容易混淆的点流匹配和扩散模型不是对立的。扩散模型可以看成流匹配的一个特例当轨迹是特定的弯曲路径时。两者的训练目标都可以写成回归某个向量场区别在于目标轨迹的设计。3.4 一步生成的现实边界说句实在话目前一步生成在复杂数据集比如高分辨率自然图像上质量还是打不过多步采样。原因有几个轨迹拉直会损失一部分分布拟合的精度一步映射要求网络容量足够大否则欠拟合蒸馏过程中误差会累积所以现在工业界的主流做法还是少步采样4-8 步而不是真的一步。一步生成更适合对速度极度敏感、对质量容忍度高的场景比如实时预览、草稿生成。4. 落到代码连续时间框架怎么实现4.1 噪声调度的参数化在连续时间框架下VP-SDE 通常这样参数化import torch def vp_sde_coefficients(t, beta_min0.1, beta_max20.0): # t: [0, 1] # 线性 beta 调度 beta_t beta_min t * (beta_max - beta_min) # 累积噪声量 integral_beta beta_min * t 0.5 * (beta_max - beta_min) * t**2 # 漂移和扩散系数 f -0.5 * beta_t g torch.sqrt(beta_t) # 边缘分布的均值和方差 mean_coef torch.exp(-0.5 * integral_beta) var 1 - torch.exp(-integral_beta) return f, g, mean_coef, var这段代码的关键是integral_beta它对应前向过程里累积的噪声量。mean_coef是 x_0 的衰减系数var是噪声的方差。训练时给定 x_0 和 t可以直接采样 x_tdef forward_sample(x0, t, noiseNone): if noise is None: noise torch.randn_like(x0) _, _, mean_coef, var vp_sde_coefficients(t) mean_coef mean_coef.view(-1, 1, 1, 1) std torch.sqrt(var).view(-1, 1, 1, 1) return mean_coef * x0 std * noise这个forward_sample就是连续时间框架下加噪的全部。注意 t 是连续的你可以取任意值不像离散框架只能取整数步。4.2 训练目标分数匹配训练目标是让网络预测分数函数。实践中常用的是去噪分数匹配Denoising Score Matching等价于预测噪声def training_loss(model, x0): batch_size x0.shape[0] t torch.rand(batch_size, devicex0.device) noise torch.randn_like(x0) x_t forward_sample(x0, t, noise) # 网络预测噪声 noise_pred model(x_t, t) # 加权 MSE权重来自 SDE 的扩散系数 _, g, _, var vp_sde_coefficients(t) weight var / (g**2 1e-8) loss (weight.view(-1, 1, 1, 1) * (noise_pred - noise)**2).mean() return loss这里的weight是个容易被忽略的细节。不同的 t 对应的损失量级差别很大如果不加权网络会偏向于拟合某些时间点。原始论文里这个权重是从似然下界推出来的实践中也可以简化成常数权重但效果会有差异。提示如果你发现训练后期 loss 降不下去先检查权重设计。我踩过一次坑权重写错导致模型在 t 接近 0 和 1 的区域几乎没学到东西。4.3 ODE 采样的实现概率流 ODE 的采样就是一个数值积分过程。最简单的 Euler 求解器torch.no_grad() def ode_sample(model, shape, steps50, devicecuda): x torch.randn(shape, devicedevice) dt -1.0 / steps # 从 t1 积到 t0 for i in range(steps): t 1.0 - i / steps t_tensor torch.full((shape[0],), t, devicedevice) # 网络预测噪声转成分数 noise_pred model(x, t_tensor) _, g, _, var vp_sde_coefficients(t_tensor) std torch.sqrt(var).view(-1, 1, 1, 1) score -noise_pred / (std 1e-8) # 概率流 ODE 的漂移项 f -0.5 * g**2 drift f.view(-1, 1, 1, 1) * x - 0.5 * (g**2).view(-1, 1, 1, 1) * score x x drift * dt return x这段代码里score -noise_pred / std是从预测噪声到分数函数的转换。推导很简单x_t mean_coef * x_0 std * noise对 x_t 求 log 概率的梯度就得到这个关系。换成 Heun 求解器的话每一步要算两次模型torch.no_grad() def heun_ode_sample(model, shape, steps25, devicecuda): x torch.randn(shape, devicedevice) dt -1.0 / steps for i in range(steps): t 1.0 - i / steps t_next t dt # 第一次评估 drift1 compute_drift(model, x, t, device) x_pred x drift1 * dt # 第二次评估用预测点 drift2 compute_drift(model, x_pred, t_next, device) # 平均 x x 0.5 * (drift1 drift2) * dt return xHeun 的精度明显高于 Euler25 步 Heun 的效果通常能打平 100 步 Euler。代价是每步两次前向所以实际加速比大概是 2 倍左右。4.4 一步生成的蒸馏实现一致性蒸馏的核心是让学生的输出在轨迹上保持一致。简化版的训练逻辑def consistency_distillation_loss(student, teacher, x0): t torch.rand(x0.shape[0], devicex0.device) noise torch.randn_like(x0) x_t forward_sample(x0, t, noise) # 教师用 ODE 走一步得到 x_{t} with torch.no_grad(): x_next teacher_ode_step(teacher, x_t, t) # 学生分别对 x_t 和 x_next 做预测要求一致 pred_t student(x_t, t) pred_next student(x_next, t - 1/steps) loss ((pred_t - pred_next)**2).mean() return loss实际实现里还有 EMA 目标网络、边界条件约束等细节但核心就是这个自一致性。训练收敛后推理时直接student(x_T, T)一步出结果。5. 实操中容易踩的几个坑5.1 时间步的采样分布训练时 t 从 [0, 1] 均匀采样看起来最自然但实际效果往往不好。因为 t 接近 0 时噪声很小网络学到的信息量少t 接近 1 时噪声太大梯度信号也弱。很多实现会用对数正态分布或者重要性采样来调整 t 的分布让训练更集中在信息量大的区间。我自己的经验是在小数据集上均匀采样还能凑合一旦数据复杂起来t 的采样策略对最终质量影响很大。这个细节在论文里经常一笔带过但复现时不能忽略。5.2 数值稳定性连续时间框架里有很多除法比如 score -noise_pred / std。当 t 接近 0 时 std 接近 0这个除法会炸。标准做法是加一个小的 epsilon或者对 t 做截断比如只采到 t0.001 而不是 t0。另一个坑是 ODE 求解器的步长。如果步长太大Euler 会发散如果太小计算量爆炸。实践中建议先用固定步数跑通再考虑自适应步长。5.3 训练与采样的时间定义不一致这是最隐蔽的坑。有些代码库训练时 t 从 0 到 10 是干净数据有些反过来0 是纯噪声。如果你混用了两个库的组件采样出来的结果会完全乱掉。我的建议是在代码里显式注释清楚 t 的方向并且在写采样循环时反复确认积分方向。我自己就因为这个问题浪费过一整天采样出来全是噪声最后发现是 t 的定义反了。5.4 蒸馏时的教师质量做一步生成蒸馏时教师模型的 ODE 轨迹质量直接决定学生上限。如果教师本身采样就有 artifacts学生只会学得更差。所以蒸馏前一定要先把教师模型调到最佳状态别指望学生能青出于蓝。6. 不同场景下该选 SDE 还是 ODE6.1 图像生成与编辑纯生成任务如果追求质量和多样性SDE 采样更稳。如果追求速度和可复现ODE 更合适。图像编辑比如 SDEdit 那类几乎都用 ODE因为需要保持输入和输出的对应关系。6.2 似然计算与密度估计需要计算数据似然的场景只能用 ODE。因为概率流 ODE 是一个确定性映射可以用瞬时变量变换公式精确计算 log p(x)。SDE 的似然计算要复杂得多通常只能给下界。6.3 实时应用实时场景比如交互式生成优先考虑 ODE 高阶求解器或者直接用蒸馏后的一步/少步模型。SDE 的随机性在实时场景里反而是负担。6.4 我的实际选择做研究复现时我一般先用 ODE Heun 跑通因为它确定性强调试方便。等流程稳定了再切到 SDE 看多样性。做产品原型时直接上蒸馏模型牺牲一点质量换速度。7. 从连续时间视角重新理解扩散模型把 SDE 和 ODE 这条线捋清楚之后我对扩散模型的理解变了不少。以前觉得它是一堆加噪去噪的步骤堆起来的现在更愿意把它看成一个连续时间上的概率流。训练是在学这个流的分数场采样是在解这个流的方程加速是在想办法把流拉直。这个视角的好处是很多看似无关的方法DDIM、DPM-Solver、一致性模型、流匹配都能放到同一个框架下理解。DDIM 本质上是概率流 ODE 的一个特定离散化DPM-Solver 是更高阶的求解器一致性模型是在蒸馏 ODE 的解流匹配是在重新设计 ODE 的轨迹。理解了这层关系再看新论文的时候就不会被各种名词绕晕了——它们大多是在这个连续时间框架里做文章要么改轨迹要么改求解器要么改训练目标。最后分享一个我自己调试时的小习惯遇到采样结果不对先画一下 ODE 轨迹的范数随时间的变化曲线。如果曲线在某段突然跳变多半是那段的数值求解出了问题或者网络在那个时间区间没学好。这个检查比盲目调参高效得多。