Метод обратного распространения распределения обеспечивает более быструю сходимость траектории для дистилляции диффузии на один шаг.

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.