AI在单一步骤中以30倍的速度生成高质量图像

在我们当前的人工智能时代,计算机可以通过扩散模型生成自己的“艺术”,通过迭代地为嘈杂的初始状态添加结构,直到清晰的图像或视频出现。扩散模型突然成为每个人的焦点:输入几个词,即可在现实与幻想的交汇处体验到即时、多巴胺上升的梦境。在幕后,这涉及到一个复杂、耗时的过程,需要多次迭代才能使算法完善图像。

麻省理工学院计算机科学与人工智能实验室(CSAIL)的研究人员提出了一个新的框架,将传统的扩散模型的多步骤过程简化为单步骤,解决了以前的限制。这是通过一种师生模型来实现的:教授一个新的计算机模型模仿生成图像的更复杂的原始模型的行为。这种方法被称为“分布匹配蒸馏”(DMD),保留了生成图像的质量,并且可以更快地生成。

“我们的工作是一种新颖的方法,将当前的稳定扩散和DALLE-3等扩散模型加速了30倍,”麻省理工学院电气工程和计算机科学博士生、CSAIL成员和DMD框架的首席研究员Tianwei Yin说道。“这一进展不仅显著减少了计算时间,而且保留了生成视觉内容的质量,甚至超过了质量。从理论上讲,这种方法将生成对抗网络(GANs)的原则与扩散模型的原则结合起来,实现了在单个步骤中生成视觉内容,这与当前扩散模型所需的数百个迭代改进的步骤形成了鲜明对比。它有可能成为一种在速度和质量方面都出色的新的生成建模方法。”

这种单步骤扩散模型可以增强设计工具,实现更快的内容创建,并可能支持药物发现和3D建模的进展,其中迅速和有效性至关重要。

分布梦想

DMD巧妙地有两个组成部分。首先,它使用回归损失,通过锚定映射来确保图像空间的粗略组织,使训练更加稳定。接下来,它使用分布匹配损失,确保使用学生模型生成给定图像的概率与其在现实世界中出现的频率相对应。为了做到这一点,它利用两个扩散模型作为指导,帮助系统理解真实图像和生成图像之间的差异,并使训练速度快的单步骤生成器成为可能。

该系统通过训练一个新的网络来最小化其生成图像与传统扩散模型使用的训练数据集中的图像之间的分布差异,从而实现更快的生成。“我们的关键洞察是使用两个扩散模型来近似指导新模型改进的梯度,”Yin说。“通过这种方式,我们将原始的更复杂模型的知识蒸馏到更简单、更快的模型中,同时避免了GANs中臭名昭著的不稳定性和模式崩溃问题。”

Yin和他的同事们使用预训练的网络来进行新的学生模型,简化了这个过程。通过复制和微调原始模型的参数,团队实现了新模型的快速训练收敛,该模型能够以相同的架构基础生成高质量的图像。“这使得可以与基于原始架构的其他系统优化相结合,进一步加速创建过程,”Yin补充道。

在与常规方法进行广泛的基准测试时,DMD表现出一致的性能。在基于ImageNet特定类别生成图像的流行基准测试中,DMD是第一个与原始的更复杂模型几乎相当的单步骤扩散技术,其Fréchet inception distance(FID)得分仅为0.3,这是令人印象深刻的,因为FID是评判生成图像的质量和多样性的指标。此外,DMD在工业规模的文本到图像生成方面表现出色,并实现了最先进的单步骤生成性能。在处理更棘手的文本到图像应用时仍存在一些质量差距,这表明未来还有改进的空间。

此外,DMD生成的图像的性能与蒸馏过程中使用的教师模型的能力密切相关。在目前的形式中,使用稳定扩散v1.5作为教师模型,学生继承了一些限制,例如渲染文本和小脸部的详细描绘,这表明通过更先进的教师模型可以进一步提高DMD生成的图像。

“自扩散模型诞生以来,减少迭代次数一直是其圣杯,”麻省理工学院电气工程和计算机科学教授、CSAIL首席研究员和论文的主要作者Fredo Durand说。“我们非常兴奋地终于实现了单步骤图像生成,这将大大降低计算成本并加速过程。”

“终于,一篇成功地将扩散模型的多样性和高视觉质量与GANs的实时性能相结合的论文,”加州大学伯克利分校电气工程和计算机科学教授Alexei Efros说道,他没有参与这项研究。“我期待这项工作为高质量的实时视觉编辑开辟了奇妙的可能性。”

Yin和Durand的合著者还包括麻省理工学院电气工程和计算机科学教授、CSAIL首席研究员William T. Freeman,以及Adobe研究科学家Michaël Gharbi SM ’15,PhD ’18;Richard Zhang;Eli Shechtman;和Taesung Park。他们的工作得到了美国国家科学基金会的资助(包括人工智能和基本相互作用研究所的资助)、新加坡国防科学与技术局的资助,以及光州科学技术研究院和亚马逊的资助。他们的工作将于6月在计算机视觉和模式识别会议上展示。