通过轨迹分布匹配，实现高效少步扩散模型训练，兼顾速度与图像质量。
Learning Few-Step Diffusion Models by Trajectory Distribution Matching
- 提出统一的轨迹分布匹配框架，融合分布与轨迹匹配优势。
- 4步生成器在1024分辨率下超越教师模型，仅需2小时训练。
- 适用于文本到图像、文本到视频，显著降低训练成本。
加速扩散模型采样对高效AIGC部署至关重要。现有扩散蒸馏方法虽可将采样缩减至一步，但在复杂任务如文生图中表现不足。少步生成在速度与质量间取得更好平衡，但现有方法存在固有权衡：分布匹配缺乏多步灵活性，轨迹匹配常导致图像质量下降。为此，本文提出基于轨迹分布匹配（TDM）的少步扩散模型学习方法，统一结合分布与轨迹匹配优势。方法引入无数据得分蒸馏目标，在分布层面对齐学生与教师的轨迹；同时设计采样步数感知目标，解耦不同步骤的学习目标，提升采样灵活性。该方法支持确定性采样以获得更优图像质量，并实现显著效率提升。TDM模型在SDXL、PixArt-α等骨干网络上均达到当前最优性能。特别地，其将PixArt-α蒸馏为4步生成器，在1024分辨率下优于教师模型的真实用户偏好，仅用500次迭代和2个A800小时训练——仅为教师训练成本的0.01%。此外，TDM可扩展用于加速文生视频扩散，仅用4个采样步数即在VBench上超越教师模型CogVideoX-2B，总分从80.91提升至81.65。
原文摘要 · Abstract (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-$α$, delivering superior quality and significantly reduced training costs. In particular, our method distills PixArt-$α$ 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/
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。