T-Stitch: 事前学習済み拡散モデルのサンプリングを軌道接続で高速化
T-Stitch: Accelerating Sampling in Pre-Trained Diffusion Models with Trajectory Stitching
February 21, 2024
著者: Zizheng Pan, Bohan Zhuang, De-An Huang, Weili Nie, Zhiding Yu, Chaowei Xiao, Jianfei Cai, Anima Anandkumar
cs.AI
要旨
拡散確率モデル（DPM）からのサンプリングは、高品質な画像生成においてしばしば計算コストが高く、通常、大規模なモデルを使用して多くのステップを必要とします。本論文では、サンプリング効率を向上させ、生成品質の低下を最小限に抑えるためのシンプルで効率的な手法であるTrajectory Stitching T-Stitchを紹介します。T-Stitchは、サンプリング軌跡全体にわたって大規模なDPMのみを使用するのではなく、初期ステップにおいてより小規模なDPMを大規模なDPMの代替として活用し、後段階で大規模なDPMに切り替えます。我々の重要な洞察は、異なる拡散モデルが同じ訓練データ分布の下で類似したエンコーディングを学習し、小規模なモデルが初期ステップにおいて良好なグローバル構造を生成できるということです。広範な実験により、T-Stitchが訓練不要であり、異なるアーキテクチャに一般的に適用可能で、既存の高速サンプリング手法と柔軟な速度と品質のトレードオフを補完することが示されています。例えば、DiT-XLにおいて、クラス条件付きImageNet生成において性能低下なしに、初期タイムステップの40%を10倍高速なDiT-Sに安全に置き換えることができます。さらに、我々の手法が、人気のある事前訓練済みのStable Diffusion（SD）モデルの加速だけでなく、公開モデルゾーンのスタイル化されたSDモデルのプロンプトアラインメントを改善するためのドロップインテクニックとしても使用できることを示します。コードはhttps://github.com/NVlabs/T-Stitchで公開されています。
English
Sampling from diffusion probabilistic models (DPMs) is often expensive for
high-quality image generation and typically requires many steps with a large
model. In this paper, we introduce sampling Trajectory Stitching T-Stitch, a
simple yet efficient technique to improve the sampling efficiency with little
or no generation degradation. Instead of solely using a large DPM for the
entire sampling trajectory, T-Stitch first leverages a smaller DPM in the
initial steps as a cheap drop-in replacement of the larger DPM and switches to
the larger DPM at a later stage. Our key insight is that different diffusion
models learn similar encodings under the same training data distribution and
smaller models are capable of generating good global structures in the early
steps. Extensive experiments demonstrate that T-Stitch is training-free,
generally applicable for different architectures, and complements most existing
fast sampling techniques with flexible speed and quality trade-offs. On DiT-XL,
for example, 40% of the early timesteps can be safely replaced with a 10x
faster DiT-S without performance drop on class-conditional ImageNet generation.
We further show that our method can also be used as a drop-in technique to not
only accelerate the popular pretrained stable diffusion (SD) models but also
improve the prompt alignment of stylized SD models from the public model zoo.
Code is released at https://github.com/NVlabs/T-Stitch