Apprentissage de modèles de diffusion en quelques étapes par alignement des distributions de trajectoires

Learning Few-Step Diffusion Models by Trajectory Distribution Matching

March 9, 2025

Auteurs: Yihong Luo, Tianyang Hu, Jiacheng Sun, Yujun Cai, Jing Tang

cs.AI

Résumé

L'accélération de l'échantillonnage des modèles de diffusion est cruciale pour un déploiement efficace de l'AIGC. Bien que les méthodes de distillation de diffusion — basées sur l'appariement de distributions et l'appariement de trajectoires — réduisent l'échantillonnage à aussi peu qu'une étape, elles peinent à répondre aux tâches complexes comme la génération texte-image. La génération en quelques étapes offre un meilleur équilibre entre vitesse et qualité, mais les approches existantes sont confrontées à un compromis persistant : l'appariement de distributions manque de flexibilité pour l'échantillonnage multi-étapes, tandis que l'appariement de trajectoires produit souvent une qualité d'image sous-optimale. Pour combler cette lacune, nous proposons d'apprendre des modèles de diffusion en quelques étapes par Appariement de Distribution de Trajectoire (TDM), un paradigme de distillation unifié qui combine les forces de l'appariement de distributions et de trajectoires. Notre méthode introduit un objectif de distillation de score sans données, alignant la trajectoire de l'élève avec celle du professeur au niveau de la distribution. De plus, nous développons un objectif prenant en compte les étapes d'échantillonnage, qui découple les cibles d'apprentissage à travers différentes étapes, permettant un échantillonnage plus ajustable. Cette approche prend en charge à la fois l'échantillonnage déterministe pour une qualité d'image supérieure et l'adaptation flexible multi-étapes, atteignant des performances de pointe avec une efficacité remarquable. Notre modèle, TDM, surpasse les méthodes existantes sur diverses architectures, telles que SDXL et PixArt-alpha, offrant une qualité supérieure et des coûts d'entraînement significativement réduits. En particulier, notre méthode distille PixArt-alpha en un générateur à 4 étapes qui surpasse son modèle enseignant en termes de préférence utilisateur réelle à une résolution de 1024. Ceci est accompli avec 500 itérations et 2 heures sur A800 — soit seulement 0,01 % du coût d'entraînement du modèle enseignant. De plus, notre TDM proposé peut être étendu pour accélérer la diffusion texte-vidéo. Notamment, TDM peut surpasser son modèle enseignant (CogVideoX-2B) en utilisant seulement 4 NFE sur VBench, améliorant le score total de 80,91 à 81,65. Page du projet : https://tdm-t2x.github.io/

English

Accelerating diffusion model sampling is crucial for efficient AIGC deployment. While diffusion distillation methods -- based on distribution matching and trajectory matching -- reduce sampling to as few as one step, they fall short on complex tasks like text-to-image generation. Few-step generation offers a better balance between speed and quality, but existing approaches face a persistent trade-off: distribution matching lacks flexibility for multi-step sampling, while trajectory matching often yields suboptimal image quality. To bridge this gap, we propose learning few-step diffusion models by Trajectory Distribution Matching (TDM), a unified distillation paradigm that combines the strengths of distribution and trajectory matching. Our method introduces a data-free score distillation objective, aligning the student's trajectory with the teacher's at the distribution level. Further, we develop a sampling-steps-aware objective that decouples learning targets across different steps, enabling more adjustable sampling. This approach supports both deterministic sampling for superior image quality and flexible multi-step adaptation, achieving state-of-the-art performance with remarkable efficiency. Our model, TDM, outperforms existing methods on various backbones, such as SDXL and PixArt-alpha, delivering superior quality and significantly reduced training costs. In particular, our method distills PixArt-alpha into a 4-step generator that outperforms its teacher on real user preference at 1024 resolution. This is accomplished with 500 iterations and 2 A800 hours -- a mere 0.01% of the teacher's training cost. In addition, our proposed TDM can be extended to accelerate text-to-video diffusion. Notably, TDM can outperform its teacher model (CogVideoX-2B) by using only 4 NFE on VBench, improving the total score from 80.91 to 81.65. Project page: https://tdm-t2x.github.io/