BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion

要約

安定拡散モデル (SDM) を使用したテキストから画像への生成 (T2I) には、10 億規模のパラメーターによる高いコンピューティング要求が伴います。
効率を高めるために、最近の研究では、元のアーキテクチャを維持しながら、サンプリング ステップを削減し、ネットワーク量子化を適用しました。
アーキテクチャの削減の試みが不足しているのは、このような大規模なモデルの再トレーニングに費用がかかるという懸念から生じている可能性があります。
この研究では、ブロック プルーニングの驚くべき可能性を明らかにし、低コストの汎用 T2I の蒸留を特集します。
SDM の U-Net からいくつかの残留ブロックとアテンション ブロックを削除することで、モデル サイズ、MAC、レイテンシーの 30% ～ 50% の削減を達成します。
限られたリソースの下でも蒸留の再トレーニングが効果的であることを示します。わずか 13 A100 日と小さなデータセットを使用して、コンパクトなモデルは元の SDM (6,000 A100 日を超える v1.4 および v2.1 ベース) を模倣できます。
移転された知識の恩恵を受けて、当社の BK-SDM は、ゼロショット MS-COCO で数十億のより大きなパラメータ モデルに対して競争力のある結果をもたらします。
さらに、パーソナライズされた生成と画像間の変換における軽量バックボーンの適用性を実証します。
エッジ デバイスにモデルを展開すると、4 秒の推論が達成されます。
コードとモデルは、https://github.com/Nota-NetsPresso/BK-SDM で見つけることができます。

要約(オリジナル)

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

arxiv情報

提供元, 利用サービス

arxiv.jp, Google