StableDiffusion:基于LatentDiffusionModels的文图生成新方案

文章主题：Stable Diffusion, Latent Diffusion Models, 文图生成, 潜在表示空间

引言

最近大火的AI作画吸引了很多人的目光，AI作画近期取得如此巨大进展的原因个人认为有很大的功劳归属于Stable Diffusion的开源。Stable diffusion是一个基于Latent Diffusion Models（潜在扩散模型，LDMs）的文图生成（text-to-image）模型。具体来说，得益于Stability AI的计算资源支持和LAION的数据资源支持，Stable Diffusion在LAION-5B的一个子集上训练了一个Latent Diffusion Models，该模型专门用于文图生成。

Latent Diffusion Models通过在一个潜在表示空间中迭代“去噪”数据来生成图像，然后将表示结果解码为完整的图像，让文图生成能够在消费级GPU上，在10秒级别时间生成图片，大大降低了落地门槛，也带来了文图生成领域的大火。所以，如果你想了解Stable Diffusion的背后原理，可以跟我一起深入解读一下其背后的论文High-Resolution Image Synthesis with Latent Diffusion Models（Latent Diffusion Models），同时这篇文章后续也会针对ppdiffusers的相关代码进行讲解。该论文发表于CVPR2022，第一作者是Robin Rombach，来自德国慕尼黑大学机器视觉与学习研究小组。

Stable Diffusion

再解读论文之前，首先让我深入了解一下Stable Diffusion。

Stable Diffusion是一款专为文本到图像生成任务设计的模型，其基础在于Latent Diffusion Models。目前，该模型已经推出了v1版本，也就是Stable Diffusion v1。作为Latent Diffusion Models的具体实现，Stable Diffusion v1具有特定的模型架构设定，其中自动编码器的下采样因子为8，UNet的大小为860M，而文本编码器则采用了CLIP ViT-L/14。至于官方目前提供的权重信息，如下所示：（接下来列出官方提供的权重信息）

sd-v1-1.ckpt: 237k steps at resolution 256×256 on laion2B-en. 194k steps at resolution 512×512 on laion-high-resolution (170M examples from LAION-5B with resolution >= 1024×1024).sd-v1-2.ckpt: Resumed from sd-v1-1.ckpt. 515k steps at resolution 512×512 on laion-aesthetics v2 5+ (a subset of laion2B-en with estimated aesthetics score > 5.0, and additionally filtered to images with an original size >= 512×512, and an estimated watermark probability < 0.5. The watermark estimate is from the LAION-5B metadata, the aesthetics score is estimated using the LAION-Aesthetics Predictor V2).sd-v1-3.ckpt: Resumed from sd-v1-2.ckpt. 195k steps at resolution 512×512 on “laion-aesthetics v2 5+” and 10% dropping of the text-conditioning to improve classifier-free guidance sampling.sd-v1-4.ckpt: Resumed from sd-v1-2.ckpt. 225k steps at resolution 512×512 on “laion-aesthetics v2 5+” and 10% dropping of the text-conditioning to improve classifier-free guidance sampling.

论文贡献

Diffusion model相比GAN可以取得更好的图片生成效果，然而该模型是一种自回归模型，需要反复迭代计算，因此训练和推理代价都很高。论文提出一种在潜在表示空间（latent space）上进行diffusion过程的方法，从而能够大大减少计算复杂度，同时也能达到十分不错的图片生成效果。相比于其它空间压缩方法（如），论文提出的方法可以生成更细致的图像，并且在高分辨率图片生成任务（如风景图生成，百万像素图像）上表现得也很好。论文将该模型在无条件图片生成（unconditional image synthesis）, 图片修复（inpainting）,图片超分（super-resolution）任务上进行了实验，都取得了不错的效果。论文还提出了cross-attention的方法来实现多模态训练，使得条件图片生成任务也可以实现。论文中提到的条件图片生成任务包括类别条件图片生成（class-condition）, 文图生成（text-to-image）, 布局条件图片生成（layout-to-image）。这也为日后Stable Diffusion的开发奠定了基础。

方法

Latent Diffusion Models整体框架如图，首先需要训练好一个自编码模型（AutoEncoder，包括一个编码器 $E$ \mathcal{E} 和一个解码器 $D$ \mathcal{D} ）。这样一来，我们就可以利用编码器对图片进行压缩，然后在潜在表示空间上做diffusion操作，最后我们再用解码器恢复到原始像素空间即可，论文将这个方法称之为感知压缩（Perceptual Compression）。个人认为这种将高维特征压缩到低维，然后在低维空间上进行操作的方法具有普适性，可以很容易推广到文本、音频、视频等领域。

扩散操作在潜在空间上的实现方式与标准扩散模型并无太大差异，主要使用的扩散模型为time-conditional UNet。然而，该研究在扩散操作中加入了一个关键特性，即条件机制（Conditioning Mechanisms），并通过cross-attention实现多模态训练，从而使得条件图片生成任务得以实现。

下面我们针对感知压缩、扩散模型、条件机制的具体细节进行展开。

图片感知压缩（Perceptual Image Compression）

感知压缩作为一种权衡策略，在许多扩散模型中并未得到应用。然而，传统的非感知压缩扩散模型存在一个主要问题，即在像素空间上训练模型会导致所需空间维度过高，从而使得生成高分辨率图片变得困难。为了解决这个问题，我们提出了感知压缩的概念，通过采用如VAE的自编码模型对原始图片进行处理，去除图片中的高频信息，并仅保留对图像具有重要意义的基本特征。这一方法的优势正如引文中所述，它能显著降低训练和采样阶段的计算复杂度，使文图生成等任务能在消费级GPU上，在短短10秒内生成图片，从而大幅降低了实现难度。

感知压缩的核心技术是一个预先训练的自编码模型，它能够捕捉到与图像空间感知的等价潜在表示。此种方法的优势在于，仅需训练一个通用的自编码模型，便可应用于多种扩散模型的训练，从而在各种任务中得以运用。因此，感知压缩不仅适用于传统的无条件图像生成，还便于将其扩展至各种图像间的修复任务（如超分辨率）以及文本至图像的转换任务。

显然，基于感知压缩的扩散模型训练可以分为两个阶段。在第一个阶段，我们需要训练自编码器；而在第二个阶段，才是扩散模型的训练过程。在进行自编码器的训练过程中，为了防止潜在表示空间产生显著的差异，研究者们采用了两种正则化手段，即KL正则化和VQ正则化。因此，在官方发布的训练模型中，我们可以看到这两种正则化的实现方式。在Stable Diffusion项目中，主要采用的是AutoencoderKL的实现方式。

具体来说，给定图像 $x\inRH\timesW\times3$ x \in \mathbb{R}^{H \times W \times 3} ，我们可以先利用一个编码器 $E$ \mathcal{E} 来将图像编码到潜在表示空间 $z=E(x)$ z=\mathcal{E}(x) ，其中 $z\inRh\timesw\timesc$ z \in \mathbb{R}^{h \times w \times c} ，然后再用解码器从潜在表示空间重建图片 $x~=D(z)=D(E(x))$ \tilde{x}=\mathcal{D}(z)=\mathcal{D}(\mathcal{E}(x)) 。在感知压缩压缩的过程中，下采样因子的大小为 $f=H/h=W/w$ f=H / h=W / w，它是2的次方，即 $f=2m$ f=2^m 。

潜在扩散模型（Latent Diffusion Models）

首先简要介绍一下普通的扩散模型（DM），扩散模型可以解释为一个时序去噪自编码器（equally weighted sequence of denoising autoencoders） $ϵθ(xt,t);t=1\dotsT$ \epsilon_\theta\left(x_t, t\right) ; t=1 \ldots T ，其目标是根据输入 $xt$ x_t 去预测一个对应去噪后的变体，或者说预测噪音，其中 $xt$ x_t是输入 $x$ x 的噪音版本。相应的目标函数可以写成如下形式：

$LDM=Ex,ϵ\simN(0,1),t[‖ϵ-ϵθ(xt,t)‖22]$ L_{D M}=\mathbb{E}_{x, \epsilon \sim \mathcal{N}(0,1), t}\left[\left\|\epsilon-\epsilon_\theta\left(x_t, t\right)\right\|_2^2\right]

其中 $t$ t 从 ${1,\dots,T}$ \{1, \ldots, T\} 中均匀采样获得。

而在潜在扩散模型中，引入了预训练的感知压缩模型，它包括一个编码器 $E$ \mathcal{E} 和一个解码器 $D$ \mathcal{D}。这样就可以利用在训练时就可以利用编码器得到 $zt$ z_t，从而让模型在潜在表示空间中学习，相应的目标函数可以写成如下形式：

$LLDM:=EE(x),ϵ\simN(0,1),t[‖ϵ-ϵθ(zt,t)‖22]$ L_{L D M}:=\mathbb{E}_{\mathcal{E}(x), \epsilon \sim \mathcal{N}(0,1), t}\left[\left\|\epsilon-\epsilon_\theta\left(z_t, t\right)\right\|_2^2\right]

条件机制（Conditioning Mechanisms）

除了无条件图片生成外，我们也可以进行条件图片生成，这主要是通过拓展得到一个条件时序去噪自编码器（conditional denoising autoencoder） $ϵθ(zt,t,y)$ \epsilon_\theta\left(z_t, t, y\right) 来实现的，这样一来我们就可通过 $y$ y 来控制图片合成的过程。具体来说，论文通过在UNet主干网络上增加cross-attention机制来实现 $ϵθ(zt,t,y)$ \epsilon_\theta\left(z_t, t, y\right)。为了能够从多个不同的模态预处理 $y$ y ，论文引入了一个领域专用编码器（domain specific encoder） $τθ$ \tau_\theta ，它用来将 $y$ y 映射为一个中间表示 $τθ(y)\inRM\timesdτ$ \tau_\theta(y) \in \mathbb{R}^{M \times d_\tau} ，这样我们就可以很方便的引入各种形态的条件（文本、类别、layout等等）。最终模型就可以通过一个cross-attention层映射将控制信息融入到UNet的中间层，cross-attention层的实现如下：

$Attention(Q,K,V)=softmax(QKTd)\cdotV,with$ Attention (Q, K, V)=\operatorname{softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) \cdot V, with

$Q=WQ(i)\cdotφi(zt),K=WK(i)\cdotτθ(y),V=WV(i)\cdotτθ(y)$ Q=W_Q^{(i)} \cdot \varphi_i\left(z_t\right), K=W_K^{(i)} \cdot \tau_\theta(y), V=W_V^{(i)} \cdot \tau_\theta(y)

其中 $φi(zt)\inRN\timesdϵi$ \varphi_i\left(z_t\right) \in \mathbb{R}^{N \times d_\epsilon^i} 是UNet的一个中间表征。相应的目标函数可以写成如下形式：

$LLDM:=EE(x),y,ϵ\simN(0,1),t[‖ϵ-ϵθ(zt,t,τθ(y))‖22]$ L_{L D M}:=\mathbb{E}_{\mathcal{E}(x), y, \epsilon \sim \mathcal{N}(0,1), t}\left[\left\|\epsilon-\epsilon_\theta\left(z_t, t, \tau_\theta(y)\right)\right\|_2^2\right]

实验

在本研究中，我们采用了LDMs模型，并在无条件图片生成任务中使用了LSUN、FFHQ以及CelebA-HQ数据集，而在类别条件图片生成任务中所用的数据集则包括ImageNet，以及在文图生成任务上所涉及的数据集为Conceptual Captions和LAION。为了全面验证模型的性能，我们设计了一系列的对比实验，并分别从感知压缩权衡（下采样因子）以及LDM生成效果对比两个方面进行了深入的分析。此外，我们还针对多种图片生成任务，如Super-Resolution、Inpainting和layout-condition等进行了实验。这些结果充分证明了LDMs模型中学习到的潜在表示空间具有较强的分布拟合能力，能够满足各种下游任务的需求。

感知压缩权衡（Perceptual Compression Tradeoffs）

前面提到过下采样因子 $f$ f 的大小为 $f=H/h=W/w$ f=H / h=W / w，如果 $f=1$ f=1 那就等于没有对输入的像素空间进行压缩，如果 $f$ f 越大，则信息压缩越严重，可能会噪声图片失真，但是训练资源占用的也越少。论文对比了 $f$ f 在分别 ${1,2,4,8,16,32}$ \{1,2, 4, 8, 16, 32\} 下的效果，发现 $f$ f 在 ${4-16}$ \{4-16\} 之间可以比较好的平衡压缩效率与视觉感知效果。作者重点推荐了LDM-4和LDM-8。