T-Stitch : Accélération de l'échantillonnage dans les modèles de diffusion pré-entraînés par couture de trajectoires
T-Stitch: Accelerating Sampling in Pre-Trained Diffusion Models with Trajectory Stitching
February 21, 2024
Auteurs: Zizheng Pan, Bohan Zhuang, De-An Huang, Weili Nie, Zhiding Yu, Chaowei Xiao, Jianfei Cai, Anima Anandkumar
cs.AI
Résumé
L'échantillonnage à partir de modèles probabilistes de diffusion (DPM) est souvent coûteux pour la génération d'images de haute qualité et nécessite généralement de nombreuses étapes avec un modèle de grande taille. Dans cet article, nous présentons T-Stitch, une technique simple mais efficace pour améliorer l'efficacité de l'échantillonnage avec peu ou pas de dégradation de la génération. Au lieu d'utiliser uniquement un grand DPM pour l'ensemble de la trajectoire d'échantillonnage, T-Stitch exploite d'abord un DPM plus petit dans les étapes initiales comme un remplacement économique du DPM plus grand, puis bascule vers le DPM plus grand à un stade ultérieur. Notre idée clé est que différents modèles de diffusion apprennent des encodages similaires sous la même distribution de données d'entraînement, et que les modèles plus petits sont capables de générer de bonnes structures globales dans les premières étapes. Des expériences approfondies démontrent que T-Stitch ne nécessite pas d'entraînement, est généralement applicable à différentes architectures, et complète la plupart des techniques d'échantillonnage rapide existantes avec des compromis flexibles entre vitesse et qualité. Sur DiT-XL, par exemple, 40 % des premières étapes temporelles peuvent être remplacées en toute sécurité par un DiT-S 10 fois plus rapide sans perte de performance sur la génération conditionnelle par classe d'ImageNet. Nous montrons en outre que notre méthode peut également être utilisée comme une technique de remplacement pour non seulement accélérer les modèles pré-entraînés populaires de diffusion stable (SD), mais aussi améliorer l'alignement des prompts des modèles SD stylisés provenant du zoo de modèles publics. Le code est disponible à l'adresse suivante : 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