分佈式回溯為一步擴散蒸餾建立了更快的收斂軌跡。

Distribution Backtracking Builds A Faster Convergence Trajectory for One-step Diffusion Distillation

August 28, 2024

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

cs.AI

摘要

加速擴散模型的取樣速度仍然是一個重要挑戰。最近的分數蒸餾方法將一個龐大的教師模型蒸餾成一個一步生成器學生模型，通過計算兩個分數函數在學生模型生成的樣本上的差異來進行優化。然而，在蒸餾過程的早期階段存在分數不匹配問題，因為現有方法主要集中於使用預先訓練的擴散模型的端點作為教師模型，忽略了學生生成器與教師模型之間的收斂軌跡的重要性。為了解決這個問題，我們通過引入教師模型的整個收斂軌跡擴展了分數蒸餾過程，並提出了分布回溯蒸餾（DisBack）用於蒸餾學生生成器。DisBack由兩個階段組成：退化記錄和分布回溯。退化記錄旨在獲取教師模型的收斂軌跡，記錄了從訓練有素的教師模型到未訓練的初始學生生成器的退化路徑。這個退化路徑隱含地代表了教師模型的中間分布。然後，分布回溯訓練一個學生生成器來回溯中間分布，以逼近教師模型的收斂軌跡。大量實驗表明，DisBack實現了比現有蒸餾方法更快更好的收斂，並實現了可比的生成性能。值得注意的是，DisBack易於實現，並且可以應用於現有蒸餾方法以提高性能。我們的代碼公開在https://github.com/SYZhang0805/DisBack。

English

Accelerating the sampling speed of diffusion models remains a significant challenge. Recent score distillation methods distill a heavy teacher model into an one-step student generator, 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) for distilling student generators. DisBask is composed of two stages: Degradation Recording and Distribution Backtracking. Degradation Recording is designed to obtain the convergence trajectory of teacher models, which records the degradation path from the trained teacher model to the untrained initial student generator. The degradation path implicitly represents the intermediate distributions of teacher models. Then Distribution Backtracking trains a student generator to backtrack the intermediate distributions for approximating 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. 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.