BK-SDM：轻量、快速且低成本的Stable Diffusion版本
机器学习
2024-12-03 v4
摘要
使用Stable Diffusion模型（SDMs）的文本到图像（T2I）生成因十亿规模参数而涉及高计算需求。为提升效率，近期研究在保留原始架构的同时减少了采样步数并应用了网络量化。缺乏架构约简尝试可能源于对如此大规模模型昂贵重训练的担忧。本工作中，我们揭示了块剪枝与特征蒸馏在低成本的通用T2I中的惊人潜力。通过从SDMs的U-Net中移除若干残差和注意力块，我们实现了模型大小、MACs和延迟的30%~50%缩减。我们表明，即使在有限资源下蒸馏重训练也有效：仅使用13 A100天和微小数据集，我们的紧凑模型即可模仿原始SDMs（v1.4和v2.1-base，耗费超6000 A100天）。得益于迁移知识，我们的BK-SDMs在零样本MS-COCO上相较更大的多十亿参数模型给出了有竞争力的结果。我们进一步展示了轻量骨干网络在个性化生成和图像到图像翻译中的适用性。我们的模型在边缘设备上的部署实现了4秒推理。代码与模型见：https://github.com/Nota-NetsPresso/BK-SDM
引用
@article{arxiv.2305.15798,
title = {BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion},
author = {Bo-Kyeong Kim and Hyoung-Kyu Song and Thibault Castells and Shinkook Choi},
journal= {arXiv preprint arXiv:2305.15798},
year = {2024}
}
备注
ECCV 2024 Camera-Ready Version