数ステップ拡散モデルの学習：軌道分布マッチングによるアプローチ

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の展開において極めて重要です。拡散蒸留法は、分布マッチングと軌跡マッチングに基づいてサンプリングをわずか1ステップにまで削減しますが、テキストから画像生成のような複雑なタスクでは不十分です。数ステップ生成は速度と品質のバランスをより良くしますが、既存のアプローチでは持続的なトレードオフに直面しています。分布マッチングは多ステップサンプリングに対して柔軟性に欠け、軌跡マッチングはしばしば最適でない画像品質をもたらします。このギャップを埋めるため、我々は軌跡分布マッチング（TDM）による数ステップ拡散モデルの学習を提案します。これは分布マッチングと軌跡マッチングの強みを組み合わせた統一された蒸留パラダイムです。我々の手法は、データフリーのスコア蒸留目的関数を導入し、生徒の軌跡を教師の軌跡と分布レベルで整合させます。さらに、異なるステップ間で学習目標を分離するサンプリングステップを意識した目的関数を開発し、より調整可能なサンプリングを可能にします。このアプローチは、優れた画像品質のための決定論的サンプリングと柔軟な多ステップ適応の両方をサポートし、驚異的な効率で最先端の性能を達成します。我々のモデル、TDMは、SDXLやPixArt-alphaなどの様々なバックボーンにおいて既存の手法を上回り、優れた品質と大幅に削減されたトレーニングコストを提供します。特に、我々の手法はPixArt-alphaを4ステップ生成器に蒸留し、1024解像度での実際のユーザー選好において教師モデルを上回ります。これは500イテレーションと2A800時間で達成され、教師モデルのトレーニングコストのわずか0.01%です。さらに、提案したTDMはテキストからビデオ拡散の加速にも拡張可能です。特に、TDMはVBenchにおいてわずか4NFEで教師モデル（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/