Weinig-staps diffusiemodellen leren door trajectverdelingsovereenstemming

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

Samenvatting

Het versnellen van diffusiemodel sampling is cruciaal voor een efficiënte inzet van AIGC. Hoewel diffusiedistillatiemethoden – gebaseerd op distributie- en trajectmatching – het sampling kunnen reduceren tot slechts één stap, schieten ze tekort bij complexe taken zoals tekst-naar-beeldgeneratie. Generatie in enkele stappen biedt een betere balans tussen snelheid en kwaliteit, maar bestaande benaderingen kampen met een hardnekkige afweging: distributiematching mist flexibiliteit voor multi-step sampling, terwijl trajectmatching vaak suboptimale beeldkwaliteit oplevert. Om deze kloof te overbruggen, stellen we voor om diffusiemodellen in enkele stappen te leren via Trajectory Distribution Matching (TDM), een geïntegreerd distillatieparadigma dat de sterke punten van distributie- en trajectmatching combineert. Onze methode introduceert een data-vrij score-distillatiedoel, dat het traject van de student afstemt op dat van de leraar op distributieniveau. Verder ontwikkelen we een sampling-stappen-bewust doel dat leerdoelen ontkoppelt over verschillende stappen, waardoor meer aanpasbaar sampling mogelijk wordt. Deze benadering ondersteunt zowel deterministisch sampling voor superieure beeldkwaliteit als flexibele multi-step aanpassing, wat resulteert in state-of-the-art prestaties met opmerkelijke efficiëntie. Ons model, TDM, overtreft bestaande methoden op verschillende backbones, zoals SDXL en PixArt-alpha, en levert superieure kwaliteit met aanzienlijk gereduceerde trainingskosten. In het bijzonder destilleert onze methode PixArt-alpha tot een 4-staps generator die zijn leraar overtreft op basis van echte gebruikersvoorkeuren bij 1024 resolutie. Dit wordt bereikt met 500 iteraties en 2 A800 uren – slechts 0,01% van de trainingskosten van de leraar. Daarnaast kan onze voorgestelde TDM worden uitgebreid om tekst-naar-video diffusie te versnellen. Opmerkelijk is dat TDM zijn lerarmodel (CogVideoX-2B) kan overtreffen door slechts 4 NFE te gebruiken op VBench, waarbij de totale score verbetert van 80,91 naar 81,65. Projectpagina: 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/