【深度学习】Consistency Models

Posted by ShawnD on June 5, 2023

Abstract

扩散模型在图像、音频和视频生成方面取得了重大突破,但它们依赖于迭代生成过程导致采样速度缓慢并限制其实时应用的潜力。

为了克服这一局限性,我们提出了一致性模型,这是一个新的生成模型家族,可以在没有对抗性训练的情况下实现高样本质量。

它们通过设计支持快速一步生成,同时仍然允许几步采样来权衡计算和样本质量。

他们还支持零样本编辑,如图像绘制、着色和超分辨率,而无需对这些任务进行明确训练。

一致性模型可以作为蒸馏预训练的扩散模型的一种方式进行训练,也可以作为独立的生成模型进行训练。

通过广泛的实验,我们证明了它们在一步和几步的生成中优于现有的扩散模型蒸馏技术。

例如,我们在CIFAR-10上实现了3.55,在ImageNet 64 x 64上实现了 6.20 的新最先进的 FID,用于一步生成。

当作为独立的生成模型进行训练时,一致性模型在CIFAR-10、ImageNet 64 x 64和LSUN 256 x 256等标准基准上也优于单步、非对抗性生成模型。

Introduction

扩散模型(Sohl-Dickstein等人,2015年;Song & Ermon,2019年;2020年;Ho等人,2020年;Song等人,2021年),也被称为基于分数的生成模型,在多个领域取得了前所未有的成功,包括图像生成(Dhariwal和Nichol,2021年;Nichol等人,2021年;Ramesh等人),音频生成((Kong et al., 2020; Chen et al., 2021; Popov et al., 2021)和视频生成 (Ho et al., 2022b;a)。与生成对抗网络(GANs,Goodfellow等人(2014))不同,这些模型不依赖于对抗性训练,因此不太容易出现训练不稳定和模式崩溃等问题。此外,扩散模型对模型架构施加的严格约束与自回归模型(Bengio和Bengio,1999年;Uria等人,2013年;2016年;Van Den Oord等人,2016年)、变分自编码器(VAEs,Kingma和Welling(2014年);Rezende等人(2014年))或归一化流(Dinh et al., 2015; 2017; Kingma & Dhariwal, 2018)。

扩散模型成功的关键是其迭代采样过程,该过程逐渐从随机噪声矢量中去除噪声。这种迭代细化过程重复评估扩散模型,允许对样本质量和计算量进行权衡:通过使用额外的计算进行更多的迭代,小尺寸的模型可以展开到更大的计算图中,生成更高质量的样本。迭代生成对于扩散模型的零样本编辑能力也至关重要,使他们能够解决具有挑战性的逆问题,从图像绘制、着色、笔画引导图像编辑到计算机断层扫描和磁共振成像(Song & Ermon,2019年;Song等人,2021年;2022年;2023年;Kawar等人,2021年;2022年; Chung et al., 2023; Meng et al., 2021)。然而,与GAN、VAE和规范化流等单步生成模型相比,扩散模型的迭代生成过程通常需要10-2000倍的计算(Song & Ermon,2020年;Ho等人,2020年;Song等人,2021年;Zhang和陈,2022年;Lu等人,2022年),导致推理缓慢并限制了它们的实时应用。

我们的目标是创建生成模型,在不牺牲迭代改进的重要优势的情况下,促进高效的单步生成。这些优势包括在必要时对样本质量和计算进行权衡,以及执行零样本编辑任务的能力。如图1所示,我们在连续时间扩散模型(Song等人,2021年)中以概率流(PF)常微分方程(ODE)为基础,其轨迹将数据分布平稳地过渡到可处理的噪声分布。我们提出学习一个模型,在任何时候将任何点映射到轨迹的起点。我们模型的一个显著属性是自一致性:同一轨迹图上的点到同一初始点。因此,我们将此类模型称为一致性模型。一致性模型允许我们通过转换随机噪声矢量(ODE轨迹的端点,例如图1中的$x_T$)来生成数据样本(ODE轨迹的初始点,例如图1中的 $x_0$)。重要的是,通过在多个时间步骤中链接一致性模型的输出,我们可以提高样本质量,并以更多计算的成本执行零样本编辑,类似于迭代细化对扩散模型的影响。

为了训练一致性模型,我们提供了两种基于强制执行自一致性属性的方法。第一种方法依赖于使用数值ODE求解器和预训练的扩散模型来生成PF ODE轨迹上的相邻点对。通过最大限度地减少这些对的模型输出之间的差异,我们可以有效地将扩散模型蒸馏成一致性模型,从而允许通过一次网络评估生成高质量的样本。相比之下,我们的第二种方法完全消除了对预训练的扩散模型的需求,允许我们独立地训练一致性模型。这种方法将一致性模型定位为一个独立的生成模型家族。至关重要的是,这两种方法都不需要对抗性训练,而且这两种训练方法都允许灵活的神经网络架构来实现一致性模型。

我们在几个具有挑战性的图像基准上展示了一致性模型的有效性,包括CIFAR-10(Krizhevsky等人,2009年)、ImageNet 64 x 64(Deng等人,2009年)和LSUN 256 x 256(Yu等人,2015年)。经验上,我们观察到,作为一种蒸馏方法,一致性模型在各种数据集和采样步骤中优于渐进式蒸馏(Salimans & Ho,2022年)。在CIFAR-10上,一致性模型在一步和两步生成中达到了3.55和2.93的新的最先进的FID。在ImageNet 64 x 64上,它分别通过一次和两次网络评估,实现了创纪录的 6.20 和 4.70 的FID。当作为独立的生成模型进行训练时,一致性模型实现了与单步生成的渐进蒸馏相当的性能,尽管无法使用预训练的扩散模型。他们能够超越许多GAN,以及跨多个数据集的所有其他非对抗性单步生成模型。我们还表明,一致性模型可用于执行零样本编辑任务,包括图像去噪、插值、内绘、着色、超分辨率和笔画引导图像编辑(SDEdit、Meng等人)。(2021))。