T-Stitch: Acelerando a Amostragem em Modelos de Difusão Pré-Treinados com Costura de Trajetórias
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
Resumo
A amostragem de modelos probabilísticos de difusão (DPMs) é frequentemente custosa para a geração de imagens de alta qualidade e geralmente requer muitas etapas com um modelo grande. Neste artigo, introduzimos o T-Stitch, uma técnica simples, porém eficiente, para melhorar a eficiência da amostragem com pouca ou nenhuma degradação na geração. Em vez de usar apenas um DPM grande para toda a trajetória de amostragem, o T-Stitch primeiro aproveita um DPM menor nas etapas iniciais como uma substituição barata do DPM maior e alterna para o DPM maior em uma etapa posterior. Nossa principal percepção é que diferentes modelos de difusão aprendem codificações semelhantes sob a mesma distribuição de dados de treinamento e que modelos menores são capazes de gerar boas estruturas globais nas etapas iniciais. Experimentos extensivos demonstram que o T-Stitch não requer treinamento, é geralmente aplicável a diferentes arquiteturas e complementa a maioria das técnicas de amostragem rápida existentes com trade-offs flexíveis entre velocidade e qualidade. No DiT-XL, por exemplo, 40% das etapas iniciais podem ser seguramente substituídas por um DiT-S 10x mais rápido sem queda de desempenho na geração condicional por classe do ImageNet. Mostramos ainda que nosso método também pode ser usado como uma técnica de substituição para não apenas acelerar os populares modelos pré-treinados de difusão estável (SD), mas também melhorar o alinhamento de prompts de modelos SD estilizados do repositório público de modelos. O código está disponível em 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