Aprendizaje de Modelos de Difusión de Pocos Pasos mediante Emparejamiento de Distribución de Trayectorias

Learning Few-Step Diffusion Models by Trajectory Distribution Matching

March 9, 2025

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

cs.AI

Resumen

Acelerar el muestreo de modelos de difusión es crucial para un despliegue eficiente de AIGC. Si bien los métodos de destilación de difusión —basados en la coincidencia de distribuciones y la coincidencia de trayectorias— reducen el muestreo a tan solo un paso, no alcanzan un rendimiento óptimo en tareas complejas como la generación de texto a imagen. La generación en pocos pasos ofrece un mejor equilibrio entre velocidad y calidad, pero los enfoques existentes enfrentan un dilema persistente: la coincidencia de distribuciones carece de flexibilidad para el muestreo en múltiples pasos, mientras que la coincidencia de trayectorias a menudo produce una calidad de imagen subóptima. Para cerrar esta brecha, proponemos aprender modelos de difusión en pocos pasos mediante Coincidencia de Distribución de Trayectoria (TDM), un paradigma de destilación unificado que combina las fortalezas de la coincidencia de distribuciones y trayectorias. Nuestro método introduce un objetivo de destilación de puntuación sin datos, alineando la trayectoria del estudiante con la del profesor a nivel de distribución. Además, desarrollamos un objetivo consciente de los pasos de muestreo que desacopla los objetivos de aprendizaje en diferentes pasos, permitiendo un muestreo más ajustable. Este enfoque admite tanto el muestreo determinista para una calidad de imagen superior como la adaptación flexible en múltiples pasos, logrando un rendimiento de vanguardia con una eficiencia notable. Nuestro modelo, TDM, supera a los métodos existentes en varias arquitecturas, como SDXL y PixArt-alpha, ofreciendo una calidad superior y costos de entrenamiento significativamente reducidos. En particular, nuestro método destila PixArt-alpha en un generador de 4 pasos que supera a su profesor en la preferencia de usuarios reales a una resolución de 1024. Esto se logra con 500 iteraciones y 2 horas en A800 —un mero 0.01% del costo de entrenamiento del profesor. Además, nuestro TDM propuesto puede extenderse para acelerar la difusión de texto a video. Notablemente, TDM puede superar a su modelo profesor (CogVideoX-2B) utilizando solo 4 NFE en VBench, mejorando la puntuación total de 80.91 a 81.65. Página del proyecto: 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/