Difusão Mais Rápida: Repensando o Papel do Codificador UNet em Modelos de Difusão

Faster Diffusion: Rethinking the Role of UNet Encoder in Diffusion Models

December 15, 2023

Autores: Senmao Li, Taihang Hu, Fahad Shahbaz Khan, Linxuan Li, Shiqi Yang, Yaxing Wang, Ming-Ming Cheng, Jian Yang

cs.AI

Resumo

Um dos componentes-chave nos modelos de difusão é a UNet para previsão de ruído. Embora vários trabalhos tenham explorado as propriedades básicas do decodificador da UNet, seu codificador permanece amplamente inexplorado. Neste trabalho, realizamos o primeiro estudo abrangente do codificador da UNet. Analisamos empiricamente as características do codificador e fornecemos insights para questões importantes sobre suas mudanças durante o processo de inferência. Em particular, descobrimos que as características do codificador mudam suavemente, enquanto as características do decodificador exibem variações substanciais em diferentes passos de tempo. Essa descoberta nos inspirou a omitir o codificador em certos passos de tempo adjacentes e reutilizar ciclicamente as características do codificador dos passos de tempo anteriores para o decodificador. Com base nessa observação, introduzimos um esquema de propagação do codificador simples, porém eficaz, para acelerar a amostragem de difusão em um conjunto diversificado de tarefas. Ao nos beneficiarmos do nosso esquema de propagação, conseguimos executar em paralelo o decodificador em certos passos de tempo adjacentes. Além disso, introduzimos um método de injeção de ruído prévio para melhorar os detalhes de textura na imagem gerada. Além da tarefa padrão de texto para imagem, também validamos nossa abordagem em outras tarefas: texto para vídeo, geração personalizada e geração guiada por referência. Sem utilizar nenhuma técnica de destilação de conhecimento, nossa abordagem acelera a amostragem dos modelos Stable Diffusion (SD) e DeepFloyd-IF em 41% e 24%, respectivamente, mantendo um desempenho de geração de alta qualidade. Nosso código está disponível em https://github.com/hutaiHang/Faster-Diffusion{FasterDiffusion}.

English

One of the key components within diffusion models is the UNet for noise prediction. While several works have explored basic properties of the UNet decoder, its encoder largely remains unexplored. In this work, we conduct the first comprehensive study of the UNet encoder. We empirically analyze the encoder features and provide insights to important questions regarding their changes at the inference process. In particular, we find that encoder features change gently, whereas the decoder features exhibit substantial variations across different time-steps. This finding inspired us to omit the encoder at certain adjacent time-steps and reuse cyclically the encoder features in the previous time-steps for the decoder. Further based on this observation, we introduce a simple yet effective encoder propagation scheme to accelerate the diffusion sampling for a diverse set of tasks. By benefiting from our propagation scheme, we are able to perform in parallel the decoder at certain adjacent time-steps. Additionally, we introduce a prior noise injection method to improve the texture details in the generated image. Besides the standard text-to-image task, we also validate our approach on other tasks: text-to-video, personalized generation and reference-guided generation. Without utilizing any knowledge distillation technique, our approach accelerates both the Stable Diffusion (SD) and the DeepFloyd-IF models sampling by 41% and 24% respectively, while maintaining high-quality generation performance. Our code is available in https://github.com/hutaiHang/Faster-Diffusion{FasterDiffusion}.