BK-SDM：轻量级 Stable Diffusion 的知识蒸馏之路

一、一个让显卡"喘口气"的梦想

想象一下：你在一个预算有限的工作站上，想要生成一张精美的AI画作，却因为模型太大，显存爆满、程序崩溃。这种痛苦每一个研究Stable Diffusion的人都不陌生——原始SD-v1.4的U-Net模型参数量高达8.6亿，整个SD系统超过10亿参数，普通人根本玩不起。

2023年，韩国AI公司Nota（现NetsPresso）的团队决定做一件事：把Stable Diffusion"瘦"下来。他们提出了BK-SDM（Block-removed Knowledge-distilled Stable Diffusion），通过结构化剪枝与知识蒸馏的结合，在显著压缩模型的同时保持生成质量。这项成果最终发表于计算机视觉顶会ECCV 2024，并在ICCV 2023 Demo Track和ICML 2023 ES-FoMo Workshop上做了早期演示。

图1：BK-SDM 对 U-Net 架构的渐进式压缩策略（来源：论文官方MODEL_CARD）

二、核心技术：从"大卸八块"到"精准蒸馏"

2.1 结构化剪枝：给U-Net做"抽脂手术"

Stable Diffusion的核心是U-Net，它负责去噪过程。BK-SDM团队发现，U-Net中的残差块（Residual Blocks）和注意力块（Attention Blocks）存在大量冗余——某些块去掉后，模型仍然能"凑合"工作。关键问题是如何优雅地裁剪这些块，同时保证输出质量。

BK-SDM设计了三种渐进式剪枝策略：

• BK-SDM-Base（0.76B参数）：移除U-Net外围阶段的一些残差块和注意力块，U-Net参数量从0.86B降至0.58B
• BK-SDM-Small（0.66B参数）：在Base基础上再移除中间阶段的块，U-Net降至0.49B
• BK-SDM-Tiny（0.50B参数）：在最内层也做裁剪，U-Net仅0.33B

这种"层层剥皮"的策略，让开发者可以根据硬件条件灵活选择模型大小。

2.2 知识蒸馏：让学生模仿老师

剪枝后的学生网络需要恢复能力。BK-SDM采用多层蒸馏策略：

图2：BK-SDM 的知识蒸馏预训练流程（来源：论文官方MODEL_CARD）

1. 特征级蒸馏：将教师网络（原始SD）的中间层特征作为监督信号，让学生网络学习逼近
2. 输出级蒸馏：让学生网络的输出分布逼近教师网络的输出
3. 去噪任务损失：保留原始SD的去噪损失函数，确保基本能力不退化

令人惊讶的是，训练数据量非常小——仅用21.3万对（LAION-Aesthetics V2 6.5+）图像-文本对，在单张A100 80GB GPU上训练5万步，就能得到一个可用的压缩模型。这意味着普通研究机构也能复现。

三、性能表现：体积小了，质量如何？

MS-COCO Benchmark（Zero-shot，512×512生成后缩至256×256）

BK-SDM-Tiny的参数量仅为原始SD的48%，但FID指标差距仅约2个点。对于iPhone 14上的Core ML部署场景（苹果官方4秒出图），这种精度损失完全可接受。

四、代码架构：站在HuggingFace的肩上

BK-SDM的代码库完全基于HuggingFace Diffusers构建，核心依赖为：

torch==2.0.1
transformers==4.27.4
diffusers==0.15.0
open_clip_torch==2.18.0
accelerate==0.18.0
peft==0.2.0

主要模块分工清晰：

• kd_train_text_to_image.py：知识蒸馏训练核心，包含蒸馏损失计算和权重迁移逻辑
• count_macs_params.py：用thop库计算模型参数量和MACs，供对比分析
• generate.py / generate_single_prompt.py：推理脚本，调用自定义的InferencePipeline
• dreambooth_finetune.py：支持DreamBooth微调，使用PEFT（LoRA）做轻量微调
• eval_clip_score.py：用OpenCLIP评估CLIP分数（ViT-g/14）
• utils/inference_pipeline.py：封装Diffusers Pipeline，提供统一的推理接口

权重发布在HuggingFace Hub：nota-ai/bk-sdm-base、nota-ai/bk-sdm-small、nota-ai/bk-sdm-tiny等系列。

五、局限与思考

1. GPU依赖：蒸馏训练阶段至少需要16GB显存的GPU（V100/A100），普通消费级显卡难以承受
2. 生成质量差距：压缩模型与原版SD在细节纹理上仍有可见差距，尤其在复杂提示词场景
3. 版本锁定：依赖diffusers==0.15.0等旧版本库，新版Diffusers可能有API不兼容问题
4. 数据局限：蒸馏数据量仅21万对，相比原SD训练数据（数十亿对）规模极小

六、行业意义：让AI画图"飞入寻常百姓家"

BK-SDM的意义不仅在于论文本身，更在于它代表了一个方向——"大模型压缩以适配边缘设备"。从苹果Core ML4秒推理、到移动端部署、再到嵌入式设备，这条路正在被BK-SDM、LCM（Latent Consistency Models）等技术逐步打通。Nota公司后续还推出了EdgeFusion（CVPR 2024 Workshop），将BK-SDM-Tiny与LCM结合，进一步加速推理。

对于普通AI爱好者而言，BK-SDM-Small是一个很好的切入点：既能在消费级GPU上训练，又能以可接受的推理速度生成高质量图片。上手门槛比原始SD低很多，值得一试。