近年来，人工智能社区目睹了更大、更高性能的语言模型的显著增长，如Falcon 40B、LLaMa-2 70B、Falcon 40B、MPT 30B，以及像SD2.1和SDXL这样的图像领域模型。这些进展无疑推动了人工智能可实现的界限，实现了高度多功能和先进的图像生成和语言理解能力。然而，当我们对这些模型的强大和复杂性感到惊叹时，我们必须认识到越来越需要将人工智能模型变得更小、更高效、更易于使用，特别是通过开源的方式。
在Segmind，我们一直在努力提高生成式人工智能模型的速度和成本效益。去年，我们开源了我们的加速SD-WebUI库，名为voltaML，它是一个基于AITemplate/TensorRT的推理加速库，推理速度提高了4-6倍。为了实现使生成模型更快、更小、更便宜的目标，我们正在开源我们的压缩SD模型的权重和训练代码；SD-Small和SD-Tiny。预训练的检查点可以在Huggingface上找到🤗
知识蒸馏
我们的新压缩模型是使用知识蒸馏（KD）技术进行训练的，这项工作主要基于这篇论文。作者描述了一种块去除知识蒸馏方法，其中移除了一些UNet层，并训练了学生模型的权重。使用论文中描述的KD方法，我们能够训练出两个压缩模型，分别使用了🧨 diffusers库；Small和Tiny，与基础模型相比，它们的参数数量分别减少了35%和55%，同时实现了与基础模型相当的图像保真度。我们在这个存储库中开源了我们的蒸馏代码，并在Huggingface上提供了预训练的检查点🤗。
知识蒸馏训练神经网络类似于教师逐步指导学生。首先，大型教师模型在大量数据上进行预训练，然后在较小的数据集上训练较小的模型，以模仿大型模型的输出，同时还进行传统的数据集训练。
在这种特定类型的知识蒸馏中，学生模型被训练来执行从纯噪声中恢复图像的正常扩散任务，但同时，模型被要求与较大的教师模型的输出匹配。匹配的输出发生在U-Net的每个块中，因此模型质量大多保留。因此，使用前面的类比，我们可以说在这种蒸馏过程中，学生不仅尝试从问题和答案中学习，还从教师的答案以及逐步获得答案的方法中学习。我们在损失函数中有3个组成部分来实现这一点，首先是目标图像的潜在特征和生成图像的潜在特征之间的传统损失。其次是教师生成的图像的潜在特征和学生生成的图像的潜在特征之间的损失。最后，也是最重要的组成部分，是特征级别的损失，即教师和学生每个块的输出之间的损失。
所有这些组合在一起构成了知识蒸馏训练。下面是论文中描述的用于KD的块去除UNet的架构。
图片来源于Shinkook等人的论文“On Architectural Compression of Text-to-Image Diffusion Models”
我们选择了Realistic-Vision 4.0作为基础教师模型，并使用LAION艺术美学数据集进行训练，该数据集中的图像评分高于7.5，因为其具有高质量的图像描述。与论文不同，我们选择对Small模型进行100K步的1M图像训练，对Tiny模型进行125K步的1M图像训练。蒸馏训练的代码可以在这里找到。
模型使用
可以使用🧨 diffusers的DiffusionPipeline来使用模型
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained("segmind/small-sd", torch_dtype=torch.float16)
prompt = "漂亮女孩的肖像"
negative_prompt = "（变形的虹膜，变形的瞳孔，半逼真的，CGI，3D，渲染，素描，卡通，绘画，动漫：1.4），文本，特写，裁剪，超出画面，最差质量，低质量，JPEG伪影，丑陋，重复，恐怖，残缺，额外的手指，突变手，画得不好的手，画得不好的脸，变异，变形，模糊，脱水，解剖不良，比例不良，额外的肢体，克隆脸，畸形，比例不协调，畸形肢体，缺臂，缺腿，额外的手臂，额外的腿，融合的手指，太多的手指，长脖子"
image = pipeline(prompt, negative_prompt = negative_prompt).images[0]
image.save("my_image.png")推理延迟方面的速度
我们观察到，精简模型比原始基础模型快100%。基准测试代码可以在这里找到。
潜在限制
精简模型处于早期阶段，输出可能还不具备生产质量。这些模型可能不是最好的通用模型。它们最适合用于针对特定概念/风格进行微调或LoRA训练。精简模型在组合性或多概念方面尚不太好。
在人像数据集上对SD-tiny模型进行微调
我们已经对我们的sd-tiny模型进行了对由Realistic Vision v4.0模型生成的人像图像进行微调。以下是使用的微调参数。
- 步骤：131000
- 学习率：1e-4
- 批大小：32
- 梯度累积步数：4
- 图像分辨率：768
- 数据集大小-7k张图像
- 混合精度：fp16
我们能够产生与原始模型产生的图像质量接近的图像，参数减少了近40%，下面的示例结果可以说明问题：
微调基础模型的代码可以在这里找到。
LoRA训练
在精简模型上进行LoRA训练的一个优势是训练速度更快。以下是我们在一些抽象概念上对经过精简模型进行的第一个LoRA训练的一些图像。LoRA训练的代码可以在这里找到。
结论
我们邀请开源社区帮助我们改进和实现对这些精简SD模型的更广泛采用。用户可以加入我们的Discord服务器，我们将在那里宣布这些模型的最新更新，发布更多的检查点和一些令人兴奋的新LoRA。如果您喜欢我们的工作，请给我们的Github点个赞。