关于人工智能:开源-SDSmall-和-SDTiny-知识蒸馏代码与权重

<p align=”center”>
<img src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202308251921008.png” width=500>
</p>

最近，人工智能社区在开发更大、更高性能的语言模型方面获得了显著的停顿，例如 Falcon 40B、LLaMa-2 70B、Falcon 40B、MPT 30B; 以及在图像畛域的模型，如 SD2.1 和 SDXL。这些提高无疑推动了人工智能的倒退，使其具备高度多功能和最先进的图像生成和语言理解能力。然而，在咱们惊叹于这些模型的弱小和复杂性之余，必须意识到一个日益增长的需要: 使人工智能模型体量更小、运行更高效、更易于拜访，特地是通过开源它们来共建生态。

在 Segmind，咱们始终致力于如何使生成式 AI 更快、更便宜。去年，咱们开源了咱们减速的 SD-WebUI 库 voltaML，它是一个基于 AITemplate/TensorRT 的推理减速库，推理速度进步了 4-6 倍。为了持续实现使生成模型更快、更小、更便宜的指标，咱们正在开源咱们压缩的 SD 模型：SD-Small 和 SD-Tiny 的权重和训练代码。预训练的检查点可在 Hugging Face 🤗 上获取。

常识蒸馏

<p align=”center”>
<img src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202308251922874.png” width=500>
</p>

咱们的新压缩模型曾经通过常识蒸馏 (KD) 技术的训练，这项工作次要基于 这篇论文。作者形容了一种块移除常识蒸馏办法，其中一些 UNet 层被移除，学生模型权重被训练。应用论文中形容的 KD 办法，咱们可能应用 🧨 diffusers 库训练两个压缩模型; Small（渺小版本） 和 Tiny（极小版本），别离比根底模型少 35% 和 55% 的参数，同时实现与根底模型相当的图像保真度。咱们曾经在这个 repo 中开源了咱们的蒸馏代码，并将预训练检查点上传到了 Hugging Face 🤗。

常识蒸馏训练神经网络相似于老师一步一步领导学生。一个大的老师模型 (teacher model) 事后在大量数据上训练，而后一个较小的模型在较小的数据集上训练，以模拟大模型的输入并在数据集上进行经典训练。

在这种非凡类型的常识蒸馏中，学生模型被训练来实现从纯噪声复原图像的失常扩散工作，但同时，模型被迫与更大的老师模型的输入匹配。输入匹配产生在 U-nets 的每个块，因而模型品质根本放弃不变。所以，应用后面的类比，咱们能够说，在这种蒸馏过程中，学生不仅会试图从问题和答案中学习，还会从老师的答案以及逐渐失去答案的办法中学习。咱们在损失函数中有 3 个组成部分来实现这一点，首先是指标图像隐变量和生成图像隐变量之间的传统损失。其次是老师生成的图像隐变量和学生生成的图像隐变量之间的损失。最初，也是最重要的组成部分，是特色级损失，即老师和学生每个块输入之间的损失。

联合所有这些形成了常识蒸馏训练。上面是论文中形容的用于 KD 的块移除 UNet 架构。

<p align=”center”>

<img src="https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202308251922371.png" width=500>

</p>

图片来自 Shinkook 等人的 论文“On Architectural Compression of Text-to-Image Diffusion Models”。

咱们以 Realistic-Vision 4.0 为根底老师模型，并在 LAION Art Aesthetic 数据集 上训练，图像分数高于 7.5，因为它们具备高质量的图像形容。与论文不同，咱们抉择别离为 Small 和 Tiny 模式训练两个模型，别离在 1M 张图像上进行 100K 步和 125K 步的训练。蒸馏训练的代码能够在 这里 找到。

模型应用

模型能够通过 🧨 diffusers 中的 DiffusionPipeline 来应用。

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("segmind/small-sd", torch_dtype=torch.float16)
prompt = "Portrait of a pretty girl"
negative_prompt = "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck"
image = pipeline(prompt, negative_prompt = negative_prompt).images[0]
image.save("my_image.png")

推理提早方面的速度体现

咱们察看到，蒸馏模型比原始根底模型快了一倍。基准测试代码能够在 这里 找到。

<p align=”center”>
<img src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202308251922479.jpeg” width=500>
</p>

潜在的局限性

蒸馏模型处于晚期阶段，输入可能还不具备生产程度的品质。这些模型可能不是最好的通用模型，它们最好用作针对特定概念 / 格调进行微调或 LoRA 训练。蒸馏模型目前还不太善于组合性或多概念。

在人像数据集上微调 SD-tiny 模型

咱们曾经在 Realistic Vision v4.0 模型生成的人像图像上微调了咱们的 sd-tiny 模型。上面是应用的微调参数。

咱们可能产生靠近原始模型产生的图像品质，参数缩小了近 40%，上面的样本后果不言自明:

<p align=”center”>
<img src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202308251922835.png” width=500>
</p>

微调根底模型的代码能够在 这里 找到。

LoRA 训练

在蒸馏模型上进行 LoRA 训练的一个长处是训练更快。上面是咱们在蒸馏模型上对一些抽象概念进行的第一个 LoRA 训练的一些图像。LoRA 训练的代码能够在 这里 找到。

<p align=”center”>
<img src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202308251922486.png” width=500>
</p>

论断

咱们邀请开源社区帮忙咱们改良并实现这些蒸馏 SD 模型的更宽泛采纳。用户能够退出咱们的 Discord 服务器，在那里咱们将发表这些模型的最新更新，公布更多的检查点和一些令人兴奋的新 LoRAs。如果你喜爱咱们的工作，请在咱们的 Github 上点一下 star。

英文原文: https://hf.co/blog/sd_distillation

原文作者: Yatharth Gupta

译者: innovation64

审校 / 排版: zhongdongy (阿东)