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 — простой, но эффективный способ повышения эффективности выборки с минимальной или нулевой деградацией качества генерации. Вместо использования исключительно крупной DPM на всем протяжении траектории выборки, T-Stitch сначала задействует меньшую DPM на начальных этапах в качестве экономичной замены более крупной модели, а затем переключается на крупную DPM на более поздних стадиях. Наше ключевое наблюдение заключается в том, что различные диффузионные модели обучаются схожим представлениям при одинаковом распределении обучающих данных, а меньшие модели способны генерировать качественные глобальные структуры на ранних этапах. Многочисленные эксперименты показывают, что T-Stitch не требует дополнительного обучения, применим к различным архитектурам и дополняет большинство существующих методов быстрой выборки, обеспечивая гибкий баланс между скоростью и качеством. Например, на модели DiT-XL 40% начальных временных шагов могут быть безопасно заменены на 10 раз более быструю DiT-S без потери производительности при генерации изображений на основе классов в ImageNet. Мы также демонстрируем, что наш метод может использоваться как встраиваемый подход не только для ускорения популярных предобученных моделей 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