트레이젝토리 분포 매칭을 통한 Few-Step 확산 모델 학습

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은 VBench에서 단 4 NFE만 사용하여 교사 모델(CogVideoX-2B)을 능가할 수 있으며, 총 점수를 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/