[論文レビュー] BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion

Bo-Kyeong Kim, Hyoung‐Kyu Song+2 more|arXiv (Cornell University)|May 25, 2023

Generative Adversarial Networks and Image Synthesis被引用数 5

ひとこと要約

本稿では、U-Netアーキテクチャから残留接続およびアテンションブロックをプルーニングし、その後に特徴量レベルの知識蒸留再訓練を実施することで、モデルサイズを最大51%削減し、遅延を43%改善する軽量で高速かつ低コストなStable Diffusionの変種、BK-SDMを紹介する。この手法は、わずか13 A100日分の計算リソースと小さなデータセットを用いて実施され、競争力のあるゼロショットテキスト-to-画像生成を実現するとともに、4秒未塔の推論時間でエッジデバイスへの効率的なデプロイが可能である。

ABSTRACT

Text-to-image (T2I) generation with Stable Diffusion models (SDMs) involves high computing demands due to billion-scale parameters. To enhance efficiency, recent studies have reduced sampling steps and applied network quantization while retaining the original architectures. The lack of architectural reduction attempts may stem from worries over expensive retraining for such massive models. In this work, we uncover the surprising potential of block pruning and feature distillation for low-cost general-purpose T2I. By removing several residual and attention blocks from the U-Net of SDMs, we achieve 30%~50% reduction in model size, MACs, and latency. We show that distillation retraining is effective even under limited resources: using only 13 A100 days and a tiny dataset, our compact models can imitate the original SDMs (v1.4 and v2.1-base with over 6,000 A100 days). Benefiting from the transferred knowledge, our BK-SDMs deliver competitive results on zero-shot MS-COCO against larger multi-billion parameter models. We further demonstrate the applicability of our lightweight backbones in personalized generation and image-to-image translation. Deployment of our models on edge devices attains 4-second inference. Code and models can be found at: https://github.com/Nota-NetsPresso/BK-SDM

研究の動機と目的

• 大規模なテキスト-to-画像拡散モデル（例：Stable Diffusion）の計算コストおよびトレーニングコストを、性能を損なわず低減すること。
• 大規模な拡散モデルにおいて、再トレーニングのコストが高いため未だ十分に調査されていないが、アーキテクチャの圧縮としてのブロックプルーニングの可能性を検討すること。
• 限られたデータおよび計算リソースを用いても、大規模な事前学習済みモデルからコンパクトでプルーニングされたバージョンへの知識蒸留が効果的に機能することを示すこと。
• Jetson AGX Orin や iPhone 14 などのエッジデバイスに、4秒未塔の推論時間で効率的にデプロイ可能なコンパクトな拡散モデルを実現すること。
• 研究者や開発者が現実的なリソース予算で、小規模だが強力な拡散モデルをトレーニングおよびデプロイできる実用的でアクセスしやすい道筋を提供すること。

提案手法

• Stable DiffusionのU-Netアーキテクチャにブロックプルーニングを適用し、複数の残留接続およびアテンションブロックを削除することで、モデルサイズとMACs（Multiply-Accumulate Operations）を削減する。
• 特徴量レベルの知識蒸留を用いてプルーニング後のモデルを再トレーニングし、元の大きなモデルからの知識を保持することで性能を維持する。
• 蒸留プロセスでは、0.22MのLAIONペアから成る小さなデータセットと、わずか13 A100日分のトレーニング計算リソースを用いることで、再トレーニングコストを顕著に削減する。
• 本手法はSDM-v1.4およびSDM-v2.1-baseの両方のモデルに適用され、異なるモデルバージョンにおいて一貫した性能向上が得られた。
• プルーニングおよび蒸留されたモデルは、ゼロショットテキスト-to-画像生成、パーソナライズド生成、画像to画像変換、エッジデプロイの各タスクで評価された。
• 顔生成を目的とした条件付きLDM（CelebA-HQ）を用いた実験により、本手法はテキスト-to-画像タスクにとどまらず、一般化性を確認した。

実験結果

リサーチクエスチョン

• RQ1 Stable DiffusionのU-Netにおける残留接続およびアテンションブロックのアーキテクチャ的プルーニングは、性能を劣化させることなく、モデルサイズおよび推論遅延を顕著に削減できるか？
• RQ2 限られたデータおよび計算リソースを用いても、大規模な事前学習済み拡散モデルの能力を、コンパクトでプルーニングされたバージョンに知識蒸留で効果的に移転できるか？
• RQ3 軽量な拡散モデルは、数十億パラメータを持つモデルと比較して、競争力あるゼロショットテキスト-to-画像生成性能を達成できるか？
• RQ4 プルーニングおよび蒸留されたモデルは、微調整コストを削減しつつ、パーソナライズドなテキスト-to-画像生成および画像to画像変換に効果的に使用できるか？
• RQ5 このようなコンパクトなモデルを、リアルタイムの推論速度を満たすエッジデバイスにデプロイすることは可能か？

主な発見

• BK-SDMは、CPUおよびGPU上ですべて最大51%のモデルサイズ削減と43%の遅延改善を達成し、MACsおよび推論時間についても30–50%の削減を実現した。
• 蒸留されたモデルは、わずか13 A100日分の計算リソースと0.22MのLAIONペアを用いてMS-COCO 30KでFID14.61を達成し、はるかに多くのリソースを要する大規模モデルと同等の性能を示した。
• BK-SDM-Smallは、Jetson AGX OrinおよびiPhone 14の両方で512×512の画像生成を4秒未塔で実行でき、エッジデプロイの実現可能性を強く示した。
• DreamBoothを用いたパーソナライズド生成において、BK-SDMは元のSDMと比較して95–99%の被写体およびプロンプト忠実度を維持しながら、微調整コストを削減した。
• 本手法は他の拡散モデルに対しても一般化可能であり、187Mパラメータを持つ条件付きLDMを圧縮し、高品質な顔生成を達成した。
• パラメータ数およびトレーニングデータを削減しても、BK-SDMモデルはCLIPスコアおよびISスコアにおいて競争力のある生成品質を維持した。