T-Stitch: Accelerazione del Campionamento nei Modelli di Diffusione Pre-Addestrati con Cucitura delle Traiettorie
T-Stitch: Accelerating Sampling in Pre-Trained Diffusion Models with Trajectory Stitching
February 21, 2024
Autori: Zizheng Pan, Bohan Zhuang, De-An Huang, Weili Nie, Zhiding Yu, Chaowei Xiao, Jianfei Cai, Anima Anandkumar
cs.AI
Abstract
Il campionamento da modelli probabilistici di diffusione (DPM) è spesso costoso per la generazione di immagini di alta qualità e richiede tipicamente molti passaggi con un modello di grandi dimensioni. In questo articolo, introduciamo il campionamento Trajectory Stitching T-Stitch, una tecnica semplice ma efficiente per migliorare l'efficienza del campionamento con un degrado minimo o nullo nella generazione. Invece di utilizzare esclusivamente un DPM di grandi dimensioni per l'intera traiettoria di campionamento, T-Stitch sfrutta inizialmente un DPM più piccolo nei primi passi come sostituto economico del DPM più grande, per poi passare al DPM più grande in una fase successiva. La nostra intuizione chiave è che diversi modelli di diffusione apprendono codifiche simili sotto la stessa distribuzione di dati di addestramento e che modelli più piccoli sono in grado di generare buone strutture globali nei primi passi. Esperimenti estesi dimostrano che T-Stitch è privo di addestramento, generalmente applicabile a diverse architetture e complementare alla maggior parte delle tecniche di campionamento veloce esistenti, con compromessi flessibili tra velocità e qualità. Su DiT-XL, ad esempio, il 40% dei primi intervalli di tempo può essere sostituito in modo sicuro con un DiT-S 10 volte più veloce senza perdita di prestazioni nella generazione condizionata alla classe su ImageNet. Mostriamo inoltre che il nostro metodo può essere utilizzato come tecnica di sostituzione non solo per accelerare i popolari modelli pre-addestrati di diffusione stabile (SD), ma anche per migliorare l'allineamento del prompt nei modelli SD stilizzati provenienti dal modello pubblico zoo. Il codice è rilasciato all'indirizzo 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