Meta最新成果：前所未有的加速Emu推理 | Imagine Flash：新型蒸馏框架发布

文章地址：https://arxiv.org/pdf/2405.05224

扩散模型是一个强大的生成框架，但推理过程非常昂贵。现有的加速方法通常会在极低步骤情况下牺牲图像质量，或者在复杂条件下失败。这项工作提出了一个新颖的蒸馏框架，旨在通过仅使用一到三个步骤实现高保真度和多样化样本生成。方法包括三个关键组成部分：

(i) 反向蒸馏，通过校准学生模型自身的反向轨迹来缓解训练推理差异；
(ii) 迁移重构损失，根据当前时间步动态调整知识传递；
(iii) 噪声校正，一种增强样本质量的推理时技术，通过解决噪声预测中的奇异性来改善样本质量。

通过大量实验证明，我们的方法在定量指标和人类评估方面优于现有竞争对手。值得注意的是，我们的方法仅使用三个去噪步骤就实现了与教师模型相当的性能，实现了高效的高质量生成。

生成的效果图

介绍

随着去噪扩散模型（DMs）的出现，生成建模已经发生了一场范式转变。这些模型在各个领域树立了新的基准，提供了前所未有的逼真性和多样性，同时确保了稳定的训练。然而，去噪过程的顺序性质带来了显著挑战。从DMs中抽样是一个耗时和昂贵的过程，所需时间在很大程度上取决于两个因素：（i）每步神经网络评估的延迟，以及 （ii）去噪步骤的总数。

大量的研究工作致力于加速抽样过程。对于文本到图像的合成，提出的方法涵盖了广泛的技术，包括高阶求解器、用于曲率降低的修改扩散公式，以及引导、步骤和一致性蒸馏。这些方法取得了令人印象深刻的改进，在接近10步的区间内达到了非常高的质量。最近，利用蒸馏和对抗性损失两者的混合方法[17,31,39]将界限推动到了不到五步。尽管这些方法在简单提示和简单样式（如动画）上取得了令人印象深刻的质量，但它们在真实照片般的图像上，特别是在复杂的文本条件下，样本质量下降。

上述方法中的一个共同主题是试图将少步骤的学生模型与复杂的教师路径对齐，尽管学生模型具有明显较低的容量（即步骤）。我们认识到这是一个限制，因此提出了一种新颖的蒸馏框架，旨在使教师模型沿着自己的扩散路径改进学生模型。总之，本文贡献有三个方面：

首先，方法引入了反向蒸馏，这是一种设计用于在其自身上游反向轨迹上校准学生模型的蒸馏过程，从而减少了训练和推理分布之间的差距，并确保在所有时间步长上训练期间没有数据泄漏。
其次，提出了一种迁移重构损失，它动态地调整来自教师模型的知识迁移。具体而言，该损失被设计为在高时间步长时从教师模型中提炼全局结构信息，而在低时间步长时专注于渲染细粒度细节和高频成分。这种自适应方法使学生能够有效地模拟教师的生成过程在扩散轨迹的不同阶段。
最后，提出了噪声校正，这是一种在推理时修改，通过解决在初始采样步骤中存在的噪声预测模型中的奇异性来增强样本质量的技术。这种无需训练的技术减轻了对比度和颜色强度的降级，这种降级通常会在极低数量的去噪步骤中操作时出现。

通过将这三个新颖组件协同结合起来，将我们的蒸馏框架应用到基线扩散模型Emu中，从而实现了Imagine Flash，在极低步长区间内实现了高质量的生成，而不会影响样本质量或条件保真度（下图2）。通过大量实验和人类评估，展示了我们的方法在各种任务和模态之间实现了抽样效率和生成质量之间的有利折衷。

方法

Imagine Flash，这是一种新颖的蒸馏技术，旨在快速进行文本到图像的生成，它建立在Emu 的基础上，但不局限于此。与原始的Emu模型不同，后者需要至少50个神经函数评估（NFEs）才能生成高质量的样本，Imagine Flash仅需少量NFEs即可实现可比较的结果。提出的蒸馏方法包括三个新颖的关键组成部分：

反向蒸馏，这是一种蒸馏过程，确保了在所有时间点t上的训练期间零数据泄漏。
迁移重构损失（SRL），这是一种自适应损失函数，旨在最大化从教师模型中的知识迁移。
噪声校正，这是一种无需训练的推理修改，用于改善在噪声预测模式下训练的少步骤方法的样本质量。

接下来，假设可以访问一个预训练的扩散模型，它预测噪声估计。这个教师模型可以在图像空间或潜在空间中操作。我们的目标是将的知识蒸馏到一个学生模型中，同时减少抽样步骤的总数，并提供额外允许中每一步的高质量增加。如果模型使用无分类器的引导(cfg)，那么也将这些知识蒸馏到学生模型中，并且消除cfg的需要。

反向蒸馏

总的来说，反向蒸馏消除了所有时间步骤t上的信息泄漏，防止模型依赖于真值信号。这是通过在训练期间模拟推理过程来实现的，这也可以解释为在其自身的上游反向路径上校准学生模型。

SRL：迁移重构损失

通过反向扩散进行图像生成的过程中，早期阶段（其中t接近T）对于构建图像的整体结构和组成至关重要。相反，后期阶段（其中t接近0）对于添加高级细节至关重要。借鉴这一观察结果，对默认的知识蒸馏损失进行了改进，鼓励学生模型学习教师模型的结构组成和细节渲染能力。这涉及将教师去噪的起始点从学生的起始点t移开，因此我们将这种方法称为迁移重构损失（SRL）。下图4概述了提出的Loss。

这种方法鼓励学生在早期后向步骤中优先提取结构知识，并在最后后向步骤中集中精力生成更精细的细节。

噪声校正

我们提出了一个简单的、无需训练的替代方案，即切换到zero-SNR速度预测，这允许使用噪声预测模型而不产生前述的偏差。具体来说，通过将视为一个特殊情况，并用真实噪声替换，更新得到了校正。

观察到，这个小的修改可以显著改善估计的颜色，使色调更加生动饱满。当推理步骤数量较少时，这种效应尤为明显。

实验

为了确保公平性，将Emu模型作为所有实验的基础。Emu是一个拥有27亿参数和分辨率为768×768的最先进模型。将我们的结果与之前的蒸馏方法进行比较，例如Step Distillation、LCM和ADD，并将它们直接应用在Emu上。所有模型都在一个委托的图像数据集上进行了替换训练。由于ADD训练没有公开可用的代码，根据论文中提供的细节自行实现了它。

与最新技术的定量比较

使用FID 、CLIP分数和CompBench将Imagine Flash与之前的方法进行比较。FID和CLIP分别衡量图像质量和提示对齐度，并根据来自COCO2017 的5k样本分割进行评估，遵循[31]中的评估协议。CompBench是一个基准，单独衡量属性绑定（颜色、形状和纹理）和对象关系（空间、非空间和复杂）。

在CompBench验证集中为每个提示生成2张图像（总共300个提示）。对于LCM和Imagine Flash，计算1、2和3步的指标。对于ADD，计算4步的指标，因为该方法专门针对4步推理进行了调优和配置，以确保公平比较。还评估了4步的Step Distillation，以提供更直接的比较。表1显示了结果。

我们的3步Imagine Flash在FID方面优于Step Distillation和ADD，即使使用了少一步。它在1、2和3步的FID上也比LCM低。我们的3步模型的CLIP分数高于所有变体的ADD和LCM，并与4步Step Distillation模型的得分（30.2）相匹配。与Step Distillation和ADD相比，它们相应地降低了FID分别为10.1和3.4，而我们的3步和2步Imagine Flash则保持了FID并略有改善。

对于CompBench，我们的1、2或3步Imagine Flash在所有类别中都优于以前的方法，除了颜色外，其中4步Step Distillation和ADD的得分与我们的类似。这突显了Imagine Flash的优越提示对齐性。

与最新技术的定性比较

在图5中，展示了Imagine Flash与当前最先进技术（SOTA）的定性比较：Step Distillation、LCM和ADD，它们都对相同的基线Emu模型进行了蒸馏，以进行公平比较。

观察到ADD生成的图像比Step Distillation和LCM更清晰，这是因为它使用了对抗损失。虽然Imagine Flash和ADD都使用了鉴别器，但Imagine Flash生成的图像比ADD更锐利和更详细。Imagine Flash的增强锐度和细节是由我们提出的SRL导致的，它有效地精炼了学生预测的高频细节，如下图6的最后一行所示。

另一方面，对于ADD，目标图像可能展示出明显不同的色谱，出现颜色伪影（见图6），并且颜色在训练迭代过程中可能出现不可预测的波动。我们假设，为了在期望中最小化L2重构损失，ADD模型最好通过预测接近零的颜色值来达到目标，从而导致图像苍白和轮廓模糊。除了改善局部细节外，SRL还可以纠正学生的文本对齐错误，如图6右侧所示（1步），其中小熊猫被转换回了狗。

与公开模型的比较

还将Imagine Flash的性能与ADD-LDMXL和Lightning-LDMXL发布的公开模型进行了比较。为此，按照前面详细说明的方法计算CLIP和FID分数，并与基准模型进行相对增益/下降的比较。请在附录A中找到表格。我们的方法在文本对齐能力上与ADD和Lightning保持类似，但在FID增加方面显示出更为有利的结果，尤其是对于两步和三步。

此外，进行了广泛的人类评估。为此，使用了三个推理步骤为OUI数据集中随机抽样的1,000个提示生成了所有方法的图像。成对的图像呈现给42名受过训练的人类标注者中的五名，他们的任务是为更具视觉吸引力的图像投票。通过多数投票汇总的结果显示在表2中，明显表明了对Imagine Flash的偏好。