[翻译] 图解Stable Diffusion

[翻译] 图解Stable Diffusion

原文：https://jalammar.github.io/illustrated-stable-diffusion/

作者：https://jalammar.github.io/

AI 图像生成是最近最火出圈的 AI 应用。它能够根据文本描述产生令人震撼的、质量极高的图像，人类创作艺术的方式正在发生改变。Stable Diffusion的发布是最近 AI 绘画领域的一个里程碑式的事件。它生成的图像效果极佳，除此之外，它的速度很快，并且对硬件资源的要求相对较低。它的出现使得普通人也能使用高性能的图像生成模型。

在体验过 AI 图像生成的强大能力之后，你或许会好奇它是怎样工作的。本文是对 Stable Diffusion 工作原理的简单介绍。

Stable Diffusion 非常全能，它有许多种使用方式。本文中我们仅聚焦于文本生成图像（text2img）。上图展示了一个根据文本输入，得到对应生成图像的例子。除了文本生成图像之外，Stable Diffusion 另一个常见的用法是修改图像（此时输入为文本+图像）。

接下来让我们开始深入了解 Stable Diffusion，理解组件、它们如何交互以及图像生成选项/参数的含义。

Stable Diffusion组件

Stable Diffusion 是一个由多个组件和模型组成的系统， 而非一个整体的模型。我们首先从高层形象地理解 Stable Diffusion，本文将后面介绍更多具体的机器学习技术细节。

当我们深入观察时，首先看到的是一个文本理解（text-understanding）组件，它捕捉文本中的意图，将文本信息转换为模型能够理解的数值表示。

这个文本编码器是一种特殊的 Transformer 语言模型（具体技术上：CLIP 模型的文本编码器）。 它获取输入文本并输出代表文本中每个单词/token 的数值表示（每个 token 由一个向量表示）。该信息随后给图像生成器（Image Generator），图像生成器本身也由多个组件组成。

图像生成器由以下两个阶段组成：

1 图像信息生成器

图像信息生成器（Image Information Creator）是 Stable Diffusion 成功的秘诀，是它性能和效率高于之前工作的原因。

该组件需要运行多步来生成图像信息。步数就是 Stable Diffusion 界面或库中的steps 参数，通常设为 50 或 100。

图像信息生成器完全在图像信息空间（或者称为潜层空间 latent space）中进行工作。本文之后会探讨这样做的原因和好处。这使得 Stable Diffusion 相比于之前的工作在像素空间（pixel space）的扩散模型更快。技术上来讲，图像信息生成器由一个 UNet 神经网络和调度算法组成。

“扩散（diffusion）”这个词描述的就是该组件的行为。该组件通过一步一步地对信息进行处理，从而得到最终的高质量图像（由接下来的图像解码器组件生成）。

2 图像解码器

图像解码器（Image Decoder）根据图像信息生成器生成的信息画出图像。不同于多步运行的信息生成器，图像解码器仅运行一次，来生成最终的像素级图像。

从上图中，我们可以看到 Stable Diffusion 的三个主要组件，它们各自由不同的神经网络组成：

• ClipText 用于文本编码
  • 输入：文本
  • 输出：77 个 token 嵌入向量，每个向量 768 维
• UNet + Scheduler 用于在潜层空间中逐步地地处理（或者说扩散）信息
  • 输入：文本嵌入和一个高维噪声张量
  • 输出：经过处理得到的信息张量
• AutoEncoder Decoder 根据信息张量画出图像
  • 输入：信息张量（维度：(4, 64, 64)）
  • 输出：图像（维度：(3, 512, 512)）

扩散模型

究竟什么是扩散？扩散是发生在粉红色图像信息生成器组件内部的过程。 该组件的输入为用于表示输入文本信息的 token 嵌入，和一个起始的随机噪声图像信息张量，生成一个信息张量，图像解码器使用该信息张量绘制最终图像。

这个过程以多步的形式进行。每一步添加更多的相关信息。为了直观地理解整个过程，我们可以将随机潜层张量（latent）传递给视觉解码器，看它是否转换为随机视觉噪声。

扩散过程有多步，每步操作一个输入潜层张量，并生成一个新的潜层张量。新的张量更好地集成了输入文本和视觉信息，其中视觉信息来自模型训练集中的图像。

如下图所示，我们可以对这些潜层向量进行可视化，来看看每一步添加了哪些信息。

整个过程看起来很惊人，步骤2和步骤4之间发生了一些特别有趣的事情。就好像轮廓是从噪音中浮现出来的。

扩散模型如何工作？

计算机视觉模型具有强大的能力， 给定足够大的数据集，这些模型可以学习复杂的操作。 扩散模型通过将问题框定如下来处理图像生成：

即我们有一张图像，我们生成一些噪声，并将它添加到图像上。

添加噪声之后，就可以作为训练样本。同样的，我们可以使用相同的方式创造很多个训练样本，来训练我们 Stable Diffusion 系统的核心组件。

这些样本具有不同的噪声量，如图中展示了从噪声量 0 （无噪声）到噪声量 3 （完全噪声）。我们可以控制给每张图像添加多少噪声量，从而我们可以将其分散到数十个步骤中，为训练数据集中的所有图像为每个图像创建数十个训练样本。

有上图所示这样一个数据集，我们就可以训练一个噪声预测器（noise predicator），经过训练后，它就可以在特定的配置下生成图像。训练的方式如下图所示，就是一个很常见的机器学习训练过程

现在让我们来看看如何生成图像。

通过去除噪声生成图像

训练好的噪声预测器可以接收一张噪声图像，和去噪步数值，预测出噪声。

预测出噪声之后，我们从图像中减去噪声，就可以得到与训练集图像接近的图像。注意，不是真实世界图像本身，而是图像的分布 - 即像素排列的世界，天空通常是蓝色的，在地面之上；人有两只眼睛；猫有尖尖的耳朵。

如果训练数据集都是美学质量较高的图像（如 Stable Diffusion 训练所使用的 LAION Aesthetics），那么生成的图像就会有较高的美学质量。如果我们在 logo 数据集上进行训练，就会得到一个 logo 生成模型。

以上归纳了扩散模型是如何生成图像的，就如 Denoising Diffusion Probabilistic Models 描述的那样。现在你已经从直觉上了解了扩散模型，它不仅是 Stable Diffusion 的主要组件，也是 Dall-E 2 和 Imagen 的主要组件。

注意我们目前所描述的图像生成没有用到任何文本数据，所以如果部署我们目前所训练的模型，它能够产生好看的图像，但我们无法控制生成的图像内容是金字塔，还是一只猫，还是其他的。下一节我们会介绍如何将文本信息集成到整个过程中，从而控制生成图像的内容。

加速：在压缩（潜层）数据而非像素图像中扩散

为了加速图像生成的过程，Stable Diffusion 的扩散过程并非运行在像素图像上，而是压缩版本的图像。原论文 将次称为 “Departure to Latent Space”。

压缩（和之后的解压缩/绘制过程）是由一个自编码器（AutoEncoder）完成。自编码器由编码器和解码器组成，其中编码器将图像压缩至潜层空间，解码器根据压缩信息绘制出图像。

前向扩散过程（加噪声）在压缩潜层空间完成。噪声是添加到这些潜层张量中，而非像素图像中。噪声预测器实际上也是预测潜层张量的噪声。

前向处理（使用编码器）是生成加噪声的数据来训练噪声预测器。编码器训练完成后，我们可以通过反向过程（reverse process）生成图像（使用解码器）。

这两个流程在原论文图三中所示：

该图还显示了 “conditioning” 组件，在本例中是描述模型应生成的图像的文本提示。 因此，让我们深入研究文本组件。

文本编码器

基于 Transformer 的语言模型用作语言理解组件。它接收文本作为输入，输出 token 嵌入。释出的 Stable Diffusion 模型使用了 ClipText （是一个基于 GPT 的模型），而论文中使用的是 BERT。

Imagen 论文中的实验显示，语言理解模型的选择对于最终生成图像的质量有很大的影响。使用更大的语言理解组件，相比于使用更大的图像生成组件，前者对生成图像质量的提升更为显著。

之前的 Stable Diffusion 系统直接使用 OpenAI 发布的预训练 ClipText 模型，作为语言理解组件。Stable Diffusion v2 使用 OpenCLIP 中更大的语言模型，参数量从原来的 63M 提升至 354 M。

CLIP 是如何训练的

CLIP 在图像文本对数据集上进行训练。OpenAI CLIP 原论文使用了包含 400M 图像及其对应文本的训练数据。

实际上，CLIP 的训练集是通过爬取网络上带有 “alt” 标签的图像收集得到的。

CLIP 由一个文本编码器和一个图像编码器组成。他的训练过程分为三步，首先接收一张图像及其对应的文本，并分别使用图像编码器和文本编码器对其进行编码，得到图像嵌入和文本嵌入。

第二步，我们使用余弦相似度来对比图像嵌入和文本嵌入。即使图像和文本嵌入来自同一个图像文本对，也没有很高的相似度分数。

但是我们希望来自同一个图像文本对的两种嵌入有较高的相似度。第三步，我们会计算对比损失，反向传播，更新图像编码器和文本编码器。再下一次，编码器提取的两种嵌入会有更高的相似度。

在大的数据集上通过不断的重复上述过程，并使用较大的批尺寸，我们最终得到的编码器能够对一张小狗的图像和“一张小狗的图片”的文本生成相似的嵌入。与 word2vec 类似，CLIP 的训练过程需要包括负样本，负样本中的图像和文本来自不同的图像文本对，模型需要为负样本给出较低的相似度。这种训练方式称为对比学习。

将文本信息送入图像生成过程

为了将文本信息加入到图像生成过程中，我们的噪声预测器的输入需要加入文本。

我们现在训练用的数据集又多加了编码过的文本信息（下图中的蓝色方块）。

为了更好地理解 UNet 中的文本 token 是如何使用的，我们接下来稍微详细点看一下 UNet 的内部。

不含文本时的Unet噪声预测器层

我们首先看一下不含文本信息时的扩散模型 Unet，其输入输出如下图所示：

在其内部（如下图所示），我们可以看到：

• Unet 由一系列层组成，这些层实际上是残差块，用于处理潜层张量
• 每一层处理上一层的输出
• 某些输出通过残差连接送入到后面的处理网络
• 时间步（timestep，对应加入的噪声量）转换为一个嵌入向量，输入到各层中

含文本时的Unet噪声预测器层

现在我们来看一下含文本信息的 Unet 网络。

为了支持文本输入（技术术语：text conditioning），我们需要在各个残差块之间添加注意力层（attention layer）。

注意，残差块并不直接将文本作为输入。而是由注意力层整合张量中的文本信息，然后下一个残差块在进行处理时，可以利用这些文本信息。

结论

希望本文能够带给读者 Stable Diffusion 运作原理的直觉理解。涉及许多其他概念，但相信一旦您熟悉了上面的构建块，它们就会更容易理解。下列资源很适合用于进一步了解 Stable Diffusion

资源

• I have a one-minute YouTube short on using Dream Studio to generate images with Stable Diffusion.
• Stable Diffusion with 🧨 Diffusers
• The Annotated Diffusion Model
• How does Stable Diffusion work? – Latent Diffusion Models EXPLAINED [Video]
• Stable Diffusion - What, Why, How? [Video]
• High-Resolution Image Synthesis with Latent Diffusion Models [The Stable Diffusion paper]
• For a more in-depth look at the algorithms and math, see Lilian Weng’s What are Diffusion Models?
• Watch the great Stable Diffusion videos from fast.ai