35张图，直观理解Stable Diffusion

公众号新闻

2023-02-02 13:02

来源 | OSCHINA 社区 - OneFlow深度学习框架

原文链接：https://my.oschina.net/oneflow/blog/6087116

作者｜Jay Alammar

翻译｜杨婷、徐佳渝

最近，AI图像生成引人注目，它能够根据文字描述生成精美图像，这极大地改变了人们的图像创作方式。Stable Diffusion作为一款高性能模型，它生成的图像质量更高、运行速度更快、消耗的资源以及内存占用更小，是AI图像生成领域的里程碑。

在接触了AI图像生成以后，你可能会好奇这些模型背后的工作原理。

下面是对Stable Diffusion工作原理的概述。

Stable Diffusion用途多样，是一款多功能模型。首先它可以根据文本生成图像(text2img)。上图是从文本输入到图像生成的示例。除此之外，我们还可以使用Stable Diffusion来替换、更改图像（这时我们需要同时输入文本和图像）。

下面是Stable Diffusion的内部结构，了解内部结构可以让我们更好地理解Stable Diffusion的组成、各组成部分的交互方式、以及各种图像生成选项/参数的含义。

1
Stable Diffusion的组成

Stable Diffusion并不是一个单一模型，而是由多个部分和模型一起构成的系统。

从内部来看，首先我们可以看到一个文本理解组件，这个组件将文本信息转化为数字表示（numeric representation）以捕捉文本意图。

这部分主要对ML进行大概介绍，文章后续还会讲解更多细节。可以说这个文本理解组件（文本编码器）是一个特殊的Transformer语言模型（严格来说它是一个CLIP模型的文本编码器）。将文本输入到 Clip 文本编码器得到特征列表，对于文本中的每一个word/token 都有会得到一个向量特征。

然后将文本特征作为图像生成器的输入，图像生成器又由几部分组成。

图像生成器两步骤：

1-图像信息创建器（Image information creator）

图像信息创建器是Stable Diffusion特有的关键部分，也是其性能远超其他模型的原因。

图像信息创建器运行多个step生成图像信息。Stable Diffusion接口（interfaces）和库（libraries）的step参数一般默认为50或100。

图像信息创建器完全在图像信息空间（亦称潜在空间）上运行，这让Stable Diffusion比以前在像素空间（pixel space）上运行的扩散模型速度更快。从技术上讲，图像信息创建器由UNet神经网络和调度算法组成。

“扩散”一词描述了图像信息创建器中发生的事情。因为图像信息创建器对信息作了逐步处理，所以图像解码器（image decoder）才能随后产出高质量图像。

2-图像解码器（Image Decoder）

图像解码器根据图像信息创建器的信息绘制图像，它只用在过程结束时运行一次，以生成最终的像素图像。

这样就构成了Stable Diffusion的三个主要组成部分，每个部分都有自己的神经网络：

ClipText：用于文本编码。输入: 文本。输出: 77个token embeddings向量，每个向量有768维。
UNet+调度程序：在信息（潜在）空间中逐步处理信息。输入: 文本embeddings和一个初始化的多维数组（结构化的数字列表，也称为张量）组成的噪声。输出：经过处理的信息数组。
自动编码解码器（Autoencoder Decoder）：使用经过处理的信息数组绘制最终图像。输入：经过处理的信息数组（维数：(4,64,64)）输出: 生成的图像（维数：（3,512,512），即（红/绿/蓝；宽，高）)。

2
到底何为扩散（Diffusion）？

扩散是发生在粉色区域图像信息创建器组件中的过程。这一部分有一个表示输入文本的token embeddings和一个随机初始化的图像信息数组，这些数组也被称为latents，在这个过程中会产生一个信息数组，图像解码器（Image Decoder）使用这个信息数组生成最终图像。

扩散是循序渐进逐步发生的，每一步都会增加更多的相关信息。为了更加直观地了解这个过程，我们可以检查随机latents数组，看它是否转化为了视觉噪音（visual noise）。在这种情况下，视觉检查（Visual inspection）是通过图像解码器进行的。

扩散分为多个step，每个step都在输入的latents数组上运行，并且会产生另一个latents数组，这个数组更类似于输入文本以及模型在模型训练时的所有图像中获取的所有视觉信息。

我们可以对一组这样的latents数组执行可视化，看看每一步都添加了什么信息。这一过程令人叹为观止。

视频链接：

在这种情况下，步骤2和4之间发生了一些特别有意思的事情，就好像轮廓是从噪音中浮现出来的。

3
Diffusion的工作原理

扩散模型图像生成的核心是强大的计算机视觉模型。在足够大的数据集的基础上，这些模型可以学会很多复杂运算。扩散模型通过如下方式建构问题来实现图像生成：

假设我们有一个图像，我们首先生成一些噪音（noise），然后将这些噪音添加到图像上。

我们可以将这看成是一个训练示例。之后我们使用同样的公式去创建更多的训练示例，然后用这些示例去训练图像生成模型的中心组件。

虽然这个例子展示了从图像（总量0，没有噪音）到总噪音（总量4，总噪音）的一些噪音值，但是我们可以轻松控制向图像中添加的噪音，因此我们可以将其分为数十个step，为数据集中的每个图像创建数十个训练示例。

有了这个数据集，我们可以训练噪音预测器（noise predictor），并最终得到一个在特定配置下运行时可以创建图像的预测器。接触过ML的人会觉得训练步骤非常熟悉：

接下来我们来看看Stable Diffusion是如何生成图像的。

4
通过降噪绘图

经过训练的噪音预测器可以对噪音图像进行降噪处理，并且可以预测噪音。

因为样本噪音（sampled noise）被预测，所以如果我们从图像中去掉这个样本，我们得到的图像就会更接近模型训练的图像。（这个图像不是确切的图像本身，而是图像分布，也就是图像的像素排列，在像素排列中天空通常是蓝色的，高于地面，人有两只眼睛，猫有尖耳朵并且总是懒洋洋的）。

如果训练数据集中的图像比较美观，比如说Stable Diffusion训练的LAION Aesthetics，那么训练出来的图像的可观赏性也会更高。如果我们在logo图像上对其进行训练，那么我们最终会得到一个logo生成模型。

这里总结了扩散模型处理图像生成的过程，主要如论文Denoising Diffusion Probabilistic Models所述。相信你对扩散的含义有了一定的了解，知道了Stable Diffusion、Dall-E 2和谷歌Imagen的主要组件。

值得注意的是，到目前为止我们所描述的扩散过程，没有使用任何文本数据，只需运行模型就能生成精美图像。不过我们无法控制图像的内容，它可能是一座金字塔，也可能是一只猫。接下来，我们将讨论如何将文本信息融入扩散过程以控制图片类型。

5
速度提升：在压缩（Latent）数据中扩散

为了加快图像生成过程，Stable Diffusion论文没有在像素图像上进行运行，而是在图像的压缩版本上运行。论文将这称为前往潜在空间(Departure to Latent Space)。

压缩（随后是解压缩/绘图）通过编码器完成。自动编码器使用Image Encoder将图像压缩进潜空间，然后使用Image Decoder再对压缩信息进行重构。

正向扩散在潜空间上完成。噪声信息应用于潜空间，而不是应用于像素图象。因此，训练噪声预测器（noise predictor）实际上是为了预测压缩表示（compressed representation）上的噪音，这个压缩表示也被称为潜空间（latent space）。

正向扩散是使用Image Encoder生成图像数据，来训练噪声预测器。训练一旦完成，就可以执行反向扩散，使用Image Decoder生成图像。

LDM/Stable Diffusion论文的图3中提及了这两个过程：

上图还显示了“conditioning”组件，这个组件在本例中是描述模型生成图像的文本提示词（text prompts）。接下来，我们继续探讨文本组件。

6
文本编码器：一种Transformer语言模型

Transformer语言模型作为语言理解组件，能够接受文本提示词，生成token embeddings。Stable Diffusion模型使用的是ClipText（基于GPT的模型），而论文中采用的是BERT。

Imagen论文表明，语言模型的选择相当重要。相较于较大的图像生成组件，较大的语言模型组件对生成图像的质量影响更大。

较大的/更好的语言模型对图像生成模型的质量有巨大的影响。资料来源：Saharia等人所著论文Google Imagen中的图A.5。

早期的Stable Diffusion模型仅使用了OpenAI发布的预训练模型ClipText。未来模型可能转向新发布的更大的CLIP变体OpenCLIP。（更新于2022 年 11 月，详情见Stable Diffusion V2 uses OpenClip。与仅含有630万文本模型参数的ClipText相比，OpenCLIP文本模型参数多达3.54亿。）

7
如何训练CLIP

CLIP模型是在图像和图像说明数据集上训练的。我们可以设想这样一个数据集，它里面有4亿张图像以及这些图像说明的材料。

图像及图像说明数据集

实际上，CLIP是在网络上抓取的带有“alt”标签的图像上训练的。CLIP是图像编码器和文本编码器的结合。简单来说，训练CLIP就是分别对图像和图像文本说明进行编码。

然后，使用余弦相似度来比较生成的embeddings。刚开始训练时，即使文本正确描述了图像，相似度也会很低。

我们更新了这两个模型，这样下次嵌入它们时就可以得到相似的embeddings。

通过在数据集上重复此操作并使用大的batch size，最终使编码器能够生成图像和文本说明相似的embeddings。如word2vec，训练过程也需要包含不匹配的图像和文本说明作为负样本，以得到较低的相似度分数。

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

为了使文本融入图像生成，我们须调整噪声预测器来输入文本。

现在，在数据集中加入文本。因为我们是在潜空间中运行，所以输入的图像和预测的噪声都处于潜空间中。

为了更好地理解UNet中文本tokens的使用方式，下面我们将进一步探究UNet模型。

Unet噪声预测器的Layers(未使用文本)

首先来看没有使用文本的UNet，其输入和输出如下：

可以看到：

UNet是一系列用于转换latents数组的layers
每一layer都对前一个layer的输出进行操作
Some of the outputs are fed (via residual connections) into the processing later in the network
通过残差连接（residual connections），将网络前面的layer输出送入到后面的layer进行处理
时间步长被转化为embedding向量，在网络层中使用