扩散模型蒸馏的核心奥秘:为何用更少步数生成更好图像

投资策略与投资组合管理 元宇宙与虚拟现实 区块链技术

扩散模型近年来在图像生成领域取得了令人瞩目的成就,从Stable Diffusion到DALL-E系列,这些模型能够根据文本描述生成极其逼真的图像。然而,扩散模型的一个核心特征在于其生成过程需要多次迭代 - - 通常需要数十甚至数百个步骤才能从随机噪声逐步精炼出高质量图像。这一特性既是其强大生成能力的来源,也是其计算成本居高不下的根本原因。正因如此,如何在保持生成质量的前提下大幅减少采样步数,成为了当前生成式人工智能研究中最活跃的方向之一。 理解扩散模型为何需要多步采样的关键,在于把握输入空间中高维概率分布的特殊性质。扩散模型的核心任务是学习数据分布,其训练过程通过向真实图像添加不同程度的高斯噪声,然后在每个噪声水平下训练网络预测噪声或对应的干净图像。

在采样阶段,模型从纯随机噪声出发,在每一步预测当前输入应该向哪个方向移动才能变得更像是真实图像,然后沿该方向迈出一小步。这一过程在数学上等价于求解一个常微分方程或随机微分方程。 由于高维空间中概率分布的特殊几何性质,扩散模型在每一步预测的方向仅仅在局部有意义。它指向的是高似然区域的质心,而非任何一个具体的数据点。如果在一步中跨越的距离过大,输出将变成许多可能图像的平均值 - - 在视觉上表现为模糊的图像。噪声会掩盖图像的高频信息(即精细的纹理和锐利的边缘),而对这些高频信息的预测本质上是不确定的,因此模型只能给出一个平滑的、缺乏细节的预测。

这正是为何扩散采样必须采取许多小步前进的原因 - - 每一步只能解决部分不确定性,所有步骤累积起来才能逐步揭示图像的高频细节。 这一过程在输入空间中描绘出一条弯曲的轨迹。可以将其想象为连接噪声和数据分布的曲线,扩散模型在每个点上预测曲线的切线方向,沿着这条切线方向前进一小步,然后重新评估新的位置和方向。由于曲率的存在,路径实际上是分段线性的,但当步数趋于无穷时,便接近一条真正的曲线。这意味着如果步数过少,线性近似就会与真实曲线产生较大偏差,导致生成质量急剧下降。 面对这一困境,研究者们从多个角度探索了加速采样的方法。

第一类方法是改进采样算法本身,例如采用确定性采样算法而非随机性算法,利用高阶求解器提高每步的近似精度,或者通过多步法聚合多个步骤的信息来估计更高阶的导数。Heun二阶方法通过在每步进行两次模型评估来获得更精确的更新方向,虽然每步计算量翻倍,但总体所需步数大幅减少,总计算成本反而降低。类似地,DPM-Solver系列方法利用ODE求解器的理论成果,在约十步内即可获得高质量的采样结果。 然而,仅仅改进采样算法所能获得的加速是有限的。更具变革性的方法是蒸馏 - - 训练一个新模型(学生模型)来模仿现有模型(教师模型)的行为。在传统的判别模型蒸馏中,学生模型通常比教师模型更小,从而实现模型压缩。

但在扩散模型的背景下,一种特殊的现象发生了:即使学生模型与教师模型的规模完全相同,学生模型也能以快得多的速度收敛,并能够用更少的采样步产生高质量结果。这一现象揭示了扩散模型训练中的一个内在矛盾。 标准扩散模型训练的目标是让网络预测给定噪声输入对应的干净图像。然而,我们真正希望网络预测的是条件期望,即所有可能干净图像的平均值。训练数据中的单个样本只是这个平均值的一个近似,权重更新在训练过程中不断被拉向不同的方向,导致收敛速度变慢,训练曲线剧烈波动。相比之下,蒸馏过程让学生模型直接学习教师模型对条件期望的近似,从而显著降低了损失函数的方差。

这种方差的减少使得训练过程更加平稳,收敛更快。 蒸馏之所以特别有效,远超单纯改进采样算法所能达到的效果,根本原因在于它能够将整个多步采样过程压缩到一个更少步数的模型中。当我们将多步采样的计算图展开时,它实际上构成了一个非常深的神经网络,其中许多层共享权重。学生模型的任务就是学习这个展开后的深度网络的输出,但其本身却是一个浅层网络。从这个角度看,扩散蒸馏与传统的模型压缩并无本质区别 - - 我们都是用一个更浅的网络来近似一个更深的网络。 渐进式蒸馏是一种直观的蒸馏策略。

它从训练好的教师模型开始,每一轮蒸馏都将两步采样压缩为一步。例如,一个原本需要一百步的模型,经过第一轮蒸馏后变为五十步,第二轮变为二十五步,依此类推,理论上经过约七轮蒸馏即可得到单步模型。这种方法的优势在于每一轮蒸馏只需要运行两步采样来获取训练目标,可以在训练过程中在线完成,无需预先生成大量数据。此外,由于可以使用与教师模型相同架构的学生模型并以教师模型的权重初始化,训练过程可以非常高效。实践中,通常会在蒸馏到四步或八步时停止,以在速度和质量之间取得最佳平衡。 一致性模型代表了蒸馏思想的另一种深刻诠释。

扩散模型预测路径上当前点的切线方向,而一致性模型则预测路径在数据侧的终点。这意味着从路径上的任何一点出发,一致性模型都可以一步跳到终点,完全绕开曲率带来的困难。一致性模型的训练可以通过一致性蒸馏完成 - - 利用扩散教师模型的预测来提供训练目标。为了避免运行完整采样过程带来的高昂成本,一致性模型采用了一种自举机制:要求模型在路径上的相邻两点给出相同的预测。由于这一约束,如果学生在一点上的预测是正确的,那么通过教师模型执行一步后到达的点上的结果也应该是相同的。利用这一性质,训练可以在不运行完整采样的情况下进行。

BOOT方法提供了一种与一致性模型互补的视角。如果说一致性模型是从路径上任意一点预测终点,那么BOOT则是从起点(噪声)预测路径上的任意一点。这两种方法之间存在一种类似转置的关系。正如自然语言处理中word2vec的skip-gram与连续词袋模型之间的关系一样,这两种生成路径描述方式互为镜像。BOOT的一个显著优势是它无需任何真实训练数据即可进行蒸馏 - - 由于学生模型的输入始终是噪声而非真实图像,因此可以从任意噪声出发,通过教师的引导来生成训练目标。 整流流方法则采取了另一种策略。

它认识到扩散模型学习到的连接噪声和数据分布的路径本质上具有最优传输的性质,但这些路径往往是弯曲的。整流流通过一个类似蒸馏的过程,将这些弯曲的路径替换为弯曲程度更低的新路径。具体而言,它使用教师模型生成大量数据-噪声配对,然后训练一个新模型来学习连接这些配对的更直的路径。这一过程可以递归应用,每一轮都使路径更加笔直。当路径完全笔直时,切线方向就指向终点,此时扩散模型和一致性模型在数学上变得完全等价。 在文本到图像的实际应用中,一类重要的蒸馏方法是指导蒸馏。

分类器自由指导是提升扩散模型条件生成质量的关键技术,它通过对比有条件和无条件预测来增强文本与图像的对齐程度。然而,分类器自由指导需要每个采样步骤进行两次模型评估,使采样成本翻倍。指导蒸馏的核心思想是将指导的效果蒸馏到一个单一模型中,使得该模型能够直接预测经过指导调整后的方向,而无需在采样时进行两次前向传播。这种方法虽然不能减少采样步数,但可以将每步的计算成本减半,从而实现约两倍的整体加速。 对抗蒸馏代表了另一种重要的技术路线,它借鉴了生成对抗网络的思想。在扩散模型主导图像生成领域之前,生成对抗网络曾是高质量图像生成的主流方法,其优势在于单步生成速度快且图像清晰,但代价是模式多样性不足,即生成的图像往往缺乏多样性,无法完全反映训练数据的分布。

将扩散模型的训练目标与对抗损失相结合,可以让模型在减少采样步数的同时保持视觉上的锐利度。然而,这种改进通常以牺牲样本多样性为代价 - - 模型倾向于生成更少但更逼真的图像,而非覆盖数据分布的全部范围。 对抗扩散蒸馏方法将分数蒸馏采样与对抗损失相结合,使用基于DINO特征构建的判别器来提升生成图像的真实感。这一方法被用于创建SDXL Turbo等实时文本到图像系统,能够在用户输入文本的同时逐步更新生成的图像。对抗蒸馏的一个关键优势是学生模型既支持单步采样以追求极致速度,也支持多步采样以追求更高质量,使用户可以根据需要灵活选择。 所有这些蒸馏方法都面临一个共同的权衡:减少采样步数几乎总是意味着某种程度上的质量损失。

这一损失可能表现为图像的模糊、细节的不自然,或者样本多样性的下降。没有一种蒸馏方法能够完全无损地实现任意程度的加速。然而,蒸馏的真正价值不在于消除这种权衡,而在于让我们能够主动选择权衡的方式 - - 是更看重图像的清晰度,还是样本的多样性?是追求单步生成的速度,还是保留多步迭代的能力? 扩散蒸馏领域还存在一个更深层的问题:是否真的需要教师模型?一旦通过蒸馏获得了某个具有理想特性的模型,是否可以不训练教师模型而直接从零开始训练出相同的模型?答案因方法而异。对于一致性模型,原始论文同时描述了一致性蒸馏和一致性训练两种方法,后者可以直接从数据中训练一致性模型而无需先训练扩散模型。然而,实践中一致性训练通常需要更精心的超参数调整,效果也往往不如先训练扩散模型再蒸馏的路线。这一现象与判别模型蒸馏中的发现一致:先训练大模型再蒸馏为小模型,效果通常优于直接训练小模型。

扩散模型训练的另一个有趣视角是将其视为求解最优传输问题。概率流ODE形式揭示了扩散模型学习的是噪声和数据分布之间的双射映射,而这种映射近似于最优传输映射。这意味着不同的扩散模型在相似数据上训练时,往往会学习到相似的映射 - - 它们都在试图逼近同一个最优解。这一观察为理解为什么蒸馏如此有效提供了更深的理论基础,也为寻找更高效的生成模型训练方法指明了方向。 综合来看,扩散蒸馏领域虽然方法繁多,但存在若干共同的模式和技巧。几乎所有方法都依赖确定性采样算法来从教师模型获取训练目标,其中DDIM是最常用的选择。

学生模型通常以教师模型的权重初始化,这不仅加速收敛,对某些方法甚至是必要的。许多方法采用LPIPS等感知损失来减少蒸馏对低层视觉质量的影响。自举机制 - - 让学生模型从自身学习 - - 是避免运行完整采样过程的关键技巧。指数移动平均虽然在某些方法中有帮助,但在一致性模型的后续改进中已被证明并非不可或缺。 在实际应用中,扩散蒸馏与其他建模选择存在复杂的相互作用。分类器自由指导的效力与采样步数密切相关 - - 指导通过在每步中对方向进行小调整来发挥作用,当步数过少时,指导的效果会大幅减弱。

因此,在进行步数蒸馏之前先进行指导蒸馏,往往能够更好地保留指导带来的质量提升。对于潜在扩散模型,还需要考虑蒸馏损失应该应用于潜在空间还是像素空间 - - 实践表明,像素空间的损失往往能提供更好的训练稳定性。 扩散蒸馏的兴起反映了一个更广泛的机器学习原则:先以最佳方式解决问题,再寻找可接受的捷径。扩散模型在多步迭代中展现出的强大生成能力已经被充分验证,蒸馏方法的任务是在保留核心能力的同时,找到在速度和质量之间最合适的平衡点。没有免费午餐这一铁律依然成立,但蒸馏让我们能够有意识地选择在哪里妥协,这种灵活性本身就是巨大的价值。随着研究的深入,扩散模型训练作为一种可扩展的求解最优传输问题的方法,其重要性可能会超越图像生成本身,影响到更广泛的生成式建模领域。

从更宏观的角度看,扩散蒸馏不仅是一项技术,更是一种思维方式 - - 它告诉我们,即使面对看似不可调和的矛盾,通过深入理解问题的本质,总能找到创造性的解决方案。 。