arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2408.15991cs.CV

分布回溯为扩散蒸馏构建更快的收敛轨迹

Distribution Backtracking Builds A Faster Convergence Trajectory for Diffusion Distillation

College of Computer Science and Technology, Zhejiang University(浙江大学计算机科学与技术学院);School of Software Technology, Zhejiang University(浙江大学软件学院);Peking University(北京大学);Alibaba Group(阿里巴巴集团)

Shengyuan Zhang, Ling Yang, Zejian Li, An Zhao, Chenye Meng, Changyuan Yang, Guang Yang, Zhiyuan Yang, Lingyun Sun

AI总结:

针对扩散模型采样加速中分数蒸馏早期的分数不匹配问题,提出分布回溯蒸馏(DisBack),引入教师模型完整收敛轨迹,通过退化记录和分布回溯两阶段优化,实现更快收敛与相当生成性能,可推广至现有方法。

摘要:

加速扩散模型的采样速度仍是一项重大挑战。近期的分数蒸馏方法将大型教师模型蒸馏为学生生成器以实现单步生成,该过程通过计算学生模型生成样本上两个分数函数的差异进行优化。然而,蒸馏过程早期存在分数不匹配问题,因为现有方法主要聚焦于使用预训练扩散模型的端点作为教师模型,忽略了学生生成器与教师模型之间收敛轨迹的重要性。为解决这一问题,我们通过引入教师模型的完整收敛轨迹扩展分数蒸馏过程,并提出分布回溯蒸馏(DisBack)。DisBack由退化记录和分布回溯两个阶段组成:退化记录阶段旨在获取教师模型的收敛轨迹,记录从训练好的教师模型到未训练的初始学生生成器的退化路径,该退化路径隐式表示教师模型的中间分布,其反向可视为从学生生成器到教师模型的收敛轨迹;分布回溯阶段训练学生生成器沿该路径回溯中间分布,以逼近教师模型的收敛轨迹。大量实验表明,DisBack比现有蒸馏方法收敛更快、效果更好,且生成性能相当,在ImageNet 64x64数据集上的FID得分为1.38。值得注意的是,DisBack易于实现,可推广到现有蒸馏方法以提升性能。我们的代码公开于https://github.com/SYZhang0805/DisBack。

英文摘要:

Accelerating the sampling speed of diffusion models remains a significant challenge. Recent score distillation methods distill a heavy teacher model into a student generator to achieve one-step generation, which is optimized by calculating the difference between the two score functions on the samples generated by the student model. However, there is a score mismatch issue in the early stage of the distillation process, because existing methods mainly focus on using the endpoint of pre-trained diffusion models as teacher models, overlooking the importance of the convergence trajectory between the student generator and the teacher model. To address this issue, we extend the score distillation process by introducing the entire convergence trajectory of teacher models and propose Distribution Backtracking Distillation (DisBack). DisBask is composed of two stages: Degradation Recording and Distribution Backtracking. Degradation Recording is designed to obtain the convergence trajectory of the teacher model, which records the degradation path from the trained teacher model to the untrained initial student generator. The degradation path implicitly represents the teacher model's intermediate distributions, and its reverse can be viewed as the convergence trajectory from the student generator to the teacher model. Then Distribution Backtracking trains a student generator to backtrack the intermediate distributions along the path to approximate the convergence trajectory of teacher models. Extensive experiments show that DisBack achieves faster and better convergence than the existing distillation method and accomplishes comparable generation performance, with FID score of 1.38 on ImageNet 64x64 dataset. Notably, DisBack is easy to implement and can be generalized to existing distillation methods to boost performance. Our code is publicly available on https://github.com/SYZhang0805/DisBack.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑