La distribution du retour en arrière crée une trajectoire de convergence plus rapide pour la distillation de diffusion en une étape.

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

August 28, 2024

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

cs.AI

Résumé

Accélérer la vitesse d'échantillonnage des modèles de diffusion reste un défi significatif. Les récents méthodes de distillation de score distillent un modèle enseignant lourd en un générateur étudiant à une étape, qui est optimisé en calculant la différence entre les deux fonctions de score sur les échantillons générés par le modèle étudiant. Cependant, il existe un problème de désaccord de score au début du processus de distillation, car les méthodes existantes se concentrent principalement sur l'utilisation du point final des modèles de diffusion pré-entraînés en tant que modèles enseignants, en négligeant l'importance de la trajectoire de convergence entre le générateur étudiant et le modèle enseignant. Pour résoudre ce problème, nous étendons le processus de distillation de score en introduisant l'ensemble de la trajectoire de convergence des modèles enseignants et proposons la Distillation avec Retour sur Distribution (DisBack) pour distiller les générateurs étudiants. DisBack se compose de deux étapes : Enregistrement de la Détérioration et Retour sur Distribution. L'Enregistrement de la Détérioration est conçu pour obtenir la trajectoire de convergence des modèles enseignants, qui enregistre le chemin de détérioration du modèle enseignant entraîné au générateur étudiant initial non entraîné. Le chemin de détérioration représente implicitement les distributions intermédiaires des modèles enseignants. Ensuite, le Retour sur Distribution entraîne un générateur étudiant à revenir sur les distributions intermédiaires pour approximer la trajectoire de convergence des modèles enseignants. Des expériences approfondies montrent que DisBack atteint une convergence plus rapide et meilleure que la méthode de distillation existante et réalise des performances de génération comparables. Notamment, DisBack est facile à implémenter et peut être généralisé aux méthodes de distillation existantes pour améliorer les performances. Notre code est publiquement disponible sur 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.