T-Stitch: Aceleración del muestreo en modelos de difusión preentrenados mediante unión de trayectorias
T-Stitch: Accelerating Sampling in Pre-Trained Diffusion Models with Trajectory Stitching
February 21, 2024
Autores: Zizheng Pan, Bohan Zhuang, De-An Huang, Weili Nie, Zhiding Yu, Chaowei Xiao, Jianfei Cai, Anima Anandkumar
cs.AI
Resumen
El muestreo a partir de modelos probabilísticos de difusión (DPMs, por sus siglas en inglés) suele ser costoso para la generación de imágenes de alta calidad y generalmente requiere muchos pasos con un modelo grande. En este artículo, presentamos T-Stitch, una técnica simple pero eficiente para mejorar la eficiencia del muestreo con poca o ninguna degradación en la generación. En lugar de utilizar únicamente un DPM grande para toda la trayectoria de muestreo, T-Stitch aprovecha primero un DPM más pequeño en los pasos iniciales como un reemplazo económico del DPM más grande y cambia al DPM más grande en una etapa posterior. Nuestra idea clave es que diferentes modelos de difusión aprenden codificaciones similares bajo la misma distribución de datos de entrenamiento, y los modelos más pequeños son capaces de generar buenas estructuras globales en los pasos iniciales. Experimentos extensos demuestran que T-Stitch no requiere entrenamiento, es aplicable de manera general a diferentes arquitecturas y complementa la mayoría de las técnicas de muestreo rápido existentes con compensaciones flexibles entre velocidad y calidad. En DiT-XL, por ejemplo, el 40% de los pasos iniciales puede ser reemplazado de manera segura por un DiT-S 10 veces más rápido sin pérdida de rendimiento en la generación condicionada por clases de ImageNet. Además, mostramos que nuestro método también puede utilizarse como una técnica de reemplazo directo para no solo acelerar los populares modelos preentrenados de difusión estable (SD, por sus siglas en inglés), sino también mejorar la alineación con el prompt en modelos SD estilizados del repositorio público de modelos. El código está disponible en 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