小红书推出AIGC加速新算法TDD

扩散模型是当前图像生成领域的热门技术之一，但其生成过程往往需要大量的迭代计算，耗时较长。为了解决这一问题，近年来出现了多种一致性蒸馏方法，旨在加速图像生成。然而，传统方法通常在图像细节表现上有所欠缺。我们提出了一种全新的多目标蒸馏方法——目标驱动蒸馏（TDD），它通过灵活选择生成步骤并解耦训练过程中的指导信号，大大提升了图像生成速度与质量。我们还将探讨TDD如何在多个实验中表现出色，证明其在图像生成任务中的广泛应用潜力。

一致性蒸馏方法在加速扩散模型的生成任务方面取得了显著成功。然而，由于之前的一致性蒸馏方法在选择目标时间步长时采用了简单直接的策略，生成的图像往往存在模糊和细节丢失的问题。

扩散模型在图像生成方面展示了卓越的性能，能够生成高质量且多样化的图像。然而，由于扩散模型采用迭代去噪过程，生成图像通常需要较长的时间。

一致性蒸馏提出了一种有效的方法，能够在加速生成同时保持图像质量。这些方法通过遵循自一致性原则来蒸馏预训练的扩散模型，即两个相邻时间步的预测结果向同一目标时间步的结果被正规化为相同。根据目标时间步的选择，我们将一致性蒸馏方法分为单目标蒸馏和多目标蒸馏。

单目标蒸馏方法在选择目标时间步时遵循一对一映射，即每次沿着PF-ODE的轨迹到达某个时间步时，始终选择相同的目标时间步。一个简单直接的选择是将任何时间步映射到最终的0时间步。然而，这些方法通常受到远距离预测的累积误差影响。另一种选择是将整个轨迹均匀划分为几个子轨迹，并将时间步映射到其所属子轨迹的结束点。虽然通过缩短训练中的预测距离可以减少误差，但在推理阶段采用不同数量的子轨迹时，图像质量将会不尽理想。

另一方面，多目标蒸馏方法遵循一对多映射，即每次到达某个时间步时，可能选择不同的目标时间步。一个典型的选择是将当前时间步映射到前面的随机目标时间步。理论上，这些方法被训练为从任何时间步预测到任何时间步，因此在不同的时间表下通常可以实现良好的性能。然而，实际上大多数这些预测是多余的，因为在常见的去噪时间表下我们不会经过这些时间步。因此，多目标蒸馏方法通常需要较高的时间预算来训练。

为缓解上述问题，我们提出了目标驱动蒸馏（TDD），这是一种强调在蒸馏过程中精心选择目标时间步的多目标方法。我们的方法包含三个关键

(1) 首先，TDD采用精细的目标时间步长选择策略。对于任何时间步，它选择前方的一个相邻时间步，该时间步落入预定义的时间表（例如4到8步，Kmin=4，Kmax=8）的少步等距去噪时间表中，这消除了远距离预测，同时只关注推理阶段可能经过的时间步。

此外，TDD引入了一个随机偏移量eta，进一步推动所选时间步向最终目标时间步前进，以适应诸如γ采样等非确定性采样。

(2) 其次，TDD在训练过程中使用解耦的指导方式。在蒸馏无分类器指导（CFG）到蒸馏模型时，为了与使用CFG的标准训练过程保持一致，TDD还用无条件（即空白）提示替换了一部分文本条件。通过这种设计，TDD可以在推理时使用一种提出的指导尺度调整技术，允许用户在基于文本提示生成的图像内容的准确性和丰富性之间找到平衡。

同时基于这种解耦方式，我们提出Guidance Scale Tuning的推理方式，用以缓解当前大CFG蒸馏过程中造成的图像复杂度低的问题。

(3) 最后，TDD可选地配备了一种非等距采样方法。在初始步骤进行短距离预测，而在后期步骤进行长距离预测，从而提高整体图像质量。此外，TDD采用x0裁剪以防止超出边界的预测，并解决过度曝光问题。

• 王存正

  小红书AIGC实习生，浙江大学软件学院在读硕士，主要研究AIGC图像生成以及AIGC加速生成方向。

• 郭紫垣

  小红书AIGC研发工程师，主要负责小红书智能创作AIGC特效生成/AIGC商业化探索方向。

• 段宇轩

  小红书AIGC实习生，上海交通大学计算机系在读博士，主要研究图像/视频等多模态生成方向。