High-Resolution Image Synthesis with Latent Diffusion Models翻译

摘要

通过将图像形成过程分解为去噪自编码器的序列级应用，扩散模型（DM）在图像数据及其他领域达到了SOTA的合成结果。此外，它们的形式化方式允许通过一种指导机制来控制图像生成过程而无需重新训练。但是，由于这些模型通常直接在像素空间中进行操作，因此对功能强大的DM的优化通常会消耗数百个GPU天数，并且由于序列评估从而导致推理昂贵。为了使DM能够在有限计算资源下训练，同时保留其质量和灵活性，我们将其应用于预训练自编码器的潜在空间。与以前的工作相反，这种在表示下训练扩散模型能允许首先到达一个复杂度降低和细节保留的最优点，从而大大提高了视觉保真度。以卷积的方式，通过在模型结构中引入cross-attention层，我们将扩散模型转换为以通用条件为输入（例如文本，bounding boxes）的强大且灵活的生成器。我们的潜在扩散模型（LDM）在图像补全和以类别为条件的图像合成上达到SOTA分数，并在各种任务上达到可比的性能，包括无条件的图像生成，文本到图像合成以及超分辨率，同时与基于像素的DM相比，显著减少了计算量。

1.介绍

　　图像合成是最近提出的计算机视觉领域之一，但也是计算量需求最大的领域之一。尤其是复杂，自然场景下的高分辨率图像合成，目前主要由基于似然的大模型来进行生成，该模型由具有数十亿个参数自回归（AR）transformers组成。相反，由于对抗学习过程不容易扩展到建模复杂、多模态的分布，因此已经揭示出GAN结果的大部分仅限于相当有限的可变数据。最近，基于分层去噪自编码器构建的扩散模型已证明在图像合成和其它领域中取得了令人印象深刻的结果，并以类别为条件的图像合成和超分辨率上达到了SOTA。此外，与其他类型的生成模型相比，即使是无条件DM，也可以轻松地应用于诸如图像补全和着色，或笔划合成之类的任务。作为基于似然的模型，它们不会表现出类似GAN的模式崩塌和训练不稳定，并且通过大量利用参数共享，它们可以对自然图像的高复杂分布进行建模，而无需像AR模型一样涉及数十亿个参数。
　　Democratizing High-Resolution Image Synthesis。DM属于基于似然的模型，其模式-覆盖行为使它们容易花费过多的容量（并因此耗费较大计算资源）来建模数据中不可察觉的细节。尽管reweighted variational目标旨在通过下采样最初的去噪步骤来解决这一问题，但DM在计算上仍然是需求的，因为训练和评估此类模型需要在RGB图像的高维空间中重复函数评估（和梯度计算）。例如，训练强大的DM通常需要数百个GPU天（例如[15]中的150-1000 v100天），并且对输入空间的噪声版本进行了重复评估，同时推理也较为昂贵，因此生成50k样本的时间大约需要花费单个A100 GPU 5天。 这对研究界和用户有两个后果：首先，训练这种模型需要仅适用于具有大量计算资源的领域，并留下较大的carbon footprint。其次，评估一个已经训练的模型在时间和空间上也很昂贵，因为同一模型结构必须顺序运行大量步骤（例如[15]中的25-1000个步骤）。
　　为了增加此类功能强大模型的可访问性，同时减少了其大量资源消耗，需要一种方法来降低训练和采样的计算复杂性。因此，减少DM的计算需求而不损害其性能是增强其可访问性的关键。
　　Departure to Latent Space。我们的方法对已经训练的基于像素空间的扩散模型进行了分析：图2显示了一个训练模型速率失真的变化。与任何基于似然的模型一样，学习可以大致分为两个阶段：首先是一个感知压缩（perceptual compression）阶段，它可以消除高频细节，但几乎无法学习语义变化。在第二阶段，实际的生成模型学习数据的语义和组合（语义压缩 semantic compression）。因此，我们的目标是首先找到一个在具有相同感知的情况下，在计算上更合适的空间，在该空间中，我们将训练扩散模型以进行高分辨率的图像合成。
　　遵循常见方式，我们将训练分为两个不同的阶段：首先，我们训练一个自编码器，该编码器提供了低维高效的表示空间，该空间在感知上等同于数据空间。重要的是，与以前的工作相反，我们不需要依靠过多的空间压缩，因为我们在学习的潜在空间中训练DM，这在空间维度方面表现出更好的缩放特性。这种复杂性的降低还提供了单个网络通过在潜在空间中进行高效的图像生成。我们将所得模型称为Latent Diffusion Models (LDM)。
　　这种方法的一个值得注意的优点是，我们仅需训练一次通用自编码器，从而可以用于多个DM训练或探索可能完全不同的任务。这可以有效地探索大量用于各种图像到图像和文本到图像任务的扩散模型。对于后者，我们设计了一个将transformer连接到DM UNet主体的结构，并使用以token为条件的机制，请参见3.3节。
　　总的来说，我们的工作具有下列贡献：

1. 与单纯基于的transformer方法相反，我们的方法比较优雅地缩放到更高维数据，因此（a）在压缩级别上实现，提供比以前的工作更忠实，更详细的重建（见图1）；（b）能高效应用于具有百万像素的高分辨率图像合成。
2. 我们在多个任务（无条件图像合成，图像补全，随机超分辨率）和数据集上实现了可比的性能，同时大大降低了计算成本。与基于像素的扩散方法相比，我们还显着降低了推理成本。
3. 我们表明，与以前学习一个编码器/解码器结构和基于得分的工作相反，我们的方法不需要重构和生成能力之间的微妙权衡。这样可以确保重建极度忠实，并且几乎不需要对潜在空间进行正则化。
4. 我们发现，对于条件密集的任务，例如超分辨率，图像不全和语义合成，我们的模型可以以卷积方式应用，并可以渲染出约$1024^2$ px的图像。
5. 此外，我们设计了一种基于cross-attention的以通用目的为条件的机制，从而实现了多模态训练。我们使用它来训练以类别为条件，文本到图像以及布局到图像的模型。
6. 最后，我们在https://github.com/CompVis/latent-diffusion 开源了预训练的潜在扩散模型和自编码模型，这也许能在除DM训练以外的其它任务被使用。

2.相关工作

Generative Models for Image Synthesis。图像的高维特性给生成式建模带来了巨大挑战。生成对抗网络（GAN）可以高效采样出高分辨率图像，但难以优化且捕获完整的数据分布。相比之下，基于似然的方法强调了良好的密度估计，这使得渲染更加良好。变分自编码器（VAE）和基于FLOW的模型可以有效地合成高分辨率图像，但样本质量无法和GAN相当。尽管自回归模型（ARM）在密度估计中实现了强大的性能，但计算要求的结构和顺序采样过程将其限制为低分辨率图像。由于基于像素的图像表示包含了无需关注的高频细节，因此最大似然训练花费了不必要的能力来对其进行建模，从而导致较长的训练时间。为了扩展到更高的分辨率，一些两阶段的方法使用ARM来建模压缩的潜在图像空间，而不是原始像素。
　　最近，Diffusion Probabilistic Models (DM) 在密度估计和样本质量上达到了SOTA。这些模型的生成能力源于自然拟合到图像类数据的归纳偏见，同时使用UNet作为神经结构的实现。当使用reweighted objective进行训练时，通常可以达到最佳的合成质量。在这种情况下，DM相当于一个有损的压缩器，并允许将图像质量进行平衡以获得压缩能力。但是，在像素空间中评估和优化这些模型具有较低的推理速度和较高训练成本。虽然这可以通过优化采样策略和使用分层方法来部分解决，但高分辨率图像数据的训练总是需要计算昂贵的梯度。我们使用所提出的LDM来解决这两个缺点，LDM在较低维度压缩的潜在空间上工作。这使得训练在计算上更便宜，并加快了推理，几乎没有降低合成质量（见图1）。
　　Two-Stage Image Synthesis。为了缓解由单个生成模型所带来的缺点，大部分研究通过两阶段方法将更高效，更好的不同方法的优势结合到模型中。VQ-VAE使用自回归模型在离散的潜在空间上学习一个具有丰富表达能力的先验。[64]通过学习离散的图像和文本表示，将这种方法扩展到文本到图像的生成。更一般而言，[68]使用条件可逆网络在不同领域的潜在空间之间提供通用迁移。与VQ-VAE不同，VQGAN采用具有对抗和感知目标的第一阶段，将自回归transformers扩展到较大的图像。但是，高压缩率需要合理的ARM训练，这引入了数十亿可训练的参数，从而限制了这种方法的整体性能和较少的压缩，这是以高计算成本的价格为代价的。我们的工作避免了这种权衡，因为由于其卷积结构，我们提出的LDM尺度更加轻量缩放到高维潜在空间。因此，我们可以在第一个学习阶段自由选择压缩等级，而不将过多的感知压缩任务交给生成扩散模型，同时保证了高还原的重构（见图1）。尽管存在使用基于得分的先验方法来联合学习编码/解码模型，但它们仍然需要在重建和生成能力之间进行困难的平衡，并且我们的方法表现的更优（第4节）。

3.Method

　　为了降低训练扩散模型对高分辨率图像合成的计算需求，我们观察到，尽管扩散模型允许通过下采样相应的损失项来忽略感知上无关的细节，但它们仍然需要在像素空间上进心昂贵的函数评估，这会导致较大的计算时间和能源消耗需求。
　　我们提出通过引入与生成学习阶段分离的显式压缩来规避这一缺点（见图2）。为了实现这一目标，我们使用了一个自编码模型，该模型在感知上等同于图像空间，但具有显着降低的计算复杂性。
　　这样的方法提供了几种优势：（i）通过远离高维图像空间，我们获得了计算上更有效的DM，因为采样是在低维空间上进行的。（ii）我们利用了从UNet架构继承的DM的归纳偏差，这使其对具有空间结构的数据特别有效，因此减轻了以前方法中所要求的压缩等级的需求。（iii）最后，我们获得了通用目的压缩模型，其潜在空间可用于训练多个生成模型，并且也可以用于其他下游应用，例如单图像CLIP引导的合成。

3.1 Perceptual Image Compression

　　我们的感知压缩模型是基于先前的工作，由一个自编码器组成，该自编码器通过将感知损失和基于patch的对抗目标相结合而训练。这可以通过强迫局部关注并避免由像素空间损失（例如$L_1$和$L_2$损失目标）引入的模糊，从而将重构约束到图像流形。
　　更准确地说，给定一个在RGB空间中的图像$x\in {\mathbb{R}}^{H\times W\times 3}$，编码器$\mathcal{E}$将$x$编码到潜在表示$z=\mathcal{E}(x)$中，解码器$\mathcal{D}$从潜在空间中重建图像，即$\tilde{x}=\mathcal{D}(z)=\mathcal{D}(\mathcal{E}(x))$，其中$z\in {\mathbb{R}}^{h\times w\times c}$。重要的是，编码器通过一个因子$f=H/h=W/w$对图像进行下采样，我们研究了不同的下采样因子$f=2^m,m\in \mathbb{N}$。
　　为了避免高度变化的潜在空间，我们尝试了两种不同类型的正则化。第一个变体是$KL-reg$。对学到的潜在表示施加一个朝正态分布的轻微的KL惩罚，类似于VAE，第二个是$VQ-reg$，在解码器內使用了一个矢量量化。该模型可以解释为VQGAN，但在解码器中使用了量化层。由于我们后续的DM旨在与我们学到的潜在空间$z=\mathcal{E}(x)$的二维结构一起工作，因此我们可以使用相对温和的压缩率并实现非常好的重构。这与以前的工作相反，这些工作依赖于学习空间$z$的任意1D排序，以自动缩放其分布，从而忽略了$z$的固有结构。因此，我们的压缩模型可以更好地保留$x$的详细信息（请参见表8）。完整的训练目标和细节可以在附录中看到。

3.2 Latent Diffusion Models

Diffusion Models。扩散模型是概率模型，旨在通过逐渐去噪正态分布变量来学习一个数据分布$p(x)$，这相当于学习具有固定长度$T$的马尔可夫链的反向过程。对于图像合成，最成功的模型依赖对$p(x)$的变分下界的重构变体，这相当于去噪score-matching。这些模型可以被解释为去噪自编码器${ϵ}_{\theta }(x_t,t);t=1,...,T$同等权重的序列，其被训练以预测其输入$x_t$的去噪变体，其中$x_t$是输入$x$的加噪版本。相应的目标可以简化为（附录B）：
$$L_{DM}={\mathbb{E}}_{x,ϵ\sim \mathcal{N}(0,1),t}[||ϵ-{ϵ}_{\theta }(x_t,t)|{|}_2^2],\text{\hspace{1em}(1)}$$
其中$t$从 { 1 , . . . , T } \{1,...,T\} {1,...,T}中均匀采样。
　　Generative Modeling of Latent Representations。通过我们训练好的由$\mathcal{E}$和$\mathcal{D}$组成的感知压缩模型，我们现在可以更高效访问低维的潜在空间，在该空间中，高频、不可感知的细节被抽象出来。与高维像素空间相比，此空间更适合基于似然的生成模型，因为它们现在可以：（i）专注于数据的重要语义方面；（ii）在较低维度，计算上训练会更有效。
　　与以前在高度压缩，离散的潜在空间中使用自回归，基于注意力的transformer模型的工作不同，我们可以利用我们模型提供的特定图像的归纳偏差。这包括主要从2D卷积层构建的UNET的能力，并使用reweighted bound将目标进一步关注到最相关的位置，即：
$$L_{LDM}:={\mathbb{E}}_{\mathcal{E}(x),ϵ\sim \mathcal{N}(0,1),t}[||ϵ-{ϵ}_{\theta }(z_t,t)||].\text{\hspace{1em}(2)}$$
我们模型的神经结构${ϵ}_{\theta }(0,t)$是由time-conditional UNet实现的。由于前向过程是固定的，$z_t$能够在训练过程中从$\mathcal{E}$高效获取，并且来自$p(z)$的样本能够通过$\mathcal{D}$解码到图像空间。

3.3 Conditioning Mechanisms

　　与其他类型的生成模型相似，扩散模型原则上能够建模条件分布$p(z|y)$。这可以通过条件去躁自编码器${ϵ}_{\theta }(z_t,t,y)$来实现，并通过输入$y$（可以是文本，语义图或其他图像到图像翻译任务）来控制合成过程。
　　然而，在图像合成的背景下，将DM的生成能力与除类别标签以外的其他类型的条件或输入图像的模糊变体结合在一起，这是迄今为止正在探索的研究领域。
　　我们通过在UNet中使用交叉注意力机制，将DM变成更灵活的条件图像生成器，这对于学习以各种形式作为输入的基于注意力的模型更有效。为了能对来自各种模态的$y$（例如语言提示）进行处理，我们引入了一个领域特定的编码器，该编码器将$y$映射到中间表示${\tau }_{\theta }(y)\in {\mathbb{R}}^{M\times d_{\tau }}$，然后通过交叉注意力映射到UNet的中间层，该注意力被实现为$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d}})\cdot V$，其中，
$$Q=W_Q^{(i)}\cdot {\phi }_i(z_t),K=W_K^{(i)}\cdot {\tau }_{\theta }(y),V=W_V^{(i)}\cdot {\tau }_{\theta }(y).$$
${\phi }_i(z_t)\in {\mathbb{R}}^{N\times d_{ϵ}^i}$表示由UNet $ϵ$实现的展开的中间表示，并且$W_V^{(i)}\in {\mathbb{R}}^{d\times d_{ϵ}^i},W_Q^{(i)}\in {\mathbb{R}}^{d\times d_{\tau }},W_K^{(i)}\in {\mathbb{R}}^{d\times d_{\tau }}$是可学习的映射矩阵。如图3所示。
　　基于图像-条件对，我们通过下式学习条件LDM，
$$L_{LDM}:={\mathbb{E}}_{\mathcal{E}(x),y,ϵ\sim \mathcal{N}(0,1),t}[||ϵ-{ϵ}_{\theta }(z_t,t,{\tau }_{\theta }(y))|{|}_2^2],\text{\hspace{1em}(3)}$$
其中，${\tau }_{\theta }$和${ϵ}_{\theta }$通过等式3来联合优化。这种条件机制是灵活的，因为${\tau }_{\theta }$能够被参数化为特定领域。