[論文レビュー] 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 vs. 21.75）に抑えられ、高精細な結果が維持された。
• ユーザースタディーにおいて、本手法はベースラインと同等または優れた画像品質を達成しており、知覚的品質に顕著な低下は認められなかった。
• 知識蒸留を必要とせず、テキストから動画、パーソナライズド、リファレンスガイドド生成など多様なタスクにおいても本手法は有効である。