Aprendendo Modelos de Difusão de Poucos Passos por Correspondência de Distribuição de Trajetória

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

Resumo

Acelerar a amostragem de modelos de difusão é crucial para a implantação eficiente de AIGC. Embora os métodos de destilação de difusão — baseados em correspondência de distribuição e correspondência de trajetória — reduzam a amostragem para tão poucos quanto um passo, eles ficam aquém em tarefas complexas como geração de texto para imagem. A geração em poucos passos oferece um melhor equilíbrio entre velocidade e qualidade, mas as abordagens existentes enfrentam uma troca persistente: a correspondência de distribuição carece de flexibilidade para amostragem em múltiplos passos, enquanto a correspondência de trajetória frequentemente resulta em qualidade de imagem subótima. Para preencher essa lacuna, propomos aprender modelos de difusão em poucos passos por Correspondência de Distribuição de Trajetória (TDM), um paradigma unificado de destilação que combina os pontos fortes da correspondência de distribuição e de trajetória. Nosso método introduz um objetivo de destilação de pontuação sem dados, alinhando a trajetória do aluno com a do professor no nível de distribuição. Além disso, desenvolvemos um objetivo consciente dos passos de amostragem que desacopla os alvos de aprendizado em diferentes passos, permitindo uma amostragem mais ajustável. Essa abordagem suporta tanto a amostragem determinística para qualidade superior de imagem quanto a adaptação flexível em múltiplos passos, alcançando desempenho de ponta com eficiência notável. Nosso modelo, TDM, supera os métodos existentes em várias arquiteturas, como SDXL e PixArt-alpha, entregando qualidade superior e custos de treinamento significativamente reduzidos. Em particular, nosso método destila o PixArt-alpha em um gerador de 4 passos que supera seu professor na preferência do usuário real em resolução de 1024. Isso é realizado com 500 iterações e 2 horas de A800 — meros 0,01% do custo de treinamento do professor. Além disso, nosso TDM proposto pode ser estendido para acelerar a difusão de texto para vídeo. Notavelmente, o TDM pode superar seu modelo professor (CogVideoX-2B) usando apenas 4 NFE no VBench, melhorando a pontuação total de 80,91 para 81,65. Página do projeto: 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/