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
초록
확산 확률 모델(Diffusion Probabil Models, DPM)에서 샘플링은 고품질 이미지 생성을 위해 종종 비용이 많이 들며, 일반적으로 대형 모델을 사용하여 많은 단계를 거쳐야 합니다. 본 논문에서는 샘플링 효율성을 크게 향상시키면서도 생성 품질의 저하가 거의 없거나 전혀 없는 간단하면서도 효율적인 기법인 Trajectory Stitching T-Stitch를 소개합니다. T-Stitch는 전체 샘플링 궤적에 대해 대형 DPM만을 사용하는 대신, 초기 단계에서 더 작은 DPM을 대형 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