Uni-Instruct：北大与小红书联手打造的统一扩散模型蒸馏框架，一步生成超越教师模型的高质量图像

近日，来自北京大学未来技术学院、国家生物医学影像中心、小红书公司的研究团队联合发表了一篇题为《Uni-Instruct: One-step Diffusion Model through Unified Diffusion Divergence Instruction》的研究论文。这篇论文由王一飞、白卫民、张可林、张德兵、罗伟健和孙赫共同完成，于2025年5月27日发布在arXiv（arXiv:2505.20755v1）预印本平台上。这项研究提出了一种名为Uni-Instruct的统一框架，彻底改变了我们对单步扩散模型的理解和训练方式。

一、研究背景：为什么我们需要单步扩散模型？

想象一下，你有一台魔法照相机，每按一次快门就能创造出一张全新的逼真图像。这就是扩散模型的魅力所在——它能从随机噪声中逐步"提炼"出高质量的图像。然而，传统的扩散模型就像是一个复杂的化学实验，需要经过几十甚至上百步的"提纯"过程才能得到最终结果，这个过程非常耗时。

这就是为什么研究人员一直在寻找更高效的方法——能否让这个魔法照相机只需按一次快门，就能直接得到惊艳的成像结果？这就是所谓的"单步扩散模型"（one-step diffusion models）或"单步生成器"（one-step generators）。在过去几年中，众多研究团队提出了不同的方法来训练这种单步模型，比如Diff-Instruct、DMD、SIM、SiD、f-distill等等。这些方法各有特色，但都缺乏一个统一的理论框架来解释它们之间的联系。

想象这些不同的方法就像是厨师们用不同的秘方烹饪同一道菜——虽然最终都能做出美味的料理，但每个厨师都坚持自己的方法是最好的，而且不了解其他厨师的秘诀。北大和小红书的研究团队就像是一位厨艺大师，他研究了所有这些秘方，发现它们其实都是同一个基本原理的变体，然后创造了一个更全面、更强大的统一烹饪理论。

二、Uni-Instruct的理论创新：统一的扩散散度框架

Uni-Instruct的核心创新在于提出了一个基于理论的统一框架，能够将10多种现有的单步扩散蒸馏方法统一起来。就像发现了不同厨师秘方背后的共同烹饪原理一样。这个框架基于研究团队提出的"f-散度族的扩散扩展理论"（diffusion expansion theory of the f-divergence family）。

要理解这个理论，我们可以把它想象成一个翻译系统。在机器学习世界中，有许多不同类型的"散度"（divergence）用来衡量两个概率分布之间的差异。就像测量两种食谱制作出的菜肴有多相似或不同一样。常见的散度包括KL散度（Kullback-Leibler divergence）、Fisher散度等。

Uni-Instruct的天才之处在于，研究团队发现了一种方法，可以将f-散度族（一大类包括KL散度等在内的散度）沿着扩散过程展开。简单来说，他们找到了一个数学公式，可以把在原始数据空间中难以计算的f-散度，转化为沿着扩散过程的积分形式。这就像是找到了一种方法，可以把一道复杂的料理拆解成一系列简单的烹饪步骤，每个步骤都容易理解和执行。

然而，这个展开后的f-散度形式仍然难以直接优化，因为它涉及到了一些无法直接计算的项。为了解决这个问题，研究团队进一步证明了一个梯度等价定理，表明可以用另一种等价但可计算的损失函数来代替原始的展开f-散度。这就像是发现了一种更简单的烹饪方法，可以得到与复杂方法完全相同的美味结果。

这个可计算的损失函数是Diff-Instruct和SIM两种方法梯度的组合。就像是将两位顶级厨师的秘方精妙结合，取长补短，创造出更出色的料理一样。通过选择不同的f散度和权重，Uni-Instruct可以灵活地恢复现有的各种方法，如Diff-Instruct、DMD、SIM、SiD等，同时也能探索这些方法之间的新组合，创造出性能更优的单步扩散模型。

三、Uni-Instruct的算法实现：从理论到实践

Uni-Instruct的实际实现看起来很复杂，但核心思想其实很直观。想象你正在教一个学生（单步模型）模仿一位大师（多步扩散模型）的绘画技巧。传统方法可能是让学生看着大师的作品，尽可能地复制出来。而Uni-Instruct则是更系统地分析大师的绘画过程和思维方式，让学生不仅仅是模仿最终作品，而是理解整个创作过程。

具体来说，Uni-Instruct的训练过程包括三个主要角色：

首先是"教师模型"，一个预训练好的多步扩散模型，它能生成高质量的图像，但需要多步迭代。

然后是我们要训练的"学生模型"，即单步扩散模型，它的目标是一步就生成与教师模型质量相当甚至更好的图像。

最后还有一个"假分数网络"（fake score network），它的作用是帮助学生模型更好地理解教师模型的行为。

训练过程中，Uni-Instruct会维护这三个模型的状态，并交替优化它们。首先，它会使用真实数据优化判别器（用于估计密度比）；然后，优化在线假分数网络，通过分数匹配损失（score matching loss）使其行为接近教师模型；最后，使用Uni-Instruct损失优化单步扩散模型，这个损失是基于前两个模型提供的信息计算的。

值得注意的是，Uni-Instruct损失会根据所选择的散度类型而变化。例如，当选择χ?散度时，SIM部分的权重会消失，剩下的就是Diff-Instruct损失；而当选择反向KL散度（RKL）时，DI部分的权重会消失，剩下的就是SIM损失。这种灵活性使得Uni-Instruct能够统一和扩展现有的各种方法。

四、实验结果：超越教师模型的单步生成器

理论固然重要，但真正令人兴奋的是Uni-Instruct在实际生成任务上的表现。研究团队在几个标准数据集上测试了Uni-Instruct，结果令人瞩目。

在CIFAR10数据集上，Uni-Instruct创造了新的单步生成记录。对于无条件生成任务，它达到了1.46的FID（Fréchet Inception Distance，越低越好），比之前的最佳方法SiD2A（1.50）和SiDA（1.52）都要好。对于条件生成任务，它达到了1.38的FID，同样超过了之前的最佳记录1.40。

更令人惊讶的是，在ImageNet 64×64条件生成任务上，Uni-Instruct取得了1.02的FID，不仅大幅领先于其他单步方法（如SiD2A的1.10和f-distill的1.16），甚至比它的79步教师扩散模型（2.35）还要好很多，提升幅度高达1.33！这就像是学生不仅学会了老师的所有技巧，还在此基础上发扬光大，创造出比老师更出色的作品。

研究团队还进行了消融实验，探究了不同散度类型（JKL、FKL、RKL）和GAN损失对模型性能的影响。他们发现，JKL（Jeffrey-KL散度）从头训练时能达到最低的FID分数，而当从预训练的SiD模型（使用RKL）继续训练时，FKL（前向KL散度）能取得更好的结果。这表明，先用RKL进行训练直到收敛，然后再用FKL继续训练，能够同时利用RKL的寻找模式行为和FKL的覆盖模式行为，从而增强模型性能。

五、拓展应用：从图像生成到3D内容创建

Uni-Instruct的应用范围远不止于2D图像生成。研究团队受到DreamFusion和ProlificDreamer的启发，也将Uni-Instruct应用于文本到3D生成任务。

在3D生成中，他们重新实现了ProlificDreamer的代码库，在稳定扩散Unet的编码器输出添加了一个额外的判别器头。他们使用FKL蒸馏模型400个周期，通过三阶段的细化方案将NeRF转换为网格：首先使用Uni-Instruct指导训练NeRF模型300~400个周期；然后从NeRF模型获取网格表示，使用SDS损失对物体的几何外观进行微调150个周期；最后，通过Uni-Instruct指导进一步微调，再额外训练150个周期，同时使用奖励模型增强物体的外观。

实验结果表明，Uni-Instruct在3D生成中也取得了不错的效果，与之前的方法相比，生成的3D内容在质量和多样性方面都有所提高。例如，对于"一个精致的有艺术图案的花瓶"这样的提示，Uni-Instruct生成的花瓶展示了更多样的形状和更逼真的图案，比ProlificDreamer和DreamFusion的结果更加多样化和高质量。

六、Uni-Instruct的局限性与未来方向

尽管Uni-Instruct取得了令人印象深刻的成果，研究团队也坦率地指出了它的一些局限性。

首先，Uni-Instruct需要一个额外的判别器来估计密度比，这可能会带来更多的计算成本。就像是雇佣了一个额外的助手来帮助学生更好地理解大师的技巧，虽然效果好，但也增加了整个学习过程的复杂性。

其次，由于梯度公式的复杂性，如果选择不合适的f函数，Uni-Instruct可能会导致性能不佳。它复杂的梯度公式不像一些更简单的现有方法（如Diff-Instruct）那样直观。这就像是一个高级烹饪技巧，如果使用不当，可能会破坏整道菜肴。

最后，研究团队还提到了Uni-Instruct在训练速度上的局限性。由于它的复杂性，训练Uni-Instruct在2D蒸馏和文本到3D任务上需要的时间是训练DMD和ProlificDreamer的两倍。这是一个需要在未来工作中解决的问题。

展望未来，研究团队希望开发更稳定的训练技术，减少Uni-Instruct的计算成本，并将其应用于更广泛的领域，如大规模数据集和文本到图像扩散模型。他们相信，Uni-Instruct提供的理论见解和实践性能将为未来研究扩散模型的高效蒸馏和知识转移提供有益的启示。

七、总结：统一框架的力量

归根结底，Uni-Instruct的贡献在于它提供了一个理论驱动的统一框架，能够将多种现有的单步扩散蒸馏方法整合在一起。这不仅帮助我们更好地理解这些方法之间的联系，还创造了性能更优的新方法。

就像是一位大厨发现了不同烹饪技巧背后的共同原理，创造出了一种更全面、更灵活的烹饪方法，可以根据需要灵活调整，做出各种美味佳肴。Uni-Instruct框架的统一性和灵活性，使得研究人员能够更深入地理解单步扩散模型，并为未来的研究提供新的方向。

对于普通人来说，这项研究的意义在于，它让生成高质量图像和3D内容的过程变得更快更高效。未来，我们可能会看到更多基于Uni-Instruct的应用，从实时图像生成到交互式3D内容创作，让创意表达变得更加便捷和直观。

如果你对这项研究感兴趣，可以在arXiv上查阅完整论文（arXiv:2505.20755v1）。无论你是机器学习研究者还是对生成模型感兴趣的普通读者，这项研究都提供了关于如何高效训练单步扩散模型的宝贵见解。