扩散模型已成为现代生成建模的基石,在众多领域取得了先进成果(Ho, Jain, Abbeel, 2020; Karras, Aittala, Aila, Laine, 2022; Nichol, Dhariwal, 2021; Ren, Zhang, Li, Jiang, Qin, Li, Li, Zheng, Li, Zhan, & Peng, Song, Meng, Ermon, 2021a; Song, Ermon, 2019; Song, Sohl-Dickstein, Kingma, Kumar, Ermon, Poole, 2021b; Yu, Li, Zhang, & Gao)。它们的成功源于一个多步去噪过程,该过程逐步将高斯噪声细化为结构化数据样本。虽然这种迭代细化能够实现高样本保真度,但也带来了一个关键瓶颈:生成单个输出通常需要数百甚至数千步,使得扩散模型不适用于实时应用(Ho, Jain, Abbeel, 2020; Song, Ermon, 2019)。
为了解决这一限制,最近的研究转向了扩散蒸馏(Berthelot, Autef, Lin, Yap, Zhai, Hu, Zheng, Talbott, & Gu, Gu, Zhai, Zhang, Liu, Susskind, 2023; Li, Feng, Chen, Wang, 2024a; Lin, Wang, & Yang, Luhman, & Luhman, Luo, Hu, Zhang, Sun, Li, Zhang, 2023; Meng, Rombach, Gao, Kingma, Ermon, Ho, Salimans, 2023; Salimans, Ho, 2022; Sauer, Lorenz, Blattmann, Rombach, 2024; Song, Dhariwal, Chen, Sutskever, 2023; Xu, Nie, & Vahdat, Yin, Gharbi, Park, Zhang, Shechtman, Durand, Freeman, 2024a; Yin, Gharbi, Zhang, Shechtman, Durand, Freeman, Park, 2024b; Zhou, Zheng, Gu, Wang, Huang, 2025; Zhou, Zheng, Wang, Yin, Huang, 2024),该方法训练学生模型仅用少量步骤就能模仿完整扩散轨迹(即教师模型)的行为。这些方法,包括轨迹蒸馏、分布匹配蒸馏和对抗性蒸馏,在实证研究中取得了显著成功,有些方法甚至能在一步内实现高质量合成(Luhman, & Luhman, Sauer, Lorenz, Blattmann, Rombach, 2024; Yin, Gharbi, Zhang, Shechtman, Durand, Freeman, Park, 2024b)。尽管这些方法被广泛采用,但其背后的机制仍不明确。特别是,合并顺序、轨迹长度或学生模型初始化等设计选择如何影响生成质量和效率尚不清楚。这种理论上的缺乏使得难以预测这些因素的影响,通常迫使从业者依赖试错或启发式方法,而这些方法可能不适用于不同的数据集或噪声模式。更重要的是,没有明确的指导原则,系统改进蒸馏策略变得具有挑战性。
在这项工作中,我们对扩散轨迹蒸馏(Berthelot, Autef, Lin, Yap, Zhai, Hu, Zheng, Talbott, & Gu, Gu, Zhai, Zhang, Liu, Susskind, 2023; Li, Feng, Chen, Wang, 2024a; Lin, Wang, & Yang, Luhman, & Luhman, Meng, Rombach, Gao, Kingma, Ermon, Ho, Salimans, 2023; Salimans, Ho, 2022; Song, Dhariwal, Chen, Sutskever, 2023)进行了详细分析,这类方法直接将教师模型的采样轨迹压缩为更少的学生步骤。这种方法避免了辅助评分网络(如分布匹配蒸馏(Luo, Hu, Zhang, Sun, Li, Zhang, 2023; Xu, Nie, & Vahdat, Yin, Gharbi, Park, Zhang, Shechtman, Durand, Freeman, 2024a; Yin, Gharbi, Zhang, Shechtman, Durand, Freeman, Park, 2024b; Zhou, Zheng, Wang, Yin, Huang, 2024)或判别器(如对抗性蒸馏(Lin, Wang, & Yang, Sauer, Lorenz, Blattmann, Rombach, 2024; Yin, Gharbi, Park, Zhang, Shechtman, Durand, Freeman, 2024a; Zhou, Zheng, Gu, Wang, Huang, 2025)所带来的复杂性,同时抓住了众多实际蒸馏技术背后的核心思想。
为了更好地理解轨迹蒸馏,我们研究了一个简化的场景,以便进行清晰的理论分析。正如Li等人(2024b)所展示的,扩散模型在容量有限的情况下倾向于学习高斯结构。因此,我们从单峰高斯情况开始,其中教师模型和学生模型都可以表示为应用于噪声观测的线性算子。在这种情况下,轨迹蒸馏简化为一个算子合并问题,即多步教师轨迹被压缩为更少的学生步骤。我们展示了每次合并都会通过离散化和学生模型的有限优化时间导致信号收缩,明确了合并顺序如何影响信号保留。
在此基础上,我们研究了当目标是匹配教师模型的输出分布时,哪些合并计划是最优的。使用高斯输出的2-Wasserstein距离,我们证明了最优策略随着数据方差的增加而发生明显的相变(Sectin 5.1),确定了标准基线方法(如vanilla蒸馏(Luhman and Luhman, 2021)和顺序BOOT(Gu et al., 2023)在何时是最优的。在这些理论指导下,我们提出了一种帕累托动态规划算法,该算法可以计算全局最优的合并计划,同时高效地剔除较差的中间候选方案,这在数据维度不同时尤为重要。
为了解决多模性问题,我们将分析扩展到高斯混合数据。在这种情况下,每个教师算子可以表示为一个K组分的专业家混合映射,组合< />步相当于扩展到Kk个有效组分。我们描述了当K组分的学生算子蒸馏这样的复合教师时不可避免的近似误差,并展示了这些误差在合并过程中的累积情况,解释了为什么对于多模数据来说,长视野合并变得更加困难。我们在表1中提供了线性和非线性情况下的理论框架的比较总结。
本文的主要贡献如下:
•单峰高斯情况下的算子合并基础和信号收缩。在线性高斯情况下,我们推导出了封闭形式的教师算子,并将轨迹蒸馏表述为一个算子合并问题(Sectin 4.1)。这一视角通过描述有限学生优化时间下的信号收缩机制来隔离具体的退化过程(Sectin 4.2)。
•基于方差的最优合并计划相变。使用高斯输出分布之间的2-Wasserstein距离作为目标(Sectin 4.3),我们证明了最优合并策略随着数据方差λ的变化而发生明显的相变(Sectin 5.1)。当λ?≤?1时,恢复顺序BOOT;当λ???1时,恢复vanilla轨迹蒸馏,为每种基线的适用性提供了原理性解释。
•用于最优合并规划的帕累托动态规划算法。受最优性特征的启发,我们提出了一种帕累托动态规划算法,该算法可以在高维且方差混合的情况下计算最优合并计划,同时高效地剔除较差的中间候选方案(Sectin 5.2)。
•高斯混合模型的扩展与近似和误差传播界限。
我们将分析扩展到高斯混合模型,将每个教师步骤建模为一个K组分的专家混合算子,并展示了< />步组合扩展为Kk个有效组分(Sectin 6.1)。在这种情况下,我们量化了K组分学生算子在蒸馏复合教师时的内在近似误差(Sectin 6.2),并描述了这些误差在合并过程中的累积情况(Sectin 6.3)。