[论文解读] Faster Diffusion: Rethinking the Role of UNet Encoder in Diffusion Models

Senmao Li, Taihang Hu+6 more|arXiv (Cornell University)|Dec 15, 2023

Advanced Neuroimaging Techniques and Applications被引用 5

一句话总结

本文提出了一种编码器传播方案，通过在扩散模型中复用相邻时间步的UNet编码器特征，显著加速采样过程，无需知识蒸馏即可实现Stable Diffusion最高41%和DeepFloyd-IF最高24%的加速。该方法基于编码器特征随时间变化极小的观察，通过引入先验噪声注入策略，在保持高图像质量的同时实现并行解码，提升推理效率。

ABSTRACT

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 \href{https://github.com/hutaiHang/Faster-Diffusion}{FasterDiffusion}.

研究动机与目标

• 研究扩散采样过程中UNet编码器与解码器特征的动态行为。
• 解决扩散模型中在每个时间步重复计算编码器特征导致的效率低下问题。
• 开发一种轻量化、无需蒸馏的加速方法，适用于多种条件生成任务。
• 在特征复用的情况下仍保留高保真度的图像细节，特别是在纹理丰富的区域。
• 在文本到图像、文本到视频、个性化生成及参考引导生成任务中验证该方法的有效性。

提出的方法

• 作者分析了时间步之间的特征演化，发现UNet编码器特征变化平缓，而解码器特征则显著变化。
• 提出一种编码器传播方案，在特征高度相似时复用前一时间步的编码器特征，从而减少冗余计算。
• 引入非均匀关键时间步选择策略，以确定哪些时间步可以安全跳过编码器推理。
• 通过复用相同编码器特征，实现在多个相邻时间步上的并行解码，进一步提升采样速度。
• 应用先验噪声注入机制，将初始噪声 $ z_T $ 注入潜在空间，以恢复因特征复用而丢失的纹理细节。
• 该方法端到端应用于多种扩散任务，包括文本到图像、文本到视频、个性化生成及参考引导生成。

实验结果

研究问题

• RQ1 在扩散采样过程中，UNet编码器与解码器特征如何随不同时间步演化？
• RQ2 在不降低生成质量的前提下，编码器特征能在多大程度上被复用？
• RQ3 编码器特征复用是否能实现相邻时间步的并行解码以加速推理？
• RQ4 编码器传播对纹理保真度有何影响，是否可被缓解？
• RQ5 所提出的方法是否能在无需知识蒸馏或模型微调的情况下，实现多种条件生成任务中的加速？

主要发现

• 编码器特征在时间步之间变化极小，相邻步骤间具有高度相似性，而解码器特征则表现出显著差异。
• 结合并行解码，所提出的编码器传播方案使Stable Diffusion采样速度提升41%，DeepFloyd-IF提升24%。
• 在MS-COCO 2017 10K数据集上，非均匀关键时间步选择策略在FID（21.08）和Clipscore（0.783）上均优于均匀策略。
• 先验噪声注入有效减少纹理退化，保持高保真结果，FID值仅从21.08轻微上升至21.75，与基线相比可忽略。
• 用户研究表明，该方法在图像质量上与基线模型相当或更优，感知质量无显著下降。
• 该方法在多种任务中均有效，包括文本到视频、个性化生成及参考引导生成，且无需知识蒸馏。