[论文解读] 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

一句话总结

本论文提出BK-SDM，一种轻量化、快速且成本低廉的Stable Diffusion变体，通过从U-Net架构中剪枝残差块和注意力块，再结合仅使用13个A100天训练时间和小规模数据集进行特征蒸馏微调，实现了模型参数量最高减少51%、延迟降低43%。该方法在保持性能的同时，显著降低了计算与训练成本，支持在边缘设备上实现高效部署，推理时间低于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

研究动机与目标

• 在不牺牲性能的前提下，降低Stable Diffusion等大规模文本到图像扩散模型的计算与训练成本。
• 探索在大规模扩散模型中通过块剪枝实现架构压缩的可行性，该方向因对昂贵微调的担忧而长期被忽视。
• 证明在数据和计算资源有限的条件下，知识蒸馏可有效将大预训练模型的知识迁移至紧凑的剪枝模型中。
• 实现在Jetson AGX Orin和iPhone 14等边缘设备上高效部署紧凑扩散模型，实现亚4秒内推理。
• 为研究人员和开发者提供一条切实可行、资源预算可控的路径，用于训练和部署小型但高性能的扩散模型。

提出的方法

• 对Stable Diffusion的U-Net架构应用块剪枝，移除多个残差块和注意力块，以减少模型大小和乘加操作数（MACs）。
• 采用特征级知识蒸馏对剪枝后的模型进行微调，将原始大模型的知识迁移至剪枝模型以维持性能。
• 蒸馏过程仅使用0.22M个LAION配对数据和13个A100天的训练计算资源，显著降低微调成本。
• 该方法适用于SDM-v1.4和SDM-v2.1-base两种模型版本，并在不同模型变体上均取得一致的性能提升。
• 在零样本文本到图像生成、个性化生成、图像到图像转换以及边缘设备部署等方面对剪枝并蒸馏后的模型进行评估。
• 在条件扩散模型（用于人脸生成的CelebA-HQ）上验证该方法，确认其在文本到图像任务之外也具备良好的泛化能力。

实验结果

研究问题

• RQ1 在Stable Diffusion的U-Net架构中剪枝残差块和注意力块，能否在不降低性能的前提下显著减少模型大小和推理延迟？
• RQ2 在数据和计算资源极少的条件下，知识蒸馏能否有效将大预训练扩散模型的能力迁移至紧凑的剪枝模型中？
• RQ3 轻量化扩散模型在多大程度上能实现与数十亿参数模型相媲美的零样本文本到图像生成性能？
• RQ4 剪枝并蒸馏后的模型能否在降低微调成本的前提下，有效支持个性化文本到图像生成与图像到图像转换？
• RQ5 在边缘设备上实现具有实时推理速度的紧凑模型部署是否可行？

主要发现

• BK-SDM在CPU和GPU上均实现最高51%的模型大小缩减和43%的延迟降低，MACs与推理时间减少30–50%。
• 仅使用13个A100天和0.22M个LAION配对数据，蒸馏模型在MS-COCO 30K数据集上FID达到14.61，性能媲美使用海量资源训练的大模型。
• BK-SDM-Small在Jetson AGX Orin和iPhone 14上均能在4秒内生成512×512分辨率图像，充分证明其在边缘设备上的部署可行性。
• 使用DreamBooth进行个性化生成时，BK-SDM模型在保持95–99%原始SDM模型的主体与提示保真度的同时，显著降低了微调成本。
• 该方法可推广至其他扩散模型，例如成功压缩了参数量为1.87亿的条件LDM用于CelebA-HQ人脸生成，实现高质量生成效果。
• 即使参数量和训练数据均大幅减少，BK-SDM模型仍保持具有竞争力的CLIP和IS分数，表明其具备强大的生成质量。