全文翻译
摘要
扩散模型在多个领域都取得了显著成功。然而，其生成速度缓慢仍是一个关键挑战。现有的加速方法虽以减少步骤为目标，但往往会牺牲样本质量、可控性或增加训练复杂性。为此，我们提出了RayFlow，这是一种新颖的扩散框架，旨在解决这些局限性。与以往方法不同，RayFlow引导每个样本沿着独特路径朝着特定于实例的目标分布演进。该方法在最大限度减少采样步骤的同时，保留了生成的多样性和稳定性。此外，我们引入了时间采样器（Time Sampler），这是一种重要性采样技术，通过聚焦关键时间步来提高训练效率。大量实验表明，与现有的加速技术相比，RayFlow在生成高质量图像方面具有优势，同时提升了速度、可控性和训练效率。
1. 引言
扩散模型彻底改变了生成式人工智能领域[10, 41-43]，在多个领域取得了令人瞩目的成果，从文本[15, 47]、图像[40]到3D模型[4, 22, 46]、音频[9, 19]以及修复[1, 2]等领域均有涉猎。然而，其生成速度缓慢（通常每个样本需要数十步）仍是一个显著的局限性。各种蒸馏技术试图加速这一过程，包括常规蒸馏[26, 56]、对抗性蒸馏[38, 45]、渐进式蒸馏[36]和变分分数蒸馏[29, 37, 46, 49, 53, 54]。当前的蒸馏方法尽管各有目标，但仍面临诸多挑战，包括大量的计算开销、复杂的训练方案，以及在生成速度、样本质量和有效引导方面的局限性[12, 28, 30, 31, 48, 55, 58]。这凸显了开发更具通用性的方法以充分发挥扩散模型能力的必要性。
如图1a所示，传统的扩散过程是大多数加速和蒸馏方法的基础。然而，从图中可以明显看出，该过程存在几个问题：
- （1）期望差异：反向过程中的期望在不同时间步存在差异。要获得高质量样本，需要更多的采样步骤，这使得加速采样方法无法避免质量下降。
- （2）扩散路径重叠：由于所有样本最终都会收敛到相同的标准高斯分布，扩散概率路径会发生重叠。$\mathbb{E}[\bar{\epsilon}_{t} | x_{t}]$可能代表多个概率路径的交点，导致采样结果存在显著随机性，并可能造成严重的质量损失。
- （3）采样不稳定性：即使采样点位置接近，最终生成的结果也可能存在显著差异，从而引入严重的不稳定性。
已有一些研究试图解决这些局限性：文献[23]引入了整流流（RF），使用线性常微分方程（ODE）进行直线路径采样（图1b左上角），而文献[44]将RF扩展到一阶曲线路径整流扩散（RD），用于类去噪扩散概率模型（DDPM）（图1b右上角）。这些样本-噪声匹配方法也被文献[25, 51]应用于扩散蒸馏。
然而，现有的样本-噪声匹配算法在试图解决上述挑战时带来了新的缺点：（1）路径不一致：样本-噪声匹配与实际ODE采样路径之间的差距过大，可能导致训练困难和泛化能力差；（2）多样性有限：采样概率路径受到严重限制，显著降低了模型生成的多样性；（3*）理论缺口：该方法虽直观性强，但缺乏基本的理论推导来证明其在采样稳定性方面的最优性。
为了解决传统扩散中的这三个挑战以及现有样本-噪声匹配方法带来的缺点，我们提出了RayFlow（如图1b底部所示）。一致的期望和路径：针对挑战（1）和（1），我们利用预训练模型计算所有时间步的统一噪声期望$\epsilon_{\mu}=\mathbb{E}_{t}[\mathbb{E}[\bar{\epsilon}_{t}]]$，从而实现高效的步骤压缩而不降低质量。个体路径设计：针对挑战（2）和（2），每个样本不再收敛到共同的高斯分布，而是沿着独特的扩散路径朝着其特定的目标均值演进，同时降低方差，最大限度地减少路径重叠和采样随机性。理论保证：针对挑战（3）和（3*），我们证明了我们的方法最大化了起点、目标均值和原点之间的路径概率，确保了最优的采样稳定性和原始数据点的可靠重建。
此外，为了提高训练效率，我们开发了时间采样器（Time Sampler），这是一种先进的重要性采样技术，用于在训练过程中识别关键时间步。通过将随机斯坦差异（SSD）与神经网络相结合，时间采样器近似最优采样时间步分布，以最小化训练损失估计器的方差，从而减少计算冗余并提高效率。
我们的主要贡献总结如下：
- RayFlow框架：我们引入了一种创新的扩散框架，具有实例独立的目标均值。这种方法增强了对生成过程的控制，实现了更高效、更精确的采样。
- 时间采样器：我们开发了一种利用SSD的时间步重要性采样技术。该方法能有效识别训练过程中的关键时间步，减少计算冗余并提高效率。
- 高效算法：我们提出了用于RayFlow训练和采样的实用算法，包括用于更快生成的快速单步采样变体。
- 理论分析：我们对RayFlow进行了全面的理论检验，详细推导了路径概率并优化了参数，以确保最大的采样稳定性。
通过大量实验，我们证明了RayFlow在生成高质量图像方面的有效性，与现有的加速算法相比，在速度、控制和训练效率方面都有提升。我们的工作为探索和控制扩散过程开辟了新的途径。