Lernen von Few-Step-Diffusionsmodellen durch Trajektorienverteilungsabgleich

Learning Few-Step Diffusion Models by Trajectory Distribution Matching

March 9, 2025

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

cs.AI

Zusammenfassung

Die Beschleunigung der Abtastung von Diffusionsmodellen ist entscheidend für den effizienten Einsatz von AIGC. Während Diffusionsdistillationsmethoden – basierend auf Verteilungsabgleich und Trajektorienabgleich – die Abtastung auf so wenig wie einen Schritt reduzieren, erreichen sie bei komplexen Aufgaben wie der Text-zu-Bild-Generierung ihre Grenzen. Die Generierung in wenigen Schritten bietet eine bessere Balance zwischen Geschwindigkeit und Qualität, doch bestehende Ansätze stehen vor einem anhaltenden Kompromiss: Der Verteilungsabgleich mangelt es an Flexibilität für die Mehrschritt-Abtastung, während der Trajektorienabgleich oft suboptimale Bildqualität liefert. Um diese Lücke zu schließen, schlagen wir vor, Diffusionsmodelle mit wenigen Schritten durch Trajektorien-Verteilungsabgleich (TDM) zu lernen, ein vereinheitlichtes Distillationsparadigma, das die Stärken von Verteilungs- und Trajektorienabgleich kombiniert. Unsere Methode führt ein datenfreies Score-Distillationsziel ein, das die Trajektorie des Schülers auf Verteilungsebene mit der des Lehrers abstimmt. Darüber hinaus entwickeln wir ein abtastschrittbewusstes Ziel, das Lernziele über verschiedene Schritte entkoppelt und eine anpassbarere Abtastung ermöglicht. Dieser Ansatz unterstützt sowohl deterministische Abtastung für überlegene Bildqualität als auch flexible Mehrschritt-Anpassung und erreicht dabei Spitzenleistungen mit bemerkenswerter Effizienz. Unser Modell, TDM, übertrifft bestehende Methoden auf verschiedenen Backbones wie SDXL und PixArt-alpha und liefert überlegene Qualität bei deutlich reduzierten Trainingskosten. Insbesondere destilliert unsere Methode PixArt-alpha in einen 4-Schritt-Generator, der seinen Lehrer bei der realen Benutzerpräferenz in 1024-Auflösung übertrifft. Dies wird mit 500 Iterationen und 2 A800-Stunden erreicht – lediglich 0,01 % der Trainingskosten des Lehrers. Zusätzlich kann unser vorgeschlagenes TDM erweitert werden, um die Text-zu-Video-Diffusion zu beschleunigen. Bemerkenswerterweise kann TDM sein Lehrermodell (CogVideoX-2B) übertreffen, indem es nur 4 NFE auf VBench verwendet und die Gesamtpunktzahl von 80,91 auf 81,65 verbessert. Projektseite: 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/