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. We hope this work can help build small yet powerful diffusion models with feasible training budgets. Code and models can be found at: https://github.com/Nota-NetsPresso/BK-SDM

翻译：基于Stable Diffusion模型（SDMs）的文本到图像生成因十亿级参数量而面临高昂计算需求。为提升效率，近期研究在保留原始架构的前提下，通过减少采样步数及应用网络量化来优化。架构精简尝试的匮乏，可能源于对这类大规模模型昂贵的重训成本的担忧。本文揭示了块剪枝与特征蒸馏在低成本通用文本到图像生成中的惊人潜力。通过移除SDMs中U-Net的若干残差块与注意力块，我们实现了模型尺寸、MACs与延迟的30%~50%缩减。研究表明，即便在有限资源下（仅需13个A100 GPU天与微小数据集），蒸馏重训仍能有效模仿原始SDMs（v1.4与v2.1-base版本需超过6,000个A100 GPU天）。借助迁移知识，我们的BK-SDMs在零样本MS-COCO任务上展现出与更大规模十亿级参数模型相竞争的性能。我们进一步验证了轻量级骨干网络在个性化生成与图像到图像翻译中的适用性。所提模型在边缘设备上部署可实现4秒推理。期望本工作能助力以可行训练预算构建小型高性能扩散模型。代码与模型开源地址：https://github.com/Nota-NetsPresso/BK-SDM