当前位置:AIGC资讯 > AIGC > 正文

Diffusion Model, Stable Diffusion, Stable Diffusion XL 详解

文章目录

Diffusion Model 生成模型 DDPM概述 向前扩散过程 前向扩散的逐步过程 前向扩散的整体过程 反向去噪过程 网络结构 训练和推理过程 训练过程 推理过程 优化目标 详细数学推导 数学基础 向前扩散过程 反向去噪过程 Stable Diffusion 组成结构 运行流程 网络结构 变分自编码器 (VAE) 文本编码器 (CLIP Text Encoder) 噪音预测器 (Unet) 训练过程和损失函数 VAE训练 UNet训练 CLIP 训练 Stable Diffusion XL 基础架构 模型结构 Base 模型 VAE 模型 U-Net 模型 CLIP 模型 Refiner 模型 训练和推理 基本训练流程 训练技巧 1) 图像尺寸条件化 2) 图像裁剪参数条件化 3) 多尺度训练 4) 使用Offset Noise 5) SDXL的条件注入与训练细节 推理 参考

Diffusion Model

生成模型

常见的生成模型有:

GAN VAE Flow-based

与 VAE 或Flow-based 模型不同,扩散模型是通过固定的过程学习的,latent variable 具有高维度(与原始数据相同)。

文生图模型的基本架构如下:

文本编码器(Text Encoder)用于编码文本的输入 生成模型 (Generation Model)输入文字和杂讯,生成一个中间产物(这个中间产物可以是人看得懂的一个压缩图片,也可以是噪音) 解码器 (Decoder) 解码器,输入中间产物,得到最终的输出的图片。
通常而言,上述的三个模型是分开训练的,最后组合到一起。
例如 Stable Diffusion
DALL-E
Imagen

衡量图像生成质量的指标

FID
将生成的图片和真是的图片都输入到预训练好的CNN模型中,然后计算输出高斯分布的Frechet的距离。其值越小越好。
CLIP score
将生成的图片和输入的文本提示,分别送入Image encoder 和text encoder ,计算二者输出的距离。如果图像文本匹配,则距离越近越好。如果不匹配,则其距离越远越好。CLIP score的得分越大越好。

DDPM概述

基本思想
Denoising Diffusion Probabilistic Models (DDPM)
受非平衡统计物理启发[1],是通过一个迭代的向前扩散过程 ,系统地、缓慢地破坏数据分布中的结构。然后,我们学习一个逆扩散过程,将数据中的结构恢复,从而产生一个高度灵活且易于处理的数据生成模型。

💡 原文中作者对于基本思想的描述:“The essential idea, inspired by non-equilibrium statistical physics [1], is to systematically and slowly destroy structure in a data distribution* through an iterative forward diffusion process. We then learn a reverse diffusion process that restores structure in data, yielding a highly flexible and tractable generative model of the data.”

核心步骤

基于上述思想,Denoising diffusion 模型包括两个过程:

向前扩散过程 (Forward Diffusion): 逐渐向输入添加噪声 反向去噪过程 (Reverse Denoising): 学习通过去噪生成数据

向前扩散过程

前向扩散的逐步过程

首先,我们先来介绍一下如何破坏数据分布。

前向扩散过程中的输入和输出

如果我们取任何一幅图像(上图a),它具有某种非随机分布。我们不知道这个分布,但我们的目标是破坏它,我们可以通过向其添加噪声来实现。在这个过程的最后,我们应该得到类似于纯噪声的噪声(上图b)。

T0时刻: 将噪音线性添加到图片上,得到的新图片作为T1时刻的输入 T1时刻: 将噪音线性添加到T0时刻的输出,得到的新图片作为T2时刻的输入 Tn时刻: 在最后一个step,原始的信息被完全破坏,得到一个纯噪音的图片。

前向扩散过程的每一步被定义为 q ( x t ∣ x t − 1 ) = N ( x t , 1 − β t x t − 1 , β t I ) q(x_t|x_{t-1})= \mathcal{N}\left(x_t, \sqrt{1-\beta_t} x_{t-1}, \beta_t I\right) q(xt​∣xt−1​)=N(xt​,1−βt​ ​xt−1​,βt​I)
其中 q q q是前向扩散过程, x t x_t xt​是第 t t t步前向过程的输出(自然地, x t − 1 x_{t-1} xt−1​是第 t t t步的输入)。 N \mathcal{N} N是一个正态分布, 1 − β t x t − 1 \sqrt{1-\beta_t} x_{t-1} 1−βt​ ​xt−1​是均值, β t I \beta_t I βt​I定义了一个方差。

调度 (Schedule)
β t \beta_t βt​被称为调度(schedule),其值范围从0到1不等。通常将值保持在较低水平,以防止方差爆炸。

线性调度(Linear Schedule)

DDPM论文 [2] 使用了一个线性调度,因此输出看起来如下图:

使用线性调度(linear schedule)的前向扩散模型的1000个时间步

在线性调度的情况下, β t \beta_t βt​从0.0001线性增加到0.02, 均值和方差也随着时间步线性变化(如下图)。

均值和方差随着时间步的变化过程 余弦调度(Cosine Schedule)

OpenAI的研究人员在他们的2021年论文中 [3] 发现使用线性调度并不那么高效。因此提出了余弦调度(下图)。这种改进使得步骤数减少到50步。

使用余弦调度(cosine schedule)的前向扩散模型的50个时间步

前向扩散的整体过程

使用前向扩散过程向图像添加噪声将会很慢。训练过程不使用与前向过程一致的示例,而是使用来自任意时间步 t 的样本。这意味着在每个训练步骤中,我们需要通过 t 步来生成 1 个训练样本。
q ( x 1 : T ∣ x 0 ) : = ∏ t = 1 T q ( x t ∣ x t − 1 ) q\left(x_{1: T} \mid x_0\right):=\prod_{t=1}^T q\left(x_t \mid x_{t-1}\right) q(x1:T​∣x0​):=t=1∏T​q(xt​∣xt−1​)
接下来,我们对上式进行逐步分解:
q t ( q t − 1 ( q t − 2 ( q t − 3 ( ⋯ q 1 ( x 0 ) ) ) ) ) q_t\left(q_{t-1}\left(q_{t-2}\left(q_{t-3}\left(\cdots q_1\left(x_0\right)\right)\right)\right)\right) qt​(qt−1​(qt−2​(qt−3​(⋯q1​(x0​)))))
t 1 t_1 t1​时刻
q ( x 1 ∣ x 0 ) = N ( x 1 , 1 − β 1 x 0 , β 1 I ) q\left(x_1 \mid x_0\right)=\mathcal{N}\left(x_1, \sqrt{1-\beta_1} x_0, \beta_1 I\right) q(x1​∣x0​)=N(x1​,1−β1​ ​x0​,β1​I)
t 2 t_2 t2​时刻
q ( x 2 ∣ x 1 ) = N ( x 2 , 1 − β 2 x 1 , β 2 I ) q\left(x_2 \mid x_1\right)=\mathcal{N}\left(x_2, \sqrt{1-\beta_2} x_1, \beta_2 I\right) q(x2​∣x1​)=N(x2​,1−β2​ ​x1​,β2​I)
从上面的公式可以看到,只有均值是依赖于上一步的输出,方差对于任意一步都是已知的(因为 β \beta β 对于每一个时间步都是已知的)。
因此,引入下面两个记号:
α t = 1 − β t α ˉ t : = ∏ s = 1 t a s \begin{aligned} \alpha_t & =1-\beta_t \\ \bar{\alpha}_t & :=\prod_{s=1}^t a_s \end{aligned} αt​αˉt​​=1−βt​:=s=1∏t​as​​
代入上述记号,每一步的加噪公式变为:
q ( x 1 ∣ x 0 ) = N ( x 1 , α 1 x 0 , ( 1 − α 1 ) I ) q\left(x_1 \mid x_0\right)=\mathcal{N}\left(x_1, \sqrt{\alpha_1} x_0,\left(1-\alpha_1\right) I\right) q(x1​∣x0​)=N(x1​,α1​ ​x0​,(1−α1​)I)
q ( x 2 ∣ x 1 ) = N ( x 2 , α 2 x 1 , ( 1 − α 2 ) I ) q\left(x_2 \mid x_1\right)=\mathcal{N}\left(x_2, \sqrt{\alpha_2} x_1,\left(1-\alpha_2\right) I\right) q(x2​∣x1​)=N(x2​,α2​ ​x1​,(1−α2​)I)
整个过程的均值计算如下:
μ t = α t x t − 1 = α t ∗ α t − 1 x t − 2 = α t α t − 1 x t − 2 = α t α t − 1 ∗ α t − 2 x t − 3 = α t α t − 1 α t − 2 x t − 3 = α t α t − 1 α t − 2 ⋯ ∗ α 1 x 0 = α t α t − 1 α t − 2 ⋯ α 1 x 0 = α ˉ t x 0 \begin{gathered} \mu_t=\sqrt{\alpha_t} x_{t-1} \\ =\sqrt{\alpha_t} * \sqrt{\alpha_{t-1}} x_{t-2}=\sqrt{\alpha_t \alpha_{t-1}} x_{t-2} \\ =\sqrt{\alpha_t \alpha_{t-1}} * \sqrt{\alpha_{t-2}} x_{t-3}=\sqrt{\alpha_t \alpha_{t-1} \alpha_{t-2}} x_{t-3} \\ =\sqrt{\alpha_t \alpha_{t-1} \alpha_{t-2} \cdots} * \sqrt{\alpha_1} x_0=\sqrt{\alpha_t \alpha_{t-1} \alpha_{t-2} \cdots \alpha_1} x_0 \\ =\sqrt{\bar{\alpha}_t} x_0 \end{gathered} μt​=αt​ ​xt−1​=αt​ ​∗αt−1​ ​xt−2​=αt​αt−1​ ​xt−2​=αt​αt−1​ ​∗αt−2​ ​xt−3​=αt​αt−1​αt−2​ ​xt−3​=αt​αt−1​αt−2​⋯ ​∗α1​ ​x0​=αt​αt−1​αt−2​⋯α1​ ​x0​=αˉt​ ​x0​​
最终可以写出从 0 到 t 的整个过程的公式:
q ( x t ∣ x 0 ) = N ( x t , α ˉ t x 0 , ( 1 − α ˉ t ) I ) q\left(x_t \mid x_0\right)=\mathcal{N}\left(x_t, \sqrt{\bar{\alpha}_t} x_0,\left(1-\bar{\alpha}_t\right) I\right) q(xt​∣x0​)=N(xt​,αˉt​ ​x0​,(1−αˉt​)I)

然而上述的过程是无法计算的,需要对每一步前向步骤应用重参数化技巧: N ( μ , σ 2 ) = μ + σ ∗ ϵ \mathcal{N}\left(\mu, \sigma^2\right)=\mu+\sigma * \epsilon N(μ,σ2)=μ+σ∗ϵ, 其中 ϵ \epsilon ϵ来自于 N ( 0 , 1 ) \mathcal{N}(0,1) N(0,1)
q ( x t ∣ x t − 1 ) = N ( x t , 1 − β t x t − 1 , β t I ) = 1 − β t x t − 1 + β t ϵ \begin{gathered} q\left(x_t \mid x_{t-1}\right)=\mathcal{N}\left(x_t, \sqrt{1-\beta_t} x_{t-1}, \beta_t I\right) \\ =\sqrt{1-\beta_t} x_{t-1}+\sqrt{\beta_t} \epsilon \end{gathered} q(xt​∣xt−1​)=N(xt​,1−βt​ ​xt−1​,βt​I)=1−βt​ ​xt−1​+βt​ ​ϵ​
因为我们知道 t = 1.. T t=1..T t=1..T 的公式是什么,我们可以使用这个技巧来得到可计算的方程式。
q ( x t ∣ x t − 1 ) = N ( x t , 1 − β t x t − 1 , β t I ) = 1 − β t x t − 1 + β t ϵ \begin{gathered} q\left(x_t \mid x_{t-1}\right)=\mathcal{N}\left(x_t, \sqrt{1-\beta_t} x_{t-1}, \beta_t I\right) \\ =\sqrt{1-\beta_t} x_{t-1}+\sqrt{\beta_t} \epsilon \end{gathered} q(xt​∣xt−1​)=N(xt​,1−βt​ ​xt−1​,βt​I)=1−βt​ ​xt−1​+βt​ ​ϵ​

反向去噪过程

可视化步骤

第 t 步: 从预测的噪声中取一部分,并将其从输入中减去。减法量由调度定义。第一步只减去预测噪声的一小部分,后续步骤逐渐减去更多预测的噪声。 将第t步的输出变成第t-1步的输入 第t-1步:再次生成整个噪声,这一次取更多的噪声并从输入中减去(因为 beta 增加),输出的结果作为t-2步的输入 第1步:得到最终生成的图像

数学推导
论文中定义的反向去噪的公式如下:
p θ ( x 0 : T ) : = p ( x T ) ∏ t = 1 T p θ ( x t − 1 ∣ x t ) p_\theta\left(x_{0: T}\right):=p\left(x_T\right) \prod_{t=1}^T p_\theta\left(x_{t-1} \mid x_t\right) pθ​(x0:T​):=p(xT​)t=1∏T​pθ​(xt−1​∣xt​)
其中
p θ ( x t − 1 ∣ x t ) = N ( x t − 1 , μ θ ( x t , t ) , ∑ θ ( x t , t ) ) p_\theta\left(x_{t-1} \mid x_t\right)=\mathcal{N}\left(x_{t-1}, \mu_\theta\left(x_t, t\right), \sum_\theta\left(x_t, t\right)\right) pθ​(xt−1​∣xt​)=N(xt−1​,μθ​(xt​,t),θ∑​(xt​,t))
扩散过程 p θ ( x 0 : T ) p_\theta\left(x_{0: T}\right) pθ​(x0:T​)是一个从 p ( x T ) p(x_T) p(xT​)开始的高斯转换链 (chain of gaussian transitions),使用一个扩散过程步骤的方程在 𝑇 次迭代中进行,即 p θ ( x t − 1 ∣ x t ) p_\theta\left(x_{t-1} \mid x_t\right) pθ​(xt−1​∣xt​)。

N ( x t − 1 , μ θ ( x t , t ) , ∑ θ ( x t , t ) ) \mathcal{N}\left(x_{t-1}, \mu_\theta\left(x_t, t\right), \sum_\theta\left(x_t, t\right)\right) N(xt−1​,μθ​(xt​,t),∑θ​(xt​,t)) 包含两个部分 :

第一个部分是均值 μ θ ( x t , t ) \mu_\theta\left(x_t, t\right) μθ​(xt​,t)

μ θ ( x t , t ) = 1 α t ( x t − 1 − α t 1 − α ˉ t ϵ θ ( x t , t ) ) \mu_\theta\left(x_t, t\right)=\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta\left(x_t, t\right)\right) μθ​(xt​,t)=αt​ ​1​(xt​−1−αˉt​ ​1−αt​​ϵθ​(xt​,t))

第二部分是方差 ∑ θ ( x t , t ) \sum_\theta\left(x_t, t\right) ∑θ​(xt​,t) (等于 σ t 2 I \sigma_t^2 I σt2​I)
方差是不可训练的,但是依赖于时间变化,即 β T I \beta_TI βT​I (这里的 β \beta β 和前向过程中的 β \beta β一致,都表示的是调度)

这里先不涉及具体的推导,只是简单介绍下流程,具体的推导可以参考论文中的Appendix A [2] 和Lilian Weng的博文[b.5]

这样,我们可以得到 x t − 1 x_{t-1} xt−1​
x t − 1 = N ( x t − 1 , 1 α t ( x t − β t 1 − α ˉ t ϵ θ ( x t , t ) ) , β t ϵ ) x_{t-1}=\mathcal{N}\left(x_{t-1}, \frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta\left(x_t, t\right)\right), \sqrt{\beta_t} \epsilon\right) xt−1​=N(xt−1​,αt​ ​1​(xt​−1−αˉt​ ​βt​​ϵθ​(xt​,t)),βt​ ​ϵ)
进一步的,我们计算给定时间步t时刻的输出:
x t − 1 = 1 a t ( x t − β t 1 − α ˉ ϵ θ ( x t , t ) ) + β t ϵ x_{t-1}=\frac{1}{\sqrt{a_t}}\left(x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}}} \epsilon_\theta\left(x_t, t\right)\right)+\sqrt{\beta_t} \epsilon xt−1​=at​ ​1​(xt​−1−αˉ ​βt​​ϵθ​(xt​,t))+βt​ ​ϵ
其中 ϵ θ ( x t , t ) \epsilon_\theta\left(x_t, t\right) ϵθ​(xt​,t) 是模型的输出,即预测出的噪音。

网络结构

模型架构是修改后的 U-Net 架构 [4]。它非常简单,但随着对扩散模型的进一步改进(例如,Stabel Diffusion 增加了整个潜在层用于图像数据嵌入)。在这里,我们先只讨论Diffusion Model的最初版本


首先修改的Unet 架构,将基础的ResNetBlock替换成Self-Attention Block, 或者修改后的ResNet Block。

然后,为了防止信息丢失,添加跳跃连接(skip connection)

下一步是添加关于当前时间步 t 的信息。为此,我们使用正弦嵌入(sinusoidal embedding),该信息被添加到所有下采样和上采样块中。

为了使网络是条件化的(依赖于外部输入),我们添加了一个文本嵌入器。它将为给定的文本创建嵌入,并将输出向量添加到时间步嵌入中。

ResNet Block
ResNet Block的结构如下,被用于下采样和上采样块的一部分。

Downsample Block

下采样块是第一个不仅接收来自前一层的数据,还接收有关时间步长和提示的数据的块。该块具有2个输入,并且表现为 U-Net 架构中的标准下采样。它接收输入并将其下采样到下一层的大小。
它使用 MaxPool2d 层(核大小为2),将输入大小减半(64x64 -> 32x32)。之后,连接2个 ResNet 块。
嵌入经过 Sigmoid 线性单元处理,然后通过一个简单的Linear层,以达到与 ResNet 块输出相同的形状。之后,两个张量相加,并发送到下一个块。

Self-Attention Block

所有的注意力块的结构都相同,所以将用第一个来描述所有的注意力块(就在第一个下采样之后)。

首先,它接收一个形状为 (128, 32, 32) 的下采样张量。
然后,对输入张量进行reshape: 压缩最后两个维度,然后翻转得到的张量((128,32,32) -> (128, 1024) -> (1024, 128))并通过层归一化,并将其用作所有 3 个输入张量(Q、K、V)
接着,输入到一个多头注意力(Multi-Head Attention, MHA),嵌入维度设置为 128,有 4 个注意力头。嵌入维度在注意力块之间变化(取决于输入长度),而头的数量保持不变。

在块内部,作者添加了两个跳跃连接,与自注意力输出组合。
第一个连接将reshape 后的输入添加到注意力层的输出中,并通过前向层(Normalization -> Linear -> GELU -> Linear)传递。
第二个连接从该层获取输出,并将其添加回注意力的输出。

最后,我们需要对张量进行reshape,最终输出的大小和输入大小一致((1024, 128) -> (128, 32, 32))。

Upsample block

使用一个简单的上采样层,缩放因子为2。
通过上采样层传递输入张量后,我们可以将其与残差连接(来自第一个 ResNet 块)连接起来。现在,它们两者的形状都是 (64, 64, 64)。
连接的张量然后通过 2 个 ResNet 块发送(与我们在下采样块中所做的相同)
第三个输入(再次与下采样块中相同)通过 SILU 和线性层发送,然后添加到第二个 ResNet 块的结果中。
整个架构以 Conv2d 层结束,该层使用核大小为 1,将我们的张量从 (64,64,64) 缩放回 (3, 64,64)。这是我们预测的噪声。

训练和推理过程

训练过程

上图是训练的伪代码:

数据集中采样一张图片 x 0 x_0 x0​ 采样一个时间步 t t t 从正态分布中采样一个噪音 ϵ \epsilon ϵ 前面我们定义过了如何在时间步t时刻,在无需迭代的情况下前向加噪 q ( x t ∣ x 0 ) = α ˉ t x 0 + 1 − α ˉ t ϵ q\left(x_t \mid x_0\right)=\sqrt{\bar{\alpha}_t} x_0+\sqrt{1-\bar{\alpha}_t} \epsilon q(xt​∣x0​)=αˉt​ ​x0​+1−αˉt​ ​ϵ, 然后用梯度下降去训练网络模型 ϵ θ \epsilon_\theta ϵθ​ ( ϵ θ \epsilon_\theta ϵθ​是用于学习生成噪音的模型Noise Predicter)。 然后重复整个过程,直到模型最终收敛。

下图是训练过程的第5步的图解:

下图是Noise Predicter 的图解实例,它的输入有两个:一个是时间步t, 另外一个是时间步t时刻对应的添加的噪音。

下图是实际的训练过程。input和ground truth和时间步 来自于前向加噪的过程。在实际应用中,Noise Predicter 还会有个文本提示的输入。

推理过程

推理

上图是推理的伪代码:
1) 先从标准高斯分布中采样一个噪音
2)从时间步 T T T开始,逐步推理前一步的输出 (输入t时刻的带噪图片 x t x_t xt​,生成t-1时刻的带噪音图片 x t − 1 x_{t-1} xt−1​)。

优化目标

对于真实的图片分布 P d a t a P_{data} Pdata​中采样一些图片 x 1 , x 2 , x 3 . . . x^1,x^2,x^3... x1,x2,x3... 。然后从一个标准的高斯分布中采样一个噪音,通过生成网络(需要训练的参数集是 θ \theta θ) 得到生成图片的分布。
优化的目标是,对于给定的图片 x i x^i xi, 生成网络生成 x i x^i xi的概率越大越好。

最大化生成分布和真实分布的相似性 等价于 最小化二者的KL散度,推理步骤如下:

第一行:取log 并不印象最大化的计算
第二行:根据log的计算性质进行变换 l o g ( x 1 ) ∗ l o g ( x 2 ) = l o g ( x 1 + x 2 ) log(x_1)*log(x_2)=log(x_1+x_2) log(x1​)∗log(x2​)=log(x1​+x2​); 最终的求和式子近似于计算 期望值
第三行:用连续积分的形式计算期望,后面减去的部分是和 θ \theta θ 不相关的,不会影响最大化的结果。
第四行: l o g ( x 1 − x 2 ) = l o g ( x 1 ) / l o g ( x 2 ) log(x_1-x_2)=log(x_1)/ log(x_2) log(x1​−x2​)=log(x1​)/log(x2​) 最终变换成KL散度的形式。

详细数学推导

数学基础

高斯分布
定义:随机变量X服从一个数学期望 μ \mu μ、方差为 σ \sigma σ的高斯分布,又名正态分布。
当 μ = 0 , σ = 1 \mu=0, \sigma=1 μ=0,σ=1时的正态分布是标准正态分布。
高斯分布概率密度函数(正态随机变量概率密度函数):
f ( x ) = 1 2 π σ exp ⁡ ( − ( x − μ ) 2 2 σ 2 ) f(x)=\frac{1}{\sqrt{2 \pi} \sigma} \exp \left(-\frac{(x-\mu)^2}{2 \sigma^2}\right) f(x)=2π ​σ1​exp(−2σ2(x−μ)2​)

两个高斯分布相加
已知两个独立的高斯分布:
N 1 ∼ ( u 1 , δ 1 2 ) , N 2 ∼ ( u 2 , δ 2 2 ) N_1 \sim\left(u_1, \delta_1^2\right), N_2 \sim\left(u_2, \delta_2^2\right) N1​∼(u1​,δ12​),N2​∼(u2​,δ22​)
这两个高斯分布的和,仍然是一个高斯分布
N 1 + N 2 ∼ ( u 1 + u 2 , δ 1 2 + δ 2 2 ) N_1+N_2 \sim\left(u_1+u_2, \delta_1^2+\delta_2^2\right) N1​+N2​∼(u1​+u2​,δ12​+δ22​)
同理,这两个高斯分布的差为:
N 1 − N 2 ∼ ( u 1 − u 2 , δ 1 2 + δ 2 2 ) N_1-N_2 \sim\left(u_1-u_2, \delta_1^2+\delta_2^2\right) N1​−N2​∼(u1​−u2​,δ12​+δ22​)

重参数化
N ( μ , σ 2 ) = μ + σ ∗ ϵ N\left(\mu, \sigma^2\right)=\mu+\sigma * \epsilon N(μ,σ2)=μ+σ∗ϵ
其中 ϵ \epsilon ϵ来自于 N ( 0 , 1 ) \mathcal{N}(0,1) N(0,1)

贝叶斯定理
条件概率:在已知一个事件 B B B 发生的情况下,此时另一个事件A发生的概率,符号表示为 P ( A ∣ B ) P(A \mid B) P(A∣B) ,它等于联合概率 P ( A , B ) P(A, B) P(A,B) 除以边缘分布 P ( B ) P(B) P(B), 即:
P ( A ∣ B ) = P ( A , B ) P ( B ) P(A \mid B)=\frac{P(A, B)}{P(B)} P(A∣B)=P(B)P(A,B)​
或可以写成:
P ( A , B ) = P ( A ∣ B ) ⋅ P ( B ) P(A, B)=P(A \mid B) \cdot P(B) P(A,B)=P(A∣B)⋅P(B)
表示事件A,B联合概率等于在B发生的条件下A发生的条件概率乘以B事件发生概率。
由于P(A,B)也等于P(B|A)P(A) (P(A,B)与P(B,A)含义相同,数值相等),所以推导出贝叶斯公式:
P ( A ∣ B ) = P ( A , B ) P ( B ) = P ( B ∣ A ) ⋅ P ( A ) P ( B ) P(A \mid B)=\frac{P(A, B)}{P(B)}=\frac{P(B \mid A) \cdot P(A)}{P(B)} P(A∣B)=P(B)P(A,B)​=P(B)P(B∣A)⋅P(A)​

KL散度

向前扩散过程

给定从真实数据分布中采样的数据点 x 0 ∼ q ( x ) \mathbf{x}_0 \sim q(\mathbf{x}) x0​∼q(x),让我们定义一个前向扩散过程,在其中我们在 T T T 步中向样本添加小量的高斯噪声,产生一系列带有噪声的样本 x 1 , … , x T \mathbf{x}_1, \ldots, \mathbf{x}_T x1​,…,xT​。步长由方差调度 { β t ∈ ( 0 , 1 ) } t = 1 T \left\{\beta_t \in(0,1)\right\}_{t=1}^T {βt​∈(0,1)}t=1T​控制。
q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q ( x 1 : T ∣ x 0 ) = ∏ t = 1 T q ( x t ∣ x t − 1 ) q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)=\mathcal{N}\left(\mathbf{x}_t ; \sqrt{1-\beta_t} \mathbf{x}_{t-1}, \beta_t \mathbf{I}\right) \quad q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_0\right)=\prod_{t=1}^T q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right) q(xt​∣xt−1​)=N(xt​;1−βt​ ​xt−1​,βt​I)q(x1:T​∣x0​)=t=1∏T​q(xt​∣xt−1​)
随着步长 t t t的增大,数据样本 x 0 \mathbf{x_0} x0​逐渐失去其可区分的特征。最终,当 T T T趋向于正无穷时, x t \mathbf{x_t} xt​等效于一个各向同性的高斯分布。

通过逐渐添加(去除)噪声生成样本的前向(反向)扩散过程的马尔可夫链。

我们可以使用重参数化技巧在任意时间步 t t t中以封闭形式对 x t \mathbf{x_t} xt​进行采样。
令 α t = 1 − β t \alpha_t=1-\beta_t αt​=1−βt​ , α ˉ t = ∏ i = 1 t α i \bar{\alpha}_t=\prod_{i=1}^t \alpha_i αˉt​=∏i=1t​αi​
x t = α t x t − 1 + 1 − α t ϵ t − 1 ;  where  ϵ t − 1 , ϵ t − 2 , ⋯ ∼ N ( 0 , I ) = α t α t − 1 x t − 2 + 1 − α t α t − 1 ϵ ˉ t − 2 ;  where  ϵ ‾ t − 2  merges two Gaussians  ( ∗ ) .  = … = α ˉ t x 0 + 1 − α ˉ t ϵ q ( x t ∣ x 0 ) = N ( x t ; α ˉ t x 0 , ( 1 − α ˉ t ) I ) \begin{aligned} & \mathbf{x}_t=\sqrt{\alpha_t} \mathbf{x}_{t-1}+\sqrt{1-\alpha_t} \boldsymbol{\epsilon}_{t-1} \quad \quad ; \text { where } \boldsymbol{\epsilon}_{t-1}, \boldsymbol{\epsilon}_{t-2}, \cdots \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\ & =\sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2}+\sqrt{1-\alpha_t \alpha_{t-1}} \bar{\epsilon}_{t-2} \quad ; \text { where } \overline{\boldsymbol{\epsilon}}_{t-2} \text { merges two Gaussians }\left(^*\right) \text {. } \\ & =\ldots \\ & =\sqrt{\bar{\alpha}_t} \mathbf{x}_0+\sqrt{1-\bar{\alpha}_t} \boldsymbol{\epsilon} \\ & q\left(\mathbf{x}_t \mid \mathbf{x}_0\right)=\mathcal{N}\left(\mathbf{x}_t ; \sqrt{\bar{\alpha}_t} \mathbf{x}_0,\left(1-\bar{\alpha}_t\right) \mathbf{I}\right) \\ & \end{aligned} ​xt​=αt​ ​xt−1​+1−αt​ ​ϵt−1​; where ϵt−1​,ϵt−2​,⋯∼N(0,I)=αt​αt−1​ ​xt−2​+1−αt​αt−1​ ​ϵˉt−2​; where ϵt−2​ merges two Gaussians (∗). =…=αˉt​ ​x0​+1−αˉt​ ​ϵq(xt​∣x0​)=N(xt​;αˉt​ ​x0​,(1−αˉt​)I)​

将两个具有不同方差的高斯分布
N ( 0 , σ 1 2 I ) \mathcal{N}\left(\mathbf{0}, \sigma_1^2 \mathbf{I}\right) N(0,σ12​I) 和 N ( 0 , σ 2 2 I ) \mathcal{N}\left(\mathbf{0}, \sigma_2^2 \mathbf{I}\right) N(0,σ22​I)合并,得到的新的分布是: N ( 0 , ( σ 1 2 + σ 2 2 ) I ) \mathcal{N}\left(\mathbf{0},\left(\sigma_1^2+\sigma_2^2\right) \mathbf{I}\right) N(0,(σ12​+σ22​)I).
将 x t − 1 = α t − 1 x t − 2 + 1 − α t − 1 ϵ t − 2 \mathbf{x}_{t-1}=\sqrt{\alpha_{t-1}} \mathbf{x}_{t-2}+\sqrt{1-\alpha_{t-1}} \boldsymbol{\epsilon}_{t-2} xt−1​=αt−1​ ​xt−2​+1−αt−1​ ​ϵt−2​ 代入当上式中,得到:
x t = a t ( a t − 1 x t − 2 + 1 − α t − 1 ϵ 2 ) + 1 − α t z 1 = a t a t − 1 x t − 2 + ( a t ( 1 − α t − 1 ) ϵ 2 + 1 − α t z 1 ) x_t=\sqrt{a_t}\left(\sqrt{a_{t-1}} x_{t-2}+\sqrt{1-\alpha_{t-1}} \boldsymbol{\epsilon}_2\right)+\sqrt{1-\alpha_t} z_1 \\ =\sqrt{a_t a_{t-1}} x_{t-2}+\left(\sqrt{a_t\left(1-\alpha_{t-1}\right)}\boldsymbol{\epsilon}_2+\sqrt{1-\alpha_t} z_1\right) xt​=at​ ​(at−1​ ​xt−2​+1−αt−1​ ​ϵ2​)+1−αt​ ​z1​=at​at−1​ ​xt−2​+(at​(1−αt−1​) ​ϵ2​+1−αt​ ​z1​)
由于每次加入的噪音 ϵ 1 \boldsymbol{\epsilon}_1 ϵ1​和 ϵ 2 \boldsymbol{\epsilon}_2 ϵ2​都服从标准高斯分布 N ∼ ( 0 , I ) \mathcal{N} \sim (0, I) N∼(0,I)。
则 a t ( 1 − α t − 1 ) z 2 ∼ N ( 0 , a t ( 1 − a t − 1 ) ) \sqrt{a_t\left(1-\alpha_{t-1}\right)} z_2 \sim N (0, a_t(1-a_{t-1})) at​(1−αt−1​) ​z2​∼N(0,at​(1−at−1​)) , 1 − α t z 1 ∼ N ( 1 − a t ) \sqrt{1-\alpha_t} z_1 \sim N(1-a_t) 1−αt​ ​z1​∼N(1−at​)
这两个分布相加
a t ( 1 − α t − 1 ) z 2 + 1 − α t z 1 ∼ N ( 0 , ( 1 − α t ) + α t ( 1 − α t − 1 ) ) = N ( 0 , 1 − α t α t − 1 ) \sqrt{a_t\left(1-\alpha_{t-1}\right)} z_2 + \sqrt{1-\alpha_t} z_1 \sim N(0, \left(1-\alpha_t\right)+\alpha_t\left(1-\alpha_{t-1}\right) )=N(0, 1-\alpha_t \alpha_{t-1}) at​(1−αt−1​) ​z2​+1−αt​ ​z1​∼N(0,(1−αt​)+αt​(1−αt−1​))=N(0,1−αt​αt−1​)

通常,当sample变的noise后,我们可以采用更大的更新率。因此 β 1 < β 2 < ⋯ < β T \beta_1<\beta_2<\cdots<\beta_T β1​<β2​<⋯<βT​ , α ˉ 1 > ⋯ > α ˉ T \bar{\alpha}_1>\cdots>\bar{\alpha}_T αˉ1​>⋯>αˉT​.

与随机梯度Langevin动力学的联系
Langevin 动力学是物理学中的一个概念,用于统计建模分子系统。结合随机梯度下降,随机梯度 Langevin 动力学( stochastic gradient Langevin dynamics)可以使用更新的马尔可夫链中的梯度 ∇ x log ⁡ p ( x ) \nabla_{\mathbf{x}} \log p(\mathbf{x}) ∇x​logp(x)来从概率密度 p ( x ) p(\mathbf{x}) p(x)中产生样本:
x t = x t − 1 + δ 2 ∇ x log ⁡ p ( x t − 1 ) + δ ϵ t ,  where  ϵ t ∼ N ( 0 , I ) \mathbf{x}_t=\mathbf{x}_{t-1}+\frac{\delta}{2} \nabla_{\mathbf{x}} \log p\left(\mathbf{x}_{t-1}\right)+\sqrt{\delta} \boldsymbol{\epsilon}_t, \quad \text { where } \boldsymbol{\epsilon}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) xt​=xt−1​+2δ​∇x​logp(xt−1​)+δ ​ϵt​, where ϵt​∼N(0,I)
其中 δ \delta δ 是 step size. 当 T → ∞ , ϵ → 0 , x T T \rightarrow \infty, \epsilon \rightarrow 0, \mathbf{x}_T T→∞,ϵ→0,xT​ 等同于true probability density p ( x ) p(\mathbf{x}) p(x).
与标准随机梯度下降相比,随机梯度 Langevin 动力学在参数更新中注入高斯噪声,以避免陷入局部最小值。

反向去噪过程

如果我们能够反转上述过程并从 q ( x t − 1 ∣ x t ) q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right) q(xt−1​∣xt​) 中采样,我们将能够从高斯噪声输入 x T ∼ N ( 0 , I ) \mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) xT​∼N(0,I) 中重新创建真实样本。请注意,如果 β t \beta_t βt​ 足够小, q ( x t − 1 ∣ x t ) q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right) q(xt−1​∣xt​) 也将是高斯分布。不幸的是,我们无法轻松地估计 q ( x t − 1 ∣ x t ) q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right) q(xt−1​∣xt​),因为它需要使用整个数据集,因此我们需要学习一个模型 p θ p_\theta pθ​ 来近似这些条件概率,以便运行反向扩散过程。
p θ ( x 0 : T ) = p ( x T ) ∏ t = 1 T p θ ( x t − 1 ∣ x t ) p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_\theta\left(\mathbf{x}_{0: T}\right)=p\left(\mathbf{x}_T\right) \prod_{t=1}^T p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right) \quad p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_\theta\left(\mathbf{x}_t, t\right), \mathbf{\Sigma}_\theta\left(\mathbf{x}_t, t\right)\right) pθ​(x0:T​)=p(xT​)t=1∏T​pθ​(xt−1​∣xt​)pθ​(xt−1​∣xt​)=N(xt−1​;μθ​(xt​,t),Σθ​(xt​,t))

训练扩散模型来建模 2D swiss roll 数据的示例。

当条件是 x 0 \mathcal{x_0} x0​时,反向条件概率是可处理的。

使用Bayes法则,可以得到

其中 C ( x t , x 0 ) C\left(\mathbf{x}_t, \mathbf{x}_0\right) C(xt​,x0​) 是某个不涉及 x t − 1 \mathbf{x}_{t-1} xt−1​ 的常数项,不影响。

对于上式的第一行,利用贝叶斯定理展开后的三项都是已知的:
q ( x t − 1 ∣ x 0 ) a t ‾ − 1 x 0 + 1 − a t ‾ − 1 ∼ N ( a t ‾ − 1 , x 0 , 1 − a t − 1 ‾ ) q ( x t ∣ x 0 ) α ˉ t x 0 + 1 − α ˉ t z ∼ N ( α ˉ t x 0 , 1 − α ˉ t ) q ( x t ∣ x t − 1 , x 0 ) = a t x t − 1 + 1 − α t z ∼ N ( a t x t − 1 , 1 − α t ) \begin{aligned} & q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_0\right) \quad \sqrt{\overline{a_t}-1} x_0+\sqrt{1-\overline{a_t}}-1 \quad \sim \mathcal{N}\left(\sqrt{\overline{a_t}}-1, x_0, 1-\overline{a_{t-1}}\right) \\ & q\left(\mathbf{x}_t \mid \mathbf{x}_0\right) \quad \sqrt{\bar{\alpha}_t} x_0+\sqrt{1-\bar{\alpha}_t} z \quad \sim \mathcal{N}\left(\sqrt{\bar{\alpha}_t} x_0, 1-\bar{\alpha}_t\right) \\ & q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \mathbf{x}_0\right)=\sqrt{a_t} x_{t-1}+\sqrt{1-\alpha_t} z \quad \sim \mathcal{N}\left(\sqrt{a_t} x_{t-1}, 1-\alpha_t\right) \end{aligned} ​q(xt−1​∣x0​)at​​−1 ​x0​+1−at​​ ​−1∼N(at​​ ​−1,x0​,1−at−1​​)q(xt​∣x0​)αˉt​ ​x0​+1−αˉt​ ​z∼N(αˉt​ ​x0​,1−αˉt​)q(xt​∣xt−1​,x0​)=at​ ​xt−1​+1−αt​ ​z∼N(at​ ​xt−1​,1−αt​)​
将标准正态分布展开之后,因为是指数运算,乘法相当于加法,除法相当于减法。因此可以得到第二行。
第三行和第四行进行展开成如下的格式:
exp ⁡ ( − ( x − μ ) 2 2 σ 2 ) = exp ⁡ ( − 1 2 ( 1 σ 2 x 2 − 2 μ σ 2 x + μ 2 σ 2 ) ) \exp \left(-\frac{(x-\mu)^2}{2 \sigma^2}\right)=\exp \left(-\frac{1}{2}\left(\frac{1}{\sigma^2} x^2-\frac{2 \mu}{\sigma^2} x+\frac{\mu^2}{\sigma^2}\right)\right) exp(−2σ2(x−μ)2​)=exp(−21​(σ21​x2−σ22μ​x+σ2μ2​))
这样可以进一步的进行配方,得到均值和方差。

按照标准高斯密度函数,均值和方差可以参数化如下(回想一下 α t = 1 − β t \alpha_t=1-\beta_t αt​=1−βt​ 和 α ˉ t = ∏ i = 1 T α i ) \left.\bar{\alpha}_t=\prod_{i=1}^T \alpha_i\right) αˉt​=∏i=1T​αi​)

上式均值中 x 0 x_0 x0​是未知的,但是我们可以通过从 x 0 x_0 x0​生成 x t x_t xt​的公式逆过来,得到 x 0 = 1 α ˉ t ( x t − 1 − α ˉ t ϵ t ) \mathbf{x}_0=\frac{1}{\sqrt{\bar{\alpha}_t}}\left(\mathbf{x}_t-\sqrt{1-\bar{\alpha}_t} \epsilon_t\right) x0​=αˉt​ ​1​(xt​−1−αˉt​ ​ϵt​) 代入上式,得到:

然而 ϵ t \epsilon_t ϵt​仍然是未知的,这个就是用Unet预测的每个时刻的噪音。

Stable Diffusion

Stable Diffusion(SD)模型,可以用于文生图,图生图,图像inpainting,ControlNet控制生成,图像超分等丰富的任务。

文生图任务是指将一段文本输入到SD模型中,经过一定的迭代次数,SD模型输出一张符合输入文本描述的图片。
图生图任务在输入本文的基础上,再输入一张图片,SD模型将根据文本的提示,将输入图片进行重绘以更加符合文本的描述。

组成结构

Stable Diffusion 是一个系统,由几个组件和模型构成,而非单独一个模型。

Text Encoder
功能:对输入的文本进行编码。
实现:CLIP 模型中的文本编码器
输入:文本
输出:77个表征向量,每个有768维度

Image Information Creator
功能:逐渐处理/扩散信息到潜在(信息)空间。
实现:Unet网络+Scheduler
输入:文本向量 和 起始多维数组(结构化的数字列表,也叫Tensor)由噪声组成。
输出:处理后的信息数组

输入:
1) 如果是图生图任务,我们在输入文本信息的同时,还需要将原图片通过图像编码器(VAE Encoder)生成Latent Feature(隐空间特征)作为输入。
2) 如果是文生图任务,我们只需要输入文本信息,再用random函数生成一个高斯噪声矩阵作为Latent Feature的“替代”输入

组成:
1)U-Net网络负责预测噪声,不断优化生成过程,在预测噪声的同时不断注入文本语义信息。
2)schedule算法对每次U-Net预测的噪声进行优化处理(动态调整预测的噪声,控制U-Net预测噪声的强度),从而统筹生成过程的进度。

Image Decoder
图像解码器通过前置信息生成器得到的信息来生成图像。它只在最后的运行一次来生成最终的像素图像。
输入:处理过的信息数组(维度:(4,64,64))
输出:结果图像(维度:(3,512,512)分别是红/绿/蓝,宽,高)

运行流程


Step 1: 将输入的文本提示通过Text Encoder进行编码, 作为Diffusion 的输入。
Step 2:将随机采样的噪音,也作为Diffusion的输出。
Step 3: 得到处理后的Image information 作为图像解码器的输入。

Diffusion 的过程是step-by-step 的,逐步生成清晰的图片。
下图通过图像解码器可视化每个步骤

扩散动作发生了多次,每一步都对输入潜在数组上操作以产生更相似于输入文本的新潜在数组,视觉信息来自于模型从所有图像训练语料。

下图可视化一系列的这些潜在数组:

在SD中,U-Net的迭代优化步数(Timesteps)大概是50或者100次,在这个过程中Latent Feature的质量不断的变好(纯噪声减少,图像语义信息增加,文本语义信息增加)

回顾Diffusion Model的内部工作原理
假设我们有张图片,我们生成了一些噪声,然后把噪声添加到图像上。

我们使用相同的方式来生成大量的训练样本。

使用这份数据集,我们可以训练噪声预测器

训练好的噪声预测器可以产生一个噪声图像,接来下的一系列降噪步骤,可以预测噪声。

采样噪声被预测了出来,所以如果我们从图像中减去它就能得到尽可能接近于模型训练样本的图像

逐步得到最终的生成图像

潜在空间

之前的Diffusion Model 不管是前向扩散过程还是反向扩散过程,都需要在像素级的图像上进行,当图像分辨率和Timesteps很大时,不管是训练还是前向推理,都非常的耗时。
而Stable Diffusion 将这些过程压缩在低维的Latent隐空间,这样一来大大降低了显存占用和计算复杂度

为了加速图像生成的过程,Stable Diffusion 是在图像的压缩空间上扩散操作。
压缩(和稍后的解压缩/绘制)是经由自编码器完成的。自编码器通过编码器压缩图像到潜在空间,然后只用压缩信息来解码重建。

在压缩空间上进行前向扩散过程。噪声数据是用在这些潜在空间上的,而非像素图片上的。所以,噪声预测器实际上是被训练在压缩表示上(潜在空间)来预测噪声的。
一旦训练完后,我们就可以通过反向处理(使用自编码器的解码器)来生成图像。

总结一下SD模型的前向推理过程:

网络结构

Stable Diffusion模型主要由VAE(变分自编码器,Variational Auto-Encoder),U-Net以及CLIP Text Encoder三个核心组件构成。

变分自编码器 (VAE)

VAE(变分自编码器,Variational Auto-Encoder)是基于Encoder-Decoder架构的生成模型。VAE的Encoder(编码器)结构能将输入图像压缩为低维Latent特征,并作为U-Net的输入。VAE的Decoder(解码器)结构能将低维Latent特征重建还原成像素级图像。

VAE的核心作用

1) 图像压缩和图像重建
如下图所示:

当我们输入一个尺寸为 H × W × C H\times W\times C H×W×C的数据,VAE的Encoder模块会将其编码为一个大小为 h × w × c h\times w\times c h×w×c 的低维Latent特征,其中 f = H / h = W / w f=H/h=W/w f=H/h=W/w为VAE的下采样率(Downsampling Factor)。反之,VAE的Decoder模块有一个相同的上采样率(Upsampling Factor)将低维Latent特征重建成像素级别的图像。

为什么VAE可以将图像压缩到一个非常小的Latent space(潜空间)后能再次对图像进行像素级重建呢?
因为虽然VAE对图像的压缩与重建过程是一个有损压缩与重建过程,但图像全图级特征关联并不是随机的,它们的分布具有很强的规律性:比如人脸的眼睛、鼻子、脸颊和嘴巴之间遵循特定的空间关系,又比如一只猫有四条腿,并且这是一个特定的生物结构特征。下面Rocky也使用VAE将图像重建成不同尺寸的生成图像,实验结论发现如果我们重建生成的图像尺寸在之上时,其实特征损失带来的影响非常小。

2) 切换不同微调训练版本的VAE模型,能够发现生成图片的细节与整体颜色也会随之改变

VAE的网络结构

VAE中Encoder和Decoder的结构如下:

SD VAE Encoder部分包含了三个DownBlock模块、一个ResNetBlock模块以及一个MidBlock模块,将输入图像压缩到Latent空间,转换成为Gaussian Distribution。 而VAE Decoder部分正好相反,其输入Latent空间特征,并重建成为像素级图像作为输出。其包含了三个UpBlock模块、一个ResNetBlock模块以及一个MidBlock模块。

三个基础组件:

GSC组件:GroupNorm+Swish+Conv Downsample组件:Padding+Conv Upsample组件:Interpolate+Conv
两个核心组件: ResNetBlock模块 SelfAttention模型

文本编码器 (CLIP Text Encoder)

作为文生图模型,Stable Diffusion中的文本编码模块直接决定了语义信息的优良程度,从而影响到最后图片生成的质量和与文本的一致性。
Stable Diffusion 中采用的文本编码器是多模态领域中的CLIP模型(Contrastive Language-Image Pre-training)中的文本编码器

CLIP 模型介绍
首先,CLIP模型是一个基于对比学习的多模态模型,主要包含Text EncoderImage Encoder两个模型。使用4亿个图片与标签文本对数据集进行训练,来学习图片与本文内容的对应关系。

Text Encoder用来提取文本的特征,可以使用NLP中常用的text transformer模型作为Text Encoder Image Encoder主要用来提取图像的特征,可以使用CNN/Vision transformer模型(ResNet和ViT等)作为Image Encoder。

CLIP模型的训练过程如下:

从训练集中随机取出一张图片和标签文本,接着CLIP模型的任务主要是通过Text Encoder和Image Encoder分别将标签文本和图片提取embedding向量 然后用余弦相似度(cosine similarity)来比较两个embedding向量的相似性,以判断随机抽取的标签文本和图片是否匹配 并进行梯度反向传播,不断进行优化训练。

完成CLIP的训练后,输入配对的图片和标签文本,则Text Encoder和Image Encoder可以输出相似的embedding向量,计算余弦相似度就可以得到接近1的结果。同时对于不匹配的图片和标签文本,输出的embedding向量计算余弦相似度则会接近0。

在Stable Diffusion中主要使用了Text Encoder部分。CLIP Text Encoder模型将输入的文本Prompt进行编码,转换成Text Embeddings(文本的语义信息)U-Net网络的CrossAttention模块嵌入Stable Diffusion中作为Condition条件,对生成图像的内容进行一定程度上的控制与引导,目前SD模型使用的的是CLIP ViT-L/14中的Text Encoder模型。

将编码后的文本信息,也作为噪音生成器的输入。

由于我们将要在潜在空间中操作,所以输入的图像和已经预测的噪声都是在潜在空间的。

CLIP Text Encoder的结构图

一般来说,我们提取CLIP Text Encoder模型最后一层特征作为CrossAttention模块的输入,但是开源社区的不断实践为我们总结了如下经验:

当我们生成二次元内容时,可以选择提取CLIP Text Encoder模型倒数第二层特征; 当我们生成写实场景内容时,可以选择提取CLIP Text Encoder模型最后一层的特征。

噪音预测器 (Unet)

下图展示了不用文本的扩散Unet, 输入是在潜在空间中的噪音图片和时间步长。

通过在ResNet模块之间增加一个Attention层,添加对文本输入的支持。

ResNet模块并不能直接看到文本,但是Attention层将文本表达合并到了潜在空间。现在接下来的ResNet可以使用合并过来的文本信息到它的处理过程了。

Unet模型的作用
预测噪声残差,并结合Sampling method(调度算法:DDPM、DDIM、DPM++等)对输入的特征矩阵进行重构,逐步将其从随机高斯噪声转化成图片的Latent Feature。
在前向推理过程中,SD模型通过反复调用 U-Net,将预测出的噪声残差从原噪声矩阵中去除,得到逐步去噪后的图像Latent Feature,再通过VAE的Decoder结构将Latent Feature重建成像素级图像。

Unet 模型的结构
Stable Diffusion中的U-Net,在传统深度学习时代的Encoder-Decoder结构的基础上,增加了ResNetBlock(包含Time Embedding)模块,Spatial Transformer(SelfAttention + CrossAttention + FeedForward)模块以及CrossAttnDownBlockCrossAttnUpBlockCrossAttnMidBlock模块。

(1)ResNetBlock模块

ResNetBlock模块:借鉴ResNet模型的“残差结构”,让网络能够构建的更深的同时,将Time Embedding信息嵌入模型。
如下图所示,Stable Diffusion U-Net完整结构图中展示了完整的ResNetBlock模块,其输入包括Latent FeatureTime Embedding。首先Latent Feature经过GSC(GroupNorm+SiLU激活函数+卷积)模块后和Time Embedding(经过SiLU激活函数+全连接层处理)加和操作,之后再经过GSC模块Skip Connection而来的输入Latent Feature做加和操作,进行两次特征融合后最终得到ResNetBlock模块的Latent Feature输出,增强SD模型的特征学习能力。

Time Embedding 正是输入到ResNetBlock模块中,为U-Net引入了时间信息(时间步长T,T的大小代表了噪声扰动的强度),模拟一个随时间变化不断增加不同强度噪声扰动的过程,让SD模型能够更好地理解时间相关性。
在SD模型调用U-Net重复迭代去噪的过程中,我们希望在迭代的早期,能够先生成整幅图片的轮廓与边缘特征,随着迭代的深入,再补充生成图片的高频和细节特征信息。由于在每个ResNetBlock模块中都有Time Embedding,就能告诉U-Net现在是整个迭代过程的哪一步,并及时控制U-Net够根据不同的输入特征和迭代阶段而预测不同的噪声残差。

GSC模块:Stable Diffusion U-Net中的最小组件之一,由GroupNorm+SiLU+Conv三者组成。

(2)CrossAttention模块

CrossAttention模块:将文本的语义信息与图像的语义信息进行Attention机制,增强输入文本Prompt对生成图片的控制。
CrossAttention模块是我们 使用输入文本Prompt控制SD模型图片内容生成 的关键。
CrossAttention接受两个输入:一个是ResNetBlock模块的输出的Latent Feature,另外一个是输入文本Prompt经过CLIP Text Encoder模型编码后的Context Embedding
两个输入首先经过Attention机制(将Context Embedding对应的语义信息与图片中对应的语义信息相耦合),输出新的Latent Feature,再将新输出的Latent Feature与输入的Context Embedding再做一次Attention机制,从而使得SD模型学习到了文本与图片之间的特征对应关系。

为什么Context Embedding用来生成K和V,Latent Feature用来生成Q呢?
原因也非常简单:因为在Stable Diffusion中,主要的目的是想把文本信息注入到图像信息中里,所以用图片token对文本信息做 Attention实现逐步的文本特征提取和耦合。

(3) BasicTransformer Block模块

BasicTransformer Block模块:由LayerNorm+SelfAttention+CrossAttention+FeedForward组成,是在CrossAttention子模块的基础上,增加了SelfAttention子模块和Feedforward子模块共同组成的,并且每个子模块都是一个残差结构,这样除了能让文本的语义信息与图像的语义信息更好的融合之外,还能通过SelfAttention机制让模型更好的学习图像数据的特征。

SelfAttention模块:SelfAttention模块的整体结构与CrossAttention模块相同,这是输入全部都是图像信息,不再输入文本信息。
该模块主要是为了让SD模型更好的学习图像数据的整体特征。
再者,SelfAttention可以将输入图像的不同部分(像素或图像Patch)进行交互,从而实现特征的整合和全局上下文的引入,能够让模型建立捕捉图像全局关系的能力,有助于模型理解不同位置的像素之间的依赖关系,以更好地理解图像的语义。
在此基础上,SelfAttention还能减少平移不变性问题,SelfAttention模块可以在不考虑位置的情况下捕捉特征之间的关系,因此具有一定的平移不变性。

FeedForward模块:Attention机制中的经典模块,由GeGlU+Dropout+Linear组成。

(4)Spatial Transformer模块

Spatial Transformer模块:由GroupNorm+Conv+BasicTransformer Block+Conv构成。
Spatial Transformer模块不改变输入输出的尺寸,只在图片对应的位置上融合了语义信息

Spatial Transformer模块是SD U-Net中的核心基础结构,Encoder中的CrossAttnDownBlock模块,Decoder中的CrossAttnUpBlock模块以及CrossAttnMidBlock模块都包含了大量的Spatial Transformer子模块。

在生成式模型中,GroupNorm的效果一般会比BatchNorm更好,生成式模型通常比较复杂,因此需要更稳定和适应性强的归一化方法。 而GroupNorm主要有以下一些优势,让其能够成为生成式模型的标配:

对训练中不同Batch-Size的适应性:在生成式模型中,通常需要使用不同的Batch-Size进行训练和微调。这会导致 BatchNorm在训练期间的不稳定性,而GroupNorm不受Batch-Size的影响,因此更适合生成式模型。 能适应通道数变化:GroupNorm 是一种基于通道分组的归一化方法,更适应通道数的变化,而不需要大量调整。 更稳定的训练:生成式模型的训练通常更具挑战性,存在训练不稳定性的问题。GroupNorm可以减轻训练过程中的梯度问题,有助于更稳定的收敛。 能适应不同数据分布:生成式模型通常需要处理多模态数据分布,GroupNorm 能够更好地适应不同的数据分布,因为它不像 Batch Normalization那样依赖于整个批量的统计信息。

(5) CrossAttnDownBlock/CrossAttnUpBlock/CrossAttnMidBlock模块

CrossAttnDownBlock_X模块:是Stable Diffusion U-Net中Encoder部分的主要模块,由X个(ResNetBlock模块+Spatial Transformer模块)+DownSample模块组成。

CrossAttnUpBlock_X模块:是Stable Diffusion U-Net中Decode部分的主要模块,由X个(ResNetBlock模块+Spatial Transformer模块)+UpSample模块组成。

CrossAttnMidBlock模块:是Stable Diffusion U-Net中EncoderDecoder连接的部分,由ResNetBlock+Spatial Transformer+ResNetBlock组成

DownSample模块:Stable Diffusion U-Net中的下采样组件,使用了Conv(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))进行采下采样。

UpSample模块:Stable Diffusion U-Net中的上采样组件,由插值算法(nearest)+Conv组成。
插值算法将输入的Latent Feature尺寸扩大一倍,同时通过一个卷积(kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))改变Latent Feature的通道数,以便于输入后续的模块中。

UpBlock_X模块:由X个ResNetBlock模块和一个UpSample模块组成。

训练过程和损失函数

VAE训练

在Stable Diffusion中,需要对VAE模型进行微调训练,主要采用了L1回归损失感知损失(perceptual loss,Learned Perceptual Image Patch Similarity,LPIPS)作为损失函数,同时使用了基于patch的对抗训练策略
为了防止在Latent空间的任意缩放导致的标准差过大,在训练VAE模型的过程中引入了正则化损失

L1回归损失
L1回归损失用在回归问题中衡量预测值与真实值之间的差异。
L 1 ( y , y ^ ) = ∑ i = 1 n ∣ y i − y ^ i ∣ L 1(y, \hat{y})=\sum_{i=1}^n\left|y_i-\hat{y}_i\right| L1(y,y^​)=i=1∑n​∣yi​−y^​i​∣

其中, y i y_i yi​ 是输入数据的真实值, y ^ i \hat{y}_i y^​i​ 是模型生成数据的预测值, n n n 是数据总数。

感知损失
感知损失的核心思想是比较原始图像和生成图像在传统深度学习模型(VGG、ResNet、ViT等)不同层中特征图之间的相似度,而不直接进行像素级别的对比。传统深度学习模型能够提取图像的高维语义信息的特征,如果两个图像在高维语义信息的特征上接近,那么它们在像素级别的语意上也应该是相似的。
感知损失的公式如下所示:
L perceptual  = ∑ l λ l ⋅ ∥ ϕ l ( I pred  ) − ϕ l ( I target  ) ∥ 2 2 L_{\text {perceptual }}=\sum_l \lambda_l \cdot\left\|\phi_l\left(I_{\text {pred }}\right)-\phi_l\left(I_{\text {target }}\right)\right\|_2^2 Lperceptual ​=l∑​λl​⋅∥ϕl​(Ipred ​)−ϕl​(Itarget ​)∥22​

其中: ϕ l \phi_l ϕl​ 表示在预训练模型 (比如VGG/ResNet网络) 的第 l l l 层的激活特征。 I p r e d I_{p r e d} Ipred​ 是模型生成的图像。 I target  I_{\text {target }} Itarget ​ 是真实图像。 λ l \lambda_l λl​ 是第 l l l 层的权重,可以根据实际情况设置合适值。

正则化损失

KL(Kullback-Leibler)正则化
让Latnet特征不要偏离正态分布太远,同时设置了较小的权重(~10e-6)来保证VAE的重建效果。 VQ(Vector Quantization)正则化
过在decoder模块中引入一个VQ-layer,将VAE转换成VQ-GAN,同样为了保证VAE的重建效果,设置较高的codebook采样维度(8192)。

虽然VAE模型使用了KL正则化,但是由于KL正则化的权重系数非常小,实际生成的Latent特征的标准差依旧存在比较大的情况,所以Stable Diffusion论文中提出了一种rescaling方法强化正则效果。
首先我们计算第一个batch数据中Latent特征的标准差 σ \sigma σ ,然后采用 1 / σ 1 / \sigma 1/σ 的系数来rescale后续所有的Latent特征使其标准差接近于1。同时在Decoder模块进行重建时,只需要将生成的Latent 特征除以 1 / σ 1 / \sigma 1/σ ,再进行像素级重建即可。
在SD中,U-Net模型使用的是经过rescaling后的Latent 特征,并且将rescaling系数设置为 0.18215 。

基于patch的对抗训练策略
使用PatchGAN的判别器来对VAE模型进行对抗训练,通过优化判别器损失,来提升生成图像的局部真实性(纹理和细节)与清晰度。
PatchGAN是GAN系列模型的一个变体,其判别器架构不再评估整个生成图像是否真实,而是评估生成图像中的patch部分是否真实。具体来说,PatchGAN的判别器接收一张图像,并输出一个矩阵,矩阵中的每个元素代表图像中对应区域的真实性。

总损失函数
Diffusion中VAE训练的完整损失函数:
L Autoencoder  = min ⁡ E , D max ⁡ ψ ( L r e c ( x , D ( E ( x ) ) ) − L a d v ( D ( E ( x ) ) ) + log ⁡ D ψ ( x ) + L r e g ( x ; E , D ) ) \begin{gathered} \mathcal{L}_{\text {Autoencoder }}=\min _{\mathcal{E}, \mathcal{D}} \max _\psi \\ \left(\mathcal{L}_{r e c}(x, \mathcal{D}(\mathcal{E}(x)))-\mathcal{L}_{a d v}(\mathcal{D}(\mathcal{E}(x)))+\log \mathcal{D}_\psi(x)+\mathcal{L}_{r e g}(x ; \mathcal{E}, \mathcal{D})\right) \end{gathered} LAutoencoder ​=E,Dmin​ψmax​(Lrec​(x,D(E(x)))−Ladv​(D(E(x)))+logDψ​(x)+Lreg​(x;E,D))​

其中 E ( x ) \mathcal{E}(x) E(x) 表示VAE重建的图像, L r e c \mathcal{L}_{r e c} Lrec​ 表示L1回归损失和感知损失, L a d v + log ⁡ D ψ \mathcal{L}_{a d v}+\log \mathcal{D}_\psi Ladv​+logDψ​ 表示 PachGAN的判别器损失, L r e g \mathcal{L}_{r e g} Lreg​ 表示KL正则损失。

UNet训练

在我们进行Stable Diffusion模型训练时,VAE部分和CLIP部分都是冻结的,所以说官方在训练SD系列模型的时候,训练过程一般主要训练U-Net部分。
Stable Diffusion中U-Net主要是进行噪声残差,在SD系列模型训练时和DDPM一样采用预测噪声残差的方法来训练U-Net,其损失函数如下所示:
L S D = E x 0 , ϵ ∼ N ( 0 , I ) , t [ ∥ ϵ − ϵ θ ( α ˉ t x 0 + 1 − α ˉ t ϵ , t , c ) ∥ 2 ] L_{S D}=\mathbb{E}_{\mathbf{x}_0, \epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I}), t}\left[\left\|\epsilon-\epsilon_\theta\left(\sqrt{\bar{\alpha}_t} \mathbf{x}_0+\sqrt{1-\bar{\alpha}_t} \epsilon, t, \mathbf{c}\right)\right\|^2\right] LSD​=Ex0​,ϵ∼N(0,I),t​[ ​ϵ−ϵθ​(αˉt​ ​x0​+1−αˉt​ ​ϵ,t,c) ​2]
这里的 c c c为Text Embeddings。

CLIP 训练

在SD模型的训练中,一般来说CLIP的整体性能是足够支撑我们的下游细分任务的,所以CLIP Text Encoder模型参数是冻结的,我们不需要对其重新训练。

由于CLIP训练时所采用的最大Token数是77,所以在SD模型进行前向推理时,当输入Prompt的Token数量超过77时,将通过Clip操作拉回77x768,而如果Token数不足77则会使用padding操作得到77x768。如果说全卷积网络的设计让图像输入尺寸不再受限,那么CLIP的这个设置就让输入的文本长度不再受限(可以是空文本)。无论是非常长的文本,还是空文本,最后都将得到一样维度的特征矩阵。

Stable Diffusion XL

与Stable Diffusion 1.x-2.x相比,Stable Diffusion XL主要进行如下的优化:

对Stable Diffusion 1.x-2.x的U-Net,VAE,CLIP Text Encoder三大核心模型都做了改进。 增加一个独立的基于Latent的Refiner模型,也是一个扩散模型,用来提升生成图像的精细化程度。 设计了很多训练Tricks,包括图像尺寸条件化策略、图像裁剪参数条件化策略以及多尺度训练策略等。 先发布Stable Diffusion XL 0.9测试版本,基于用户的使用体验和图片生成的反馈情况,针对性增加数据集和使用RLHF(Reinforcement Learning from Human Feedback,基于人类反馈的强化学习)技术优化训练后,推出了Stable Diffusion XL 1.0正式版。

SDXL 能够生成1024x1024分辨率及以上的高质量图片。
比起Stable Diffusion 1.x-2.x,Stable Diffusion XL的参数量增加到了66亿(Base模型35亿+Refiner模型31亿),并且先后发布了模型结构完全相同的0.9和1.0两个版本。Stable Diffusion XL 1.0在0.9版本上使用更多训练集+RLHF来优化生成图像的色彩、对比度、光线以及阴影方面,使得生成图像的构图比0.9版本更加鲜明准确。

上表是Stable Diffusion XL与之前的Stable Diffusion系列的对比,从中可以看出,Stable Diffusion 1.x的U-Net参数量只有860M,就算是Stable Diffusion 2.x,其参数量也不过865M。但等到Stable Diffusion XL,U-Net模型(Base部分)参数量就增加到2.6B,参数量增加幅度达到了3倍左右。

基础架构

Stable Diffusion XL是一个二阶段的级联扩散模型(Latent Diffusion Model),包括Base模型和Refiner模型。

Base模型的主要工作和Stable Diffusion 1.x-2.x一致,具备文生图(txt2img)、图生图(img2img)、图像inpainting等能力。
SDXL Base模型由U-Net、VAE以及CLIP Text Encoder(两个)三个模块组成。

在FP16精度下Base模型大小6.94G(FP32:13.88G),其中U-Net占5.14G、VAE模型占167M以及两个CLIP Text Encoder一大一小(OpenCLIP ViT-bigG和OpenAI CLIP ViT-L)分别是1.39G和246M。

在Base模型之后,级联了Refiner模型,对Base模型生成的图像Latent特征进行精细化提升,其本质上是在做图生图的工作。
SDXL Refiner模型同样由U-Net、VAE和CLIP Text Encoder(一个)三个模块组成.

在FP16精度下Refiner模型大小6.08G,其中U-Net占4.52G、VAE模型占167M(与Base模型共用)以及CLIP Text Encoder模型(OpenCLIP ViT-bigG)大小1.39G(与Base模型共用)。

模型结构

Base 模型

VAE 模型

Stable Diffusion XL依旧是基于Latent的扩散模型,所以VAE的Encoder和Decoder结构依旧是Stable Diffusion XL提取图像Latent特征和图像像素级重建的关键。

当输入是 图片
Stable Diffusion XL和Stable Diffusion一样,首先会使用VAE的Encoder结构将输入图像转换为Latent特征,然后U-Net不断对Latent特征进行优化,最后使用VAE的Decoder结构将Latent特征重建出像素级图像。
除了提取Latent特征和图像的像素级重建外,VAE还可以改进生成图像中的高频细节,小物体特征和整体图像色彩。 当输入是 文字
这时我们不需要VAE的Encoder结构,只需要Decoder进行图像重建。

SDXL和SD模型的区别:
Stable Diffusion XL使用了和之前Stable Diffusion系列一样的VAE结构(KL-f8),但在训练中选择了更大的Batch-Size(256 vs 9),并且对模型进行指数滑动平均操作(EMA,exponential moving average),EMA对模型的参数做平均,从而提高性能并增加模型鲁棒性。

与此同时,与Stable Diffusion一样,VAE模型在Stable Diffusion XL中除了能进行图像压缩和图像重建的工作外,通过切换不同微调训练版本的VAE模型,能够改变生成图片的细节与整体颜色(更改生成图像的颜色表现,类似于色彩滤镜)。

U-Net 模型
GSC模块:Stable Diffusion Base XL U-Net中的最小组件之一,由GroupNorm+SiLU+Conv三者组成。
DownSample模块:Stable Diffusion Base XL U-Net中的下采样组件,使用了Conv(kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))进行采下采样。
UpSample模块:Stable Diffusion Base XL U-Net中的上采样组件,由插值算法(nearest)+Conv组成。
ResNetBlock模块:借鉴ResNet模型的“残差结构”,让网络能够构建的更深的同时,将Time Embedding信息嵌入模型。
CrossAttention模块:将文本的语义信息与图像的语义信息进行Attention机制,增强输入文本Prompt对生成图像的控制。

Text Condition是三维的,而Latent Feature是四维的,那它们是怎么进行Attention机制的呢?
其实在每次进行Attention机制前,我们需要将Latent Feature从[batch_size,channels,height,width]转换到[batch_size,height*width,channels] ,这样就变成了三维特征,就能够和Text Condition做CrossAttention操作。
在完成CrossAttention操作后,我们再将Latent Feature从[batch_size,height*width,channels]转换到[batch_size,channels,height,width] ,这样就又重新回到原来的维度。

Text Condition如何跟latent Feature大小保持一致呢?
因为latent embedding不同位置的H和W是不一样的,但是Text Condition是从文本中提取的,其H和W是固定的。这里在CorssAttention模块中有一个非常巧妙的点,那就是在不同特征做Attention操作前,使用Linear层将不同的特征的尺寸大小对齐。

SelfAttention模块:SelfAttention模块的整体结构与CrossAttention模块相同,这是输入全部都是图像信息,不再输入文本信息。

FeedForward模块:Attention机制中的经典模块,由GeGlU+Dropout+Linear组成。

BasicTransformer Block模块:由LayerNorm+SelfAttention+CrossAttention+FeedForward组成,是多重Attention机制的级联,并且每个Attention机制都是一个“残差结构”。通过加深网络和多Attention机制,大幅增强模型的学习能力与图文的匹配能力。

SDXL_Spatial Transformer_X模块:由GroupNorm+Linear+X个BasicTransformer Block+Linear构成,同时也使用了ResNet模型的“残差结构”。

SDXL_DownBlock模块:由两个ResNetBlock+一个DownSample组成。

SDXL_UpBlock_X模块:由X个ResNetBlock模块组成。

CrossAttnDownBlock_X_K模块:是Stable Diffusion XL Base U-Net中Encoder部分的主要模块,由K个(ResNetBlock模块+SDXL_Spatial Transformer_X模块)+一个DownSample模块组成。

CrossAttnUpBlock_X_K模块:是Stable Diffusion XL Base U-Net中Decoder部分的主要模块,由K个(ResNetBlock模块+SDXL_Spatial Transformer_X模块)+一个UpSample模块组成。

CrossAttnMidBlock模块:是Stable Diffusion XL Base U-Net中Encoder和ecoder连接的部分,由ResNetBlock+SDXL_Spatial Transformer_10+ResNetBlock组成。

总结

其中增加的SDXL_Spatial Transformer_X模块(主要包含Self Attention + Cross Attention + FeedForward)数量占新增参数量的主要部分。 U-Net的Encoder和Decoder结构也从之前系列的4stage改成3stage([1,1,1,1] -> [0,2,10]),同时SDXL只使用两次下采样和上采样,而之前的SD系列模型都是三次下采样和上采样。 在SDXL U-Net的Encoder结构中,包含了两个CrossAttnDownBlock结构和一个SDXL_DownBlock结构;在Decoder结构中,包含了两个CrossAttnUpBlock结构和一个SDXL_UpBlock结构;与此同时,Encoder和Decoder中间存在Skip Connection,进行信息的传递与融合。

比起Stable Diffusion 1.x-2.x,Stable Diffusion XL在第一个stage中不再使用Spatial Transformer Blocks,而在第二和第三个stage中大量增加了Spatial Transformer Blocks(分别是2和10),那么这样设计有什么好处呢?

首先,在第一个stage中不使用SDXL_Spatial Transformer_X模块,可以明显减少显存占用和计算量。 然后在第二和第三个stage这两个维度较小的feature map上使用数量较多的SDXL_Spatial Transformer_X模块,能在大幅提升模型整体性能(学习能力和表达能力)的同时,优化了计算成本。
CLIP 模型

Stable Diffusion XL模型采用的Text Encoder依然是基于CLIP架构的。
CLIP模型主要包含Text Encoder和Image Encoder两个模块,Stable Diffusion 1.x系列使用的是OpenAI CLIP ViT-L/14(123.65M)中的Text Encoder模型,而Stable Diffusion 2.x系列则使用OpenCLIP ViT-H/14(354.03M)中的Text Encoder模型。
Stable Diffusion XL和Stable Diffusion 1.x-2.x系列一样,只使用Text Encoder模块从文本信息中提取Text Embeddings。

不同的是,Stable Diffusion XL与之前的系列相比使用了两个CLIP Text Encoder,分别是OpenCLIP ViT-bigG(694M)OpenAI CLIP ViT-L/14(123.65M),从而大大增强了Stable Diffusion XL对文本的提取和理解能力,同时提高了输入文本和生成图片的一致性。

SDXL OpenCLIP ViT-bigG的结构图

SDXL OpenAI CLIP ViT-L/14的结构图

由上面两个结构对比可知,OpenCLIP ViT-bigG的优势在于模型结构更深,特征维度更大,特征提取能力更强,但是其两者的基本CLIPEncoder模块是一样的。

与传统深度学习中的模型融合类似,Stable Diffusion XL分别提取两个Text Encoder的倒数第二层特征,并进行concat操作作为文本条件(Text Conditioning)。
其中OpenCLIP ViT-bigG的特征维度为77x1280,而OpenAI CLIP ViT-L/14的特征维度是77x768,所以输入总的特征维度是77x2048(77是最大的token数,2048是SDXL的context dim),再通过Cross Attention模块将文本信息传入Stable Diffusion XL的训练过程与推理过程中。

和Stable Diffusion 1.x-2.x一致的是,Stable Diffusion XL输入的最大Token数依旧是77,当输入文本的Token数量超过77后,将通过Clip操作拉回77x2048;如果Token数不足77则会通过padding操作得到77x2048。

与此同时,Stable Diffusion XL还提取了OpenCLIP ViT-bigGpooled text embedding,将其嵌入到Time Embeddings中(add操作),作为辅助约束条件(强化文本的整体语义信息),但是这种辅助条件的强度是较为微弱的。

Refiner 模型

已经有U-Net(Base)模型生成了图像的Latent特征,所以Refiner模型的主要工作是在Latent特征进行小噪声去除和细节质量提升。
Refiner模型和Base模型一样是基于Latent的扩散模型,也采用了Encoder-Decoder结构,和U-Net兼容同一个VAE模型。
不过在Text Encoder部分,Refiner模型只使用了OpenCLIP ViT-bigGText Encoder,同样提取了倒数第二层特征以及进行了pooled text embedding的嵌入。

Refiner模型主要做了图像生成图像(img2img)的工作,其具备很强的迁移兼容能力,可以作为Stable Diffusion、Midjourney、DALL-E、GAN、VAE等生成式模型的级联组件,

Refiner的模型结构

SDXL Refiner模型和SDXL Base模型在结构上的异同:

SDXL Base的Encoder和Decoder结构都采用4个stage,而SDXL Base设计的是3个stage。 SDXL Refiner和SDXL Base一样,在第一个stage中没有使用Attention模块。 在经过第一个卷积后,SDXL Refiner设置初始网络特征维度为384,而SDXL Base 采用的是320。 SDXL Refiner的Attention模块中SDXL_Spatial Transformer结构数量均设置为4。 SDXL Refiner的参数量为2.3B,比起SDXL Base的2.6B参数量略小一些。

训练和推理

基本训练流程

VAE
下表是Stable Diffusion XL的VAE在COCO2017 验证集上,图像大小为256×256像素的情况下的性能。
表中的三个VAE模型结构是一样的,不同点在于SD 2.x VAE是基于SD 1.x VAE微调训练了Decoder部分,同时保持Encoder部分权重不变,使他们有相同的Latent特征分布,所以SD 1.x和SD 2.x的VAE模型是互相兼容的。而SDXL VAE是重新从头开始训练的,所以其Latent特征分布与之前的两者不同。

由于Latent特征分布产生了变化,SDXL VAE的缩放系数也产生了变化。VAE在将Latent特征送入U-Net之前,需要对Latent特征进行缩放让其标准差尽量为1,之前的Stable Diffusion系列采用的缩放系数为0.18215,由于Stable Diffusion XL的VAE进行了全面的重训练,所以缩放系数重新设置为0.13025。

由于缩放系数的改变,Stable Diffusion XL VAE模型与之前的Stable Diffusion系列并不兼容。如果在SDXL上使用之前系列的VAE,会生成充满噪声的图片。

CLIP
和之前的系列一样,SDXL Text Encoder在官方训练时是冻结的,我们在对SDXL模型进行微调训练时,可以同步开启Text Encoder的微调训练,能够使得Text Encoder对生成图片的控制力增强,使其生成内容更加贴近训练集的分布。

Refiner
SDXL Refiner模型的训练逻辑与SDXL Base一样,不过Refiner模型只在前200个Timesteps上训练(设置的noise level较低)。

训练技巧

1) 图像尺寸条件化

之前在Stable Diffusion的训练过程中,主要分成两个阶段:

一个是在256x256的图像尺寸上进行预训练 然后在512x512的图像尺寸上继续训练。

而这两个阶段的训练过程都要对最小图像尺寸进行约束。第一阶段中,会将尺寸小于256x256的图像舍弃;同样的,在第二阶段,会将尺寸小于512x512的图像筛除。这样的约束会导致训练数据中的大量数据被丢弃,从而很可能导致模型性能和泛化性的降低。
下图展示了如果将尺寸小于256x256的图像筛除,整个数据集将减少39%的数据。如果加上尺寸小于512x512的图像,未利用数据占整个数据集的百分比将更大。

针对上述数据集利用率的问题,常规思路可以借助超分模型将尺寸过小的图像放大。但是面对对于图像尺寸过小的场景,目前的超分模型可能会在对图像超分的同时会引入一些噪声伪影,影响模型的训练,导致生成一些模糊的图像。

Stable Diffusion XL为了在解决数据集利用率问题的同时不引入噪声伪影,将U-Net(Base)模型与原始图像分辨率相关联,核心思想是将输入图像的原始高度和宽度作为额外的条件嵌入U-Net模型中,表示为 C s i z e = ( h e i g h t , w i d t h ) C_{size}=(height,width) Csize​=(height,width)
heightwidth都使用傅里叶特征编码进行独立嵌入,然后将特征concat后加在Time Embedding上,将图像尺寸作为条件引入训练过程。
这样以来,模型在训练过程中能够学习到图像的原始分辨率信息,从而在推理生成阶段更好地适应不同尺寸的图像生成,而不会产生噪声伪影的问题。

如下图所示,在使用了图像尺寸条件化策略后,Base模型已经对不同图像分辨率有了“自己的判断”。当输入低分辨率条件时,生成的图像较模糊;在不断增大分辨率条件时,生成的图像质量不断提升。

2) 图像裁剪参数条件化

之前的Stable Diffusion系列模型,由于需要输入固定的图像尺寸用作训练,很多数据在预处理阶段会被裁剪。生成式模型中典型的预处理方式是先调整图像尺寸,使得最短边与目标尺寸匹配,然后再沿较长边对图像进行随机裁剪或者中心裁剪。
虽然裁剪是一种数据增强方法,但是训练中对图像裁剪导致的图像特征丢失,可能会导致在图像生成过程中出现不符合训练数据分布的特征。

下图中展示了SD 1.4和SD 1.5的经典失败案例,生成图像中的猫出现了头部缺失的问题,龙也出现了体征不完整的情况:

其实之前NovelAI就发现了这个问题,并提出了基于分桶(Ratio Bucketing)的多尺度训练策略
其主要思想是先将训练数据集按照不同的长宽比(aspect ratio)进行分组(groups)或者分桶(buckets)
在训练过程中,每次在buckets中随机选择一个bucket并从中采样Batch个数据进行训练。将数据集进行分桶可以大量较少裁剪图像的操作,并且能让模型学习多尺度的生成能力;但相对应的,预处理成本大大增加,特别是数据量级较大的情况下。
并且尽管数据分桶成功解决了数据裁剪导致的负面影响,但如果能确保数据裁剪不把负面影响引入生成过程中,裁剪这种数据增强方法依旧能给模型增强泛化性能。

Stable Diffusion XL使用了一种简单而有效的条件化方法,即图像裁剪参数条件化策略。
其主要思想是在加载数据时,将左上角的裁剪坐标通过傅里叶编码后加在Time Embedding上,并嵌入U-Net(Base)模型中,并与原始图像尺寸一起作为额外的条件嵌入U-Net模型,从而在训练过程中让模型学习到对“图像裁剪”的认识。

从下图中可以看到,将不同的 c c r o p c_{crop} ccrop​ 坐标条件的生成图像进行了对比,当我们设置 c c r o p = ( 0 , 0 ) c_{crop}=(0,0) ccrop​=(0,0)时可以生成主要物体居中并且无特征缺失的图像,而采用其它的坐标条件则会出现有裁剪效应的图像:

图像尺寸条件化策略和图像裁剪参数条件化策略都能在SDXL训练过程中使用(在线方式应用),同时也可以很好的迁移到其他AIGC生成式模型的训练中。下图详细给出了两种策略的通用使用流程:

可以看到,SDXL在训练过程中的数据处理流程和之前的系列是一样的,只是需要再将图像原始长宽(width和height)以及图像进行crop操作时的左上角的裁剪坐标top和left作为条件输入。

3) 多尺度训练

Stable Diffusion XL采用了多尺度训练策略,并在此基础上,增加了分桶策略。

Stable Diffusion XL首先采用图像尺寸条件化和图像裁剪参数条件化这两种策略在256x256和512x512的图像尺寸上分别预训练600000步和200000步(batch size = 2048),总的数据量约等于 (600000 + 200000) x 2048 = 16.384亿。

接着Stable Diffusion XL在1024x1024的图像尺寸上采用多尺度方案来进行微调,并将数据分成不同纵横比的桶(bucket),并且尽可能保持每个桶的像素数接近1024×1024,同时相邻的bucket之间height或者width一般相差64像素左右,Stable Diffusion XL的具体分桶情况如下图所示: 其中Aspect Ratio = Height / Width,表示高宽比。

在训练过程中,一个Batch从一个桶里的图像采样,并且我们在每个训练步骤(step)中可以在不同的桶之间交替切换。
除此之外,Aspect Ratio也会作为条件嵌入到U-Net(Base)模型中,嵌入方式和上面提到的其他条件嵌入方式一致,让模型能够更好地学习到“多尺度特征”。
与此同时,SDXL在多尺度微调阶段依然使用图像裁剪参数条件化策略,进一步增强SDXL对图像裁剪的敏感性。

在完成了多尺度微调后,SDXL就可以进行不同Aspect Ratio的图像生成了,不过官方推荐生成尺寸默认为1024x1024。

4) 使用Offset Noise

在SDXL进行微调时,使用了Offset Noise操作,能够让SDXL生成的图像有更高的色彩自由度(纯黑或者纯白背景的图像)。

SD 1.x和SD 2.x一般只能生成中等亮度的图片,即生成平均值相对接近 0.5 的图像(全黑图像为 0,全白图像为 1),之所以会出现这个问题,是因为SD系列模型训练和推理过程的不一致造成的。
SD模型在训练中进行noise scheduler流程并不能将图像完全变成随机高斯噪声,但是推理过程中,SD模型是从一个随机高斯噪声开始生成的,因此就会存在训练与推理的噪声处理过程不一致。

Offset Noise操作是解决这个问题的一种直观并且有效的方法,我们只需要在SD模型的微调训练时,把额外从高斯分布中采样的偏置噪声引入图片添加噪声的过程中,这样就对图像的色彩均值造成了破坏,从而提高了SDXL生成图像的“泛化性能”。

5) SDXL的条件注入与训练细节

上面我们已经详细讲解了SDXL的四个额外的条件信息注入(pooled text embedding图像尺寸条件图像裁剪参数条件图像多尺寸条件
其中三个图像条件可以像Timestep一样采用傅立叶编码得到Embedding特征,然后再和pooled text embedding特征concat,得到维度为2816的embeddings特征。

SDXL在训练中的配置。和SD 1.x系列一样,SDXL在训练时采用了1000步的DDPM和相同的noise scheduler,同时依旧采用基于预测noise的损失函数,和SD 1.x系列一致。

推理

在Stable Diffusion XL推理阶段,输入一个prompt,通过VAE和U-Net(Base)模型生成Latent特征,接着给这个Latent特征进行扩散过程加上一定的噪音。在此基础上,再使用Refiner模型进行去噪,以提升图像的整体质量与局部细节。

参考

论文
【1】 DDPM的灵感来源 Deep Unsupervised Learning using Nonequilibrium Thermodynamics
【2】DDPM 开山之作 : Denoising Diffusion Probabilistic Models (DDPM)
【3】DDPM的改进,提出余弦调度策略 Improved Denoising Diffusion Probabilistic Models
【4】DDPM采用的最基础的Unet架构U-Net: Convolutional Networks for Biomedical Image Segmentation

博文
【b.1】通过可视化的动画,逐步讲解 Step by Step Visual Introduction to Diffusion Models
【b.2】主页里有对应的ppt和yotube链接 cvpr2022-tutorial-diffusion-models
【b.3】偏向于数学的角度介绍 How diffusion models work: the math from scratch
【b.4】Arxiv24最新的一篇关于DDPM非常详细的讲解 Tutorial on Diffusion Models for Imaging and Vision
【b.5】介绍的非常全面,包括具体的数学推导What are Diffusion Models?
【b.6】Stable Diffusion 图解 The Illustrated Stable Diffusion
【b.7】强烈推荐知乎上的一个讲解 深入浅出完整解析Stable Diffusion(SD)核心基础知识 & 深入浅出完整解析Stable Diffusion XL(SDXL)核心基础知识

视频
【1】李宏毅老师的系列讲解
【生成式AI】淺談圖像生成模型 Diffusion Model 原理
【生成式AI】Stable Diffusion、DALL-E、Imagen 背後共同的套路
【生成式AI】Diffusion Model 原理剖析 (1/4) (optional)
【生成式AI】Diffusion Model 原理剖析 (2/4) (optional)
【生成式AI】Diffusion Model 原理剖析 (3/4) (optional)
【生成式AI】Diffusion Model 原理剖析 (4/4) (optional)

【2】Yotube上的视频,不仅介绍了DDPM基础还有应用
Tutorial on Denoising Diffusion-based Generative Modeling: Foundations and Applications

更新时间 2024-06-08