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
摘要
从扩散概率模型（DPMs）中采样通常对于高质量图像生成而言成本高昂，通常需要许多步骤以及一个庞大的模型。在本文中，我们介绍了一种名为采样轨迹拼接（T-Stitch）的简单而高效的技术，以提高采样效率，减少或不降低生成质量。T-Stitch并非仅仅在整个采样轨迹中使用一个大型DPM，而是首先利用较小的DPM作为较便宜的替代品来代替较大的DPM的初始步骤，并在后期切换到较大的DPM。我们的关键见解是，在相同的训练数据分布下，不同的扩散模型学习类似的编码，并且较小的模型能够在早期步骤中生成良好的全局结构。大量实验证明，T-Stitch无需训练，在不同架构中通用，并且能够与大多数现有的快速采样技术相辅相成，具有灵活的速度和质量权衡。例如，在DiT-XL上，可以安全地用速度快10倍的DiT-S替换40%的早期时间步，而在有条件类别的ImageNet生成中不会降低性能。我们进一步展示，我们的方法不仅可以用作加速流行的预训练稳定扩散（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