Apprendimento di Modelli di Diffusione a Pochi Passi mediante Allineamento della Distribuzione delle Traiettorie

Learning Few-Step Diffusion Models by Trajectory Distribution Matching

March 9, 2025

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

cs.AI

Abstract

L'accelerazione del campionamento nei modelli di diffusione è cruciale per un'implementazione efficiente dell'AIGC. Sebbene i metodi di distillazione della diffusione -- basati sull'allineamento delle distribuzioni e sul matching delle traiettorie -- riducano il campionamento a un solo passo, essi non riescono a gestire compiti complessi come la generazione di immagini da testo. La generazione in pochi passi offre un migliore equilibrio tra velocità e qualità, ma gli approcci esistenti affrontano un persistente compromesso: l'allineamento delle distribuzioni manca di flessibilità per il campionamento multi-passo, mentre il matching delle traiettorie spesso produce una qualità dell'immagine subottimale. Per colmare questa lacuna, proponiamo di apprendere modelli di diffusione a pochi passi tramite il Trajectory Distribution Matching (TDM), un paradigma di distillazione unificato che combina i punti di forza dell'allineamento delle distribuzioni e del matching delle traiettorie. Il nostro metodo introduce un obiettivo di distillazione del punteggio senza dati, allineando la traiettoria dello studente con quella del docente a livello di distribuzione. Inoltre, sviluppiamo un obiettivo consapevole dei passi di campionamento che disaccoppia i target di apprendimento tra diversi passi, consentendo un campionamento più regolabile. Questo approccio supporta sia il campionamento deterministico per una qualità superiore dell'immagine, sia l'adattamento flessibile multi-passo, raggiungendo prestazioni all'avanguardia con un'efficienza notevole. Il nostro modello, TDM, supera i metodi esistenti su vari backbone, come SDXL e PixArt-alpha, offrendo una qualità superiore e costi di addestramento significativamente ridotti. In particolare, il nostro metodo distilla PixArt-alpha in un generatore a 4 passi che supera il suo docente nelle preferenze degli utenti reali a una risoluzione di 1024. Questo è stato ottenuto con 500 iterazioni e 2 ore su A800 -- appena lo 0,01% del costo di addestramento del docente. Inoltre, il nostro TDM proposto può essere esteso per accelerare la diffusione da testo a video. In modo significativo, TDM può superare il suo modello docente (CogVideoX-2B) utilizzando solo 4 NFE su VBench, migliorando il punteggio totale da 80,91 a 81,65. Pagina del progetto: 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/