扩散模型连续时间框架:SDE与ODE视角及一步生成解析
扩散模型这两年在生成式AI领域的热度不用我多说从图像生成到机器人动作规划背后几乎都能看到它的影子。但很多人在上手跑通Stable Diffusion或者Diffusion Policy之后对底层的数学框架其实还是一知半解——尤其是当论文里出现SDE、ODE、概率流、一步生成这些词的时候很容易就卡住了。我自己在复现几篇扩散模型相关工作时也在这块反复绕了不少弯路后来把连续时间框架这条线彻底捋清楚之后很多之前觉得割裂的概念一下子就串起来了。这篇内容主要面向已经跑过扩散模型Demo、想进一步理解其数学骨架的开发者以及正在做扩散模型加速采样、一步生成方向研究的同学。我会从SDE和ODE两条路径出发把连续时间框架怎么搭起来、两种视角各自解决什么问题、一步生成到底是怎么实现的尽量用能落地的方式讲清楚。1. 连续时间框架到底在解决什么问题1.1 从离散到连续为什么需要换视角最早接触扩散模型的时候大家学的基本都是DDPM那套离散加噪的流程定义T个时间步每一步加一点高斯噪声前向过程是一个马尔可夫链反向过程用神经网络去预测每一步的噪声。这套框架直观、好实现但它有几个让人难受的地方。第一个问题是步数太多。DDPM标准配置是1000步采样一张图要跑1000次网络前向推理成本高得离谱。虽然DDIM把确定性采样引入进来可以把步数压到50步甚至20步但本质上还是在离散的步数框架里做文章没有从根本上改变“多步迭代”这个约束。第二个问题是理论分析不方便。离散框架下你要分析误差累积、要推导最优采样步长、要证明收敛性都得在离散的递推式上做数学处理很繁琐。而连续时间框架把离散的T步取极限变成一个连续的时间变量t从0到1的演化过程这时候就可以借用随机微分方程和常微分方程这套成熟的数学工具来分析问题了。第三个问题是灵活性。离散框架下前向加噪的schedule是固定的你想换个噪声调度策略就得重新推导整个反向过程。而连续时间框架下噪声调度变成了一个连续函数你可以自由设计反向过程的形式不依赖于具体调度这就给模型设计留出了很大的空间。我个人的体会是如果你只是调包跑推理离散框架够用了。但如果你想改采样器、想做蒸馏加速、想理解为什么某些方法能work那连续时间框架是绕不过去的。1.2 SDE视角把加噪过程写成随机微分方程连续时间框架下前向加噪过程可以写成一个随机微分方程$$dx f(x, t)dt g(t)dw$$这里f(x,t)是漂移项控制确定性的演化趋势g(t)是扩散系数控制噪声注入的强度dw是标准维纳过程的增量。这个方程的意思是在每一个微小的时间间隔dt内x的变化由两部分组成——一部分是确定性的漂移一部分是随机的噪声注入。具体到扩散模型最常见的两个SDE是VP-SDEVariance Preserving漂移项f(x,t) -½β(t)x扩散系数g(t) √β(t)。这个形式保证了这个过程的方差始终有界前向过程最终收敛到标准正态分布。DDPM对应的就是VP-SDE的离散化版本。VE-SDEVariance Exploding漂移项f(x,t) 0扩散系数g(t) √(dσ²(t)/dt)。这个形式下方差随时间不断增大最终趋于无穷。NCSNNoise Conditional Score Network对应的是这个框架。这两个SDE的选择不是随便定的它们各自对应不同的噪声调度策略也影响后续反向过程的数值稳定性。VP-SDE因为方差有界数值上更稳定所以在实际应用中更常见。1.3 ODE视角去掉随机性后的确定性演化有意思的是同一个前向加噪过程除了SDE描述之外还可以用一个ODE来描述。这个ODE被称为概率流ODEProbability Flow ODE$$dx [f(x, t) - \frac{1}{2}g(t)^2 \nabla_x \log p_t(x)]dt$$这个ODE的关键在于它和对应的SDE共享同一个边缘概率密度p_t(x)的演化轨迹。换句话说如果你从同一个初始分布出发沿着SDE走和沿着ODE走在每一个时刻tx的边缘分布是一样的。区别在于SDE的轨迹是随机的每次走出来的路径不同ODE的轨迹是确定的给定初始点路径就唯一确定了。这个性质非常重要因为它意味着你可以用ODE来做采样而ODE采样有一个巨大的好处可以用高阶数值求解器来加速。SDE采样受限于随机性一般只能用一阶方法比如Euler-Maruyama步长不能太大而ODE采样可以用Runge-Kutta这类高阶方法步长可以放得很大步数可以压得很低。提示概率流ODE的推导用到了Fokker-Planck方程核心思路是让SDE的Fokker-Planck方程和ODE的连续性方程对应起来从而解出ODE的漂移项。这个推导过程在很多综述论文里都有建议自己推一遍推完之后对SDE和ODE的关系会有完全不同的理解。2. SDE与ODE的核心细节拆解2.1 反向SDE的推导与score function的角色前向SDE是从数据分布走向噪声分布采样需要的是反向过程也就是从噪声分布走回数据分布。反向SDE的形式是$$dx [f(x, t) - g(t)^2 \nabla_x \log p_t(x)]dt g(t)d\bar{w}$$这里最关键的量是∇_x log p_t(x)也就是score function即对数概率密度对x的梯度。这个量刻画了在时刻t概率密度在x处的“陡峭程度”指向概率密度增大的方向。问题在于我们并不知道真实的数据分布p_t(x)所以score function也没法直接算。扩散模型的核心思想就是用神经网络s_θ(x,t)去估计这个score function。训练目标是最小化$$\mathbb{E}{t,x_0,\epsilon}[|\epsilon - \epsilon\theta(x_t, t)|^2]$$这就是大家熟悉的噪声预测损失。有意思的是噪声预测和score估计是等价的∇_x log p_t(x) ≈ -ε_θ(x_t,t)/σ_t。所以DDPM里训练网络预测噪声本质上就是在估计score function。这个等价关系是理解扩散模型的关键。很多人训练的时候只知道要预测噪声但不知道为什么预测噪声就能实现生成答案就在这个等价关系里。2.2 概率流ODE的推导逻辑概率流ODE的推导核心是Fokker-Planck方程。对于前向SDE其对应的Fokker-Planck方程描述了概率密度p_t(x)随时间的演化$$\frac{\partial p_t}{\partial t} -\nabla \cdot (f p_t) \frac{1}{2}g^2 \Delta p_t$$而ODE对应的连续性方程是$$\frac{\partial p_t}{\partial t} -\nabla \cdot (\tilde{f} p_t)$$让这两个方程相等可以解出ODE的漂移项$$\tilde{f}(x, t) f(x, t) - \frac{1}{2}g(t)^2 \nabla_x \log p_t(x)$$这就是概率流ODE的表达式。推导本身不复杂但背后的直觉值得琢磨ODE的漂移项等于SDE的漂移项减去一半的扩散项乘以score。这个“减去一半”的系数很关键它保证了ODE和SDE的边缘分布一致。2.3 两种视角的对比与选择依据对比维度SDE视角ODE视角轨迹性质随机每次采样路径不同确定给定初始点路径唯一数值求解器一般用一阶方法可用高阶方法采样步数通常较多数百步可压缩到很少数步到数十步理论分析适合分析随机性、方差适合分析确定性映射、可逆性似然计算需要处理随机项可用瞬时变量公式精确计算一步生成较难直接实现天然适合蒸馏加速选择哪个视角取决于你要解决什么问题。如果你关注的是生成质量的理论保证、或者需要分析采样过程的随机性SDE视角更合适。如果你关注的是采样速度、需要做蒸馏加速、或者想计算精确似然ODE视角更有优势。我自己的经验是实际做项目的时候两个视角都要掌握因为它们描述的是同一个过程的两个侧面很多问题从一个视角看很复杂换到另一个视角就豁然开朗了。3. 从多步到一步ODE框架下的加速生成3.1 为什么ODE天然适合加速ODE采样之所以能加速根本原因在于它的确定性。SDE采样每一步都注入随机噪声这个噪声会引入误差步长大了误差就大所以只能用很小的步长、很多步数来保证质量。而ODE采样没有随机项误差只来自数值离散化用高阶求解器可以把离散化误差控制得很小步长就可以放得很大。举个具体的例子用DDIM本质上是概率流ODE的一阶离散化采样20步就能出不错的结果用DPM-Solver二阶或三阶求解器10步左右就能达到类似质量再往上用高阶求解器配合好的噪声调度5步甚至更少也能work。但这里有个瓶颈即使步数压到5步也还是多步迭代。对于实时生成、边缘部署这类场景5步可能还是太慢。这就引出了一步生成的需求。3.2 蒸馏加速的核心思路一步生成的核心思路是蒸馏训练一个学生网络让它一步就能从噪声映射到数据。具体做法有几种渐进式蒸馏Progressive Distillation把教师模型的N步采样过程逐步蒸馏成N/2步、N/4步直到1步。每一步蒸馏都是训练学生网络去匹配教师网络两步采样的结果。这个方法思路清晰但需要多轮蒸馏训练成本高。一致性模型Consistency Model训练一个网络使得概率流ODE轨迹上任意两点映射到同一个初始点。这样推理的时候从噪声出发一步就能映射到数据。一致性模型的关键是自一致性约束同一条ODE轨迹上的不同点经过网络映射后应该得到相同的结果。对抗蒸馏用GAN的判别器来约束学生网络的输出分布使其匹配教师模型的输出分布。这个方法在感知质量上表现很好但训练不稳定是老大难问题。3.3 一步生成的实操要点如果你要自己实现一步生成有几个点需要特别注意教师模型的选择教师模型的质量直接决定学生模型的上限。建议先用一个训练充分的、采样质量稳定的扩散模型作为教师不要用还没收敛的模型。蒸馏损失的权重蒸馏损失通常由两部分组成——匹配教师输出的蒸馏损失和保证生成质量的辅助损失比如对抗损失或感知损失。两者的权重需要仔细调蒸馏损失权重太高会导致生成结果模糊太低又学不到教师的知识。时间步的采样策略蒸馏训练时时间步t的采样策略很关键。如果均匀采样靠近数据端t小的样本会占主导因为那里的score变化剧烈。实践中常用重要性采样让训练样本在时间轴上分布更均匀。注意一步生成模型在分布覆盖度上通常不如多步模型容易出现模式坍塌。如果你的应用场景对多样性要求高建议保留多步采样的选项或者用少步2-4步而不是严格一步。4. 实操中常见的问题与排查4.1 采样质量突然变差的排查思路这是最常见的问题之一。你之前跑得好好的采样脚本换了个模型或者换了个调度器结果出来的图要么糊要么噪。排查的时候按这个顺序来先检查噪声调度不同的SDE对应不同的噪声调度VP-SDE和VE-SDE的调度函数完全不同。如果你把VP-SDE训练出来的模型用VE-SDE的采样器去跑结果肯定不对。确认训练和采样用的是同一套调度。再检查时间步的离散化连续时间框架下时间步的离散化方式会影响采样质量。特别是靠近t0和t1的地方score function变化剧烈需要更细的离散化。如果用的是均匀步长建议在两端加密。最后检查数值稳定性ODE求解器在t接近0的时候可能会数值不稳定因为score function在那里可能趋于无穷。实践中常用的做法是在tε处截断不要真的走到t0。4.2 蒸馏训练不收敛的常见原因蒸馏训练比普通扩散模型训练更难收敛我踩过的坑主要有这几个学习率太大蒸馏训练的学习率通常要比预训练小一个数量级。因为学生网络是在拟合教师网络的输出学习率太大会导致学生网络在教师输出附近震荡。教师网络的输出噪声如果教师网络本身的采样结果就有噪声比如用了SDE采样那蒸馏目标就带有随机性学生网络很难学。建议用ODE采样得到确定性的教师输出作为蒸馏目标。判别器过强如果用对抗蒸馏判别器太强会导致生成器梯度消失。常用的技巧是判别器用谱归一化或者用WGAN-GP的损失形式。4.3 常见问题速查表问题现象可能原因排查方法解决方案采样结果全黑或全白时间步截断位置不对检查t的终止值调整截断阈值通常在0.01-0.05之间生成图像模糊蒸馏损失权重过高检查损失函数配置降低蒸馏损失权重增加感知损失采样速度没提升求解器阶数不够检查ODE求解器配置换用高阶求解器如DPM-Solver模式坍塌一步生成约束太强检查生成多样性指标改用少步生成或增加多样性正则训练loss震荡学习率太大观察loss曲线降低学习率加warmup数值溢出score function在t→0处爆炸检查中间变量数值范围在tε处截断或用对数域计算4.4 几个容易被忽略的实操细节关于噪声调度的选择VP-SDE和VE-SDE不是唯一的选择还有sub-VP SDE等变体。sub-VP SDE在保持方差有界的同时让方差随时间的衰减更平滑在某些数据集上表现更好。如果你的模型在标准VP-SDE上效果不理想可以试试sub-VP。关于求解器的选择一阶Euler求解器简单但精度低二阶Heun求解器精度高但每步需要两次网络评估。实际选择的时候要权衡如果网络评估成本高用一阶但多步如果网络评估成本低用二阶但少步。DPM-Solver系列求解器专门针对扩散模型设计利用了扩散ODE的半线性结构在少步场景下表现很好。关于时间步的嵌入方式连续时间框架下时间t是连续变量需要嵌入到网络里。常用的方式是正弦位置编码但编码的频率范围需要根据t的取值范围来设计。如果t的范围是[0,1]编码频率太高会导致网络难以学习太低又区分度不够。实践中常用的是对数间隔的频率。关于数值精度扩散模型的采样过程涉及大量小数值的累加用float32可能会有精度损失。如果发现采样结果有细微的噪声可以试试用float64跑采样过程或者用Kahan求和来减少累加误差。5. 连续时间框架的扩展应用5.1 在Diffusion Policy中的应用Diffusion Policy是把扩散模型用到机器人动作生成上的工作它的核心是把动作序列当作“图像”来生成。在连续时间框架下Diffusion Policy的采样过程同样可以用SDE或ODE来描述。用ODE采样的好处是确定性同样的观测输入会得到同样的动作输出这对机器人控制很重要——如果每次采样出来的动作都不一样机器人就会抖动。实际操作中Diffusion Policy通常用DDIM即概率流ODE的一阶离散化来采样步数在10-20步左右。如果要进一步加速可以用蒸馏把步数压到1-2步但要注意动作的平滑性可能会下降。5.2 在潜在空间扩散中的应用Stable Diffusion的核心创新之一是在潜在空间做扩散而不是在像素空间。潜在空间的维度更低扩散过程的计算量更小。在连续时间框架下潜在空间的扩散同样可以用SDE和ODE来描述而且因为维度低ODE求解器的数值稳定性更好可以用更大的步长。如果你要在自己的项目里用潜在空间扩散有几个点需要注意VAE的编码器和解码器需要和扩散模型联合调优否则潜在空间的分布可能不适合扩散潜在空间的缩放因子需要仔细调太大或太小都会影响扩散过程的稳定性。5.3 在语言模型中的扩散应用扩散模型用在文本生成上是一个比较新的方向。文本是离散的而扩散模型天然是连续的所以需要把离散的token映射到连续空间或者用离散扩散的方式。在连续时间框架下连续空间的文本扩散可以用ODE采样而且因为文本的序列结构可以用一些针对序列的求解器优化。这个方向目前还在早期但已经有一些有意思的工作。如果你对这个方向感兴趣建议先从连续空间的文本扩散入手理解清楚之后再考虑离散扩散。6. 我个人的一些经验体会绕了这么大一圈说几个我自己在实际操作中觉得最有用的点。第一个是关于理解深度和实操的关系。我一开始学扩散模型的时候觉得SDE和ODE这些数学推导离实际写代码很远后来发现完全不是这样。当你的采样结果出问题的时候如果你理解SDE和ODE的关系你就能很快定位是调度的问题、求解器的问题还是数值精度的问题。数学框架不是装饰它是排查问题的地图。第二个是关于一步生成的取舍。一步生成很诱人但它不是万能的。在我的经验里一步生成在标准数据集上比如CIFAR-10、ImageNet 64x64表现不错但在高分辨率、复杂场景下和多步采样的差距还是很明显的。如果你的应用场景对质量要求极高建议还是用少步4-8步而不是严格一步。第三个是关于工具的选择。如果你只是想做实验验证想法用现成的库比如Hugging Face的diffusers就够了它已经实现了各种SDE和ODE求解器。但如果你想深入理解建议自己从零实现一遍VP-SDE的采样过程包括score function的估计、反向SDE的离散化、概率流ODE的求解。自己写一遍比看十篇论文都管用。最后分享一个小技巧在调试ODE采样器的时候可以先把步数设得很大比如1000步确认结果正确之后再逐步减少步数观察质量下降的拐点在哪里。这个拐点就是你的求解器和调度组合能达到的最小步数。不同模型、不同数据集的拐点不一样需要自己测。