Обучение диффузионных моделей с малым количеством шагов через сопоставление распределений траекторий

Learning Few-Step Diffusion Models by Trajectory Distribution Matching

March 9, 2025

Авторы: Yihong Luo, Tianyang Hu, Jiacheng Sun, Yujun Cai, Jing Tang

cs.AI

Аннотация

Ускорение выборки в диффузионных моделях имеет решающее значение для эффективного развертывания AIGC. Хотя методы дистилляции диффузии, основанные на согласовании распределений и траекторий, позволяют сократить выборку до одного шага, они не справляются со сложными задачами, такими как генерация изображений по тексту. Генерация с малым количеством шагов предлагает лучший баланс между скоростью и качеством, но существующие подходы сталкиваются с постоянным компромиссом: согласование распределений недостаточно гибко для многошаговой выборки, а согласование траекторий часто приводит к неоптимальному качеству изображений. Чтобы устранить этот разрыв, мы предлагаем обучение диффузионных моделей с малым количеством шагов с помощью метода **Согласования Распределения Траекторий (TDM)**, унифицированной парадигмы дистилляции, которая сочетает преимущества согласования распределений и траекторий. Наш метод вводит объектив дистилляции оценок без данных, согласовывая траекторию ученика с траекторией учителя на уровне распределения. Кроме того, мы разрабатываем объектив, учитывающий шаги выборки, который разделяет цели обучения для разных шагов, обеспечивая более гибкую выборку. Этот подход поддерживает как детерминированную выборку для превосходного качества изображений, так и гибкую адаптацию для многошаговой выборки, достигая передовых показателей с высокой эффективностью. Наша модель, TDM, превосходит существующие методы на различных архитектурах, таких как SDXL и PixArt-alpha, обеспечивая лучшее качество и значительно снижая затраты на обучение. В частности, наш метод дистиллирует PixArt-alpha в 4-шаговый генератор, который превосходит своего учителя по предпочтениям реальных пользователей при разрешении 1024. Это достигается за 500 итераций и 2 часа на A800 — всего 0.01% от стоимости обучения учителя. Кроме того, предложенный TDM может быть расширен для ускорения диффузии текста в видео. Примечательно, что TDM может превзойти свою учительскую модель (CogVideoX-2B), используя всего 4 NFE на VBench, улучшая общий балл с 80.91 до 81.65. Страница проекта: 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/