在当今的人工智能时代,计算机可以通过扩散模型生成自己的“艺术”,迭代地将结构添加到嘈杂的初始状态,直到出现清晰的图像或视频。扩散模型突然在每个人的餐桌上占据了一席之地:输入几句话,在现实与幻想的交汇处体验瞬间的、多巴胺激增的梦境。在幕后,它涉及一个复杂、耗时的过程,需要算法进行多次迭代才能完善图像。
麻省理工学院计算机科学和人工智能实验室(CSAIL)的研究人员引入了一个新框架,将传统扩散模型的多步骤过程简化为一个步骤,解决了以前的局限性。这是通过一种师生模型来完成的:教授一种新的计算机模型来模仿生成图像的更复杂的原始模型的行为。这种方法称为分布匹配蒸馏 (DMD),可保留生成图像的质量,并允许更快地生成图像。
“我们的工作是一种新颖的方法,可以将当前的扩散模型(如稳定扩散和DALLE-3)加速30倍,”麻省理工学院电气工程和计算机科学博士生Tianwei Yin说,CSAIL附属机构,DMD框架的首席研究员。“这一进步不仅大大减少了计算时间,而且还保持了(如果不是超过)生成的视觉内容的质量。从理论上讲,该方法将生成对抗网络 (GAN) 的原理与扩散模型的原理相结合,一步到位即可生成视觉内容——这与当前扩散模型所需的数百步迭代细化形成鲜明对比。它可能是一种在速度和质量方面表现出色的新型生成建模方法。
这种单步扩散模型可以增强设计工具,实现更快的内容创建,并可能支持药物发现和 3D 建模的进步,其中及时性和有效性是关键。
分销梦想
DMD巧妙地有两个组成部分。首先,它使用回归损失,锚定映射以确保图像空间的粗略组织,使训练更加稳定。接下来,它使用分布匹配损失,确保使用学生模型生成给定图像的概率与其真实世界的发生频率相对应。为此,它利用两个扩散模型作为指导,帮助系统了解真实图像和生成图像之间的区别,并使训练快速的一步生成器成为可能。
该系统通过训练新网络来实现更快的生成,以最小化其生成的图像与传统扩散模型使用的训练数据集中的图像之间的分布差异。“我们的关键见解是近似梯度,这些梯度使用两个扩散模型来指导新模型的改进,”Yin说。“通过这种方式,我们将原始的、更复杂的模型的知识提炼成更简单、更快的模型,同时绕过GAN中臭名昭著的不稳定性和模式崩溃问题。
Yin及其同事将预训练网络用于新的学生模型,简化了这一过程。通过对原始模型的参数进行复制和微调,该团队实现了新模型的快速训练收敛,能够生成具有相同架构基础的高质量图像。“这样可以与基于原始架构的其他系统优化相结合,以进一步加快创建过程,”Yin 补充道。
当使用各种基准测试与常规方法进行测试时,DMD 表现出一致的性能。在基于ImageNet上的特定类生成图像的流行基准测试中,DMD是第一个一步扩散技术,它产生的图片几乎与原始的、更复杂的模型中的图片相当,其超近距离Fréchet初始距离(FID)得分仅为0.3,这令人印象深刻,因为FID是关于判断生成图像的质量和多样性的。此外,DMD 在工业规模的文本到图像生成方面表现出色,并实现了最先进的一步生成性能。在处理更棘手的文本到图像应用程序时,仍然存在轻微的质量差距,这表明未来仍有一些改进空间。
此外,DMD生成图像的性能与蒸馏过程中使用的教师模型的功能有着内在的联系。在当前使用Stable Diffusion v1.5作为教师模型的形式中,学生继承了诸如渲染文本和小脸的详细描述等局限性,这表明DMD生成的图像可以通过更高级的教师模型进一步增强。
“自扩散模型诞生以来,减少迭代次数一直是扩散模型的圣杯,”麻省理工学院电气工程和计算机科学教授、CSAIL首席研究员、该论文的主要作者Fredo Durand说。“我们很高兴最终能够实现单步映像生成,这将大大降低计算成本并加快这一过程。”
“最后,这篇论文成功地将扩散模型的多功能性和高视觉质量与GAN的实时性能相结合,”加州大学伯克利分校电气工程和计算机科学教授Alexei Efros说,他没有参与这项研究。“我希望这项工作能够为高质量的实时视觉编辑开辟奇妙的可能性。”
Yin 和 Durand 的共同作者是麻省理工学院电气工程和计算机科学教授和 CSAIL 首席研究员 William T. Freeman,以及 Adobe 研究科学家 Michaël Gharbi SM '15, PhD '18;理查德·张;伊莱·谢赫特曼;和泰成公园。他们的工作部分得到了美国国家科学基金会(U.S. National Science Foundation)的资助(包括人工智能与基础交互研究所(Institute for Artificial Intelligence and Fundamental Interactions)、新加坡国防科技局(Singapore Defense Science and Technology Agency)的资助,以及光州科学技术学院(Gwangju Institute of Science and Technology)和亚马逊(Amazon)的资助。他们的工作将在6月的计算机视觉和模式识别会议上发表。