利用潜在扩散模型进行高分辨率图像合成
Title:High-Resolution Image Synthesis with Latent Diffusion Models
paper是慕尼黑大学和海德堡大学发表在CVPR 22的工作
paper地址
文章目录
- 利用潜在扩散模型进行高分辨率图像合成
- Abstract
- 1. Introduction
- 2. Related Work
- 3. Method
- 3.1. Perceptual Image Compression
- 3.2. Latent Diffusion Models
- 3.3. Conditioning Mechanisms
Abstract
通过将图像形成过程分解为去噪自编码器的顺序应用，扩散模型（DMs）在图像数据及其它领域实现了最先进的合成结果。此外，它们的公式化允许引入一个引导机制来控制图像生成过程，而无需重新训练。然而，由于这些模型通常直接在像素空间中操作，优化强大的DMs通常需要数百GPU天，且由于顺序评估，推理成本昂贵。为了在有限的计算资源上启用DM训练，同时保留它们的质量和灵活性，我们将它们应用在强大的预训练自编码器的潜在空间中。与以往的工作相比，首次在复杂度降低和细节保留之间达到近乎最优点的训练扩散模型，大大提升了视觉保真度。通过在模型架构中引入交叉注意层，我们将扩散模型转变为强大且灵活的生成器，用于通用的条件输入，如文本或边界框，并且以卷积方式实现高分辨率合成成为可能。我们的潜在扩散模型（LDMs）在图像修复和类条件图像合成方面实现了最先进结果，并在各种任务上表现出高度竞争力，包括文本到图像的合成、无条件图像生成和超分辨率，同时与基于像素的DMs相比，显著降低了计算需求。
图1。通过较不激进的下采样提高可实现质量的上限。由于扩散模型为空间数据提供了出色的归纳偏置，我们不需要相关生成模型在潜在空间中的重度空间下采样，但仍然可以通过适当的自编码模型大大减少数据的维度，见第3节。图像来自DIV2K[1]验证集，评估在 51 2 2 p x 512^2 \mathrm{px} 5122px。我们用 f f f表示空间下采样因子。重建FIDs[29]和PSNR在ImageNet-val[12]上计算。另见表8。
1. Introduction
图像合成是计算机视觉领域中最近发展最为迅猛的领域之一，但也是计算需求最大的领域之一。特别是复杂自然场景的高分辨率合成，目前主要通过扩大（scaling up）基于似然的模型来主导，这些模型可能包含数十亿参数的自回归（AR）Transformer[66,67]。相比之下，GANs[3, 27, 40]的有希望的结果被发现主要局限于数据变异性相对有限的情况，因为它们的对抗学习程序不容易扩展到复杂、多模态分布的建模。最近，扩散模型[82]，它们由去噪自编码器的层次结构构建，已经在图像合成[30,85]及其它领域[7,45,48,57]中展示出令人印象深刻的结果，并且在类条件图像合成[15,31]和超分辨率[72]中实现了最先进的水平。此外，即使是无条件的DMs也可以很容易地应用于图像修复和上色[85]或基于笔触的合成[53]等任务，与其他类型的生成模型[19,46,69]相比。作为基于似然的模型，它们不会展现出像GANs那样的模式崩溃和训练不稳定性，并且，通过大量利用参数共享，它们可以建模自然图像的高度复杂分布，而不涉及像AR模型[67]那样的数十亿参数。
民主化高分辨率图像合成 DMs属于基于似然的模型类别，其模式覆盖行为使它们倾向于在对数据的不可察觉细节上花费过多的容量（计算资源）[16, 73]。尽管重新加权的变分目标[30]旨在通过对初始去噪步骤的欠采样来解决这一问题，DMs仍然在计算上要求很高，因为训练和评估这样一个模型需要在RGB图像的高维空间中重复函数评估（和梯度计算）。例如，训练最强大的DMs通常需要数百GPU天（例如，在[15]中为150 1000 V100天），并且在输入空间的噪声版本上重复评估也使得推理变得昂贵，以至于产生50k样本大约需要5天[15]在单个A100 GPU上。这对研究社区和用户总体有两个后果：首先，训练这样一个模型需要大量的计算资源，这些资源只对该领域的一小部分人可用，并且留下了巨大的碳足迹[65, 86]。其次，评估一个已经训练好的模型在时间和内存上也是昂贵的，因为相同的模型架构必须顺序运行大量的步骤（例如，在[15]中为25 - 1000步）。
为了提高这种强大模型类的可访问性，同时减少其巨大的资源消耗，需要一种方法来降低训练和采样的计算复杂度。因此，在不损害 DM 性能的情况下减少 DM 的计算需求是提高其可访问性的关键。
出发至潜在空间 我们的方法从分析像素空间中已训练的扩散模型开始：图 2 显示了已训练模型的速率失真权衡。与任何基于可能性的模型一样，学习大致可分为两个阶段：第一阶段是感知压缩阶段，该阶段会删除高频细节，但仍会学习很少的语义变化。在第二阶段，实际的生成模型会学习数据的语义和概念组成（语义压缩）。因此，我们的目标是首先找到一个感知等效但计算上更合适的空间，我们将在其中训练用于高分辨率图像合成的扩散模型。
图 2. 说明感知和语义压缩：数字图像的大多数比特对应于不可察觉的细节。虽然 DM 允许通过最小化负责损失项来抑制这种语义上无意义的信息，但梯度（在训练期间）和神经网络主干（训练和推理）仍然需要在所有像素上进行评估，从而导致多余的计算和不必要的昂贵优化和推理。 我们提出潜在扩散模型 (LDM) 作为一种有效的生成模型和一个单独的温和压缩阶段，仅消除不可察觉的细节。数据和图像来自 [30]。
遵循常见做法[11, 23, 66, 67, 96]，我们将训练分为两个不同的阶段：首先，我们训练一个自编码器，它提供了一个低维（因此高效）的表示空间，这个空间在感知上等同于数据空间。重要的是，与之前的工作[23,66]相比，我们不需要依赖过度的空间压缩，因为我们在学习到的潜在空间中训练DMs，这展示了与空间维度相关的更好的缩放属性。减少的复杂性还提供了从潜在空间以单个网络传递进行高效图像生成的能力。我们将得到的模型类称为潜在扩散模型（LDMs）。
这种方法的一个显著优势是，我们只需要训练通用自编码阶段一次，因此可以将其重用于多个DM训练或探索可能完全不同的任务[81]。这使得能够高效地探索大量的扩散模型，用于各种图像到图像和文本到图像的任务。对于后者，我们设计了一种架构，将Transformer连接到DM的UNet骨干[71]，并支持任意类型的基于token的条件机制，见第3.3节。
总的来说，我们的工作做出了以下贡献：(i) 与纯粹基于Transformer的方法[23, 66]相比，我们的方法对高维数据的扩展更为优雅，因此可以(a) 在一个提供比以往工作更忠实和详细重建的压缩级别上工作（见图1）以及(b) 可以高效地应用于百万像素图像的高分辨率合成。
(ii) 我们在多个任务（无条件图像合成、图像修复、随机超分辨率）和数据集上实现了竞争性的表现，同时显著降低了计算成本。与基于像素的扩散方法相比，我们还显著降低了推理成本。
(iii) 我们展示了，与先前的工作[93]相比，后者同时学习编码器/解码器架构和基于评分的先验，我们的方法不需要精细的重建和生成能力的权衡。这确保了极其忠实的重建，并且对潜在空间的正则化需求非常小。
(iv) 我们发现，对于像超分辨率、图像修复和语义合成这样密集条件的任务，我们的模型可以以卷积方式应用，并渲染出大型、一致的图像，约 102 4 2 1024^2 10242个像素。
(v) 此外，我们设计了一种基于交叉注意力的通用条件机制，使得多模态训练成为可能。我们使用它来训练类条件、文本到图像和布局到图像的模型。
(vi) 最后，我们在https://github.com/CompVis/latent-diffusion上发布了预训练的潜在扩散和自编码模型，这些模型可能除了DMs的训练外，还可以用于各种任务[81]。
这篇文章也是Stable Diffusion的核心：Code链接：https://github.com/CompVis/stable-diffusion。
2. Related Work
用于图像合成的生成模型 图像的高维性质对生成建模提出了独特的挑战。生成对抗网络（GAN）[27]允许高效采样高分辨率图像，并具有良好的感知质量[3, 42]，但是它们难以优化[2, 28, 54]，并且难以捕捉完整的数据分布[55]。相比之下，基于似然的方法强调良好的密度估计，这使得优化更加规范。变分自编码器（VAE）[46]和基于流的模型[18, 19]使高分辨率图像的高效合成成为可能[9, 44, 92]，但是样本质量与GANs不相匹配。虽然自回归模型（ARM）[6, 10, 94, 95]在密度估计方面取得了强大的性能，但是计算要求高的架构[97]和序列采样过程将它们限制在低分辨率图像上。因为基于像素的图像表示包含几乎不可察觉的高频细节[16,73]，最大似然训练在对它们进行建模上花费了不成比例的大量计算资源，导致训练时间长。为了扩展到更高的分辨率，几种两阶段方法[23,67,101,103]使用ARMs来建模压缩的潜在图像空间，而不是原始像素。
最近，扩散概率模型（DM）[82]在密度估计[45]以及样本质量[15]方面取得了最先进的结果。这些模型的生成能力源于当其底层神经骨架实现为UNet[15, 30, 71, 85]时，它们与图像类数据的归纳偏置自然契合。当使用重新加权的目标[30]进行训练时，通常可以达到最佳的合成质量。在这种情况下，DM对应于一个有损压缩器，并允许在图像质量和压缩能力之间进行权衡。然而，评估和优化这些模型在像素空间中的缺点是低推理速度和非常高的训练成本。尽管前者可以通过高级采样策略[47, 75, 84]和分层方法[31, 93]部分解决，但在高分辨率图像数据上训练始终需要计算昂贵的梯度。我们提出的LDMs解决了这两个缺点，它们在更低维度的压缩潜在空间上工作。这使得训练在计算上更加便宜，并且几乎不降低合成质量的情况下加速了推理（见图1）。
两阶段图像合成 为了缓解单一生成方法的缺点，很多研究[11, 23, 67, 70, 101, 103]已经致力于通过两阶段方法将不同方法的优势结合起来，形成更高效和性能更强的模型。VQ-VAEs[67, 101]使用自回归模型学习离散潜在空间上的表达性先验。[66]通过学习离散图像和文本表示的联合分布，将这种方法扩展到文本到图像的生成。更一般地，[70]使用条件可逆网络提供多样域潜在空间之间的通用转换。与VQ-VAEs不同，VQGANs[23, 103]在第一阶段采用对抗性和感知目标，将自回归变换器扩展到更大的图像上。然而，可行的ARM训练所需的高压缩率引入了数十亿的可训练参数[23, 66]，限制了此类方法的整体性能，而较少的压缩则需要高计算成本[23, 66]。我们的工作避免了这种权衡，因为我们提出的LDMs由于其卷积骨架，更温和地扩展到更高维度的潜在空间。因此，我们可以自由选择压缩级别，以最佳方式在学习强大的第一阶段和不将过多的感知压缩留给生成扩散模型之间进行调解，同时保证高保真重建（见图1）。
虽然存在联合[93]或分别[80]学习编码/解码模型以及基于评分的先验的方法，但前者仍然需要在重建和生成能力之间进行艰难的权衡[11]，并且被我们的方法（第4节）超越，后者则专注于如人脸等高度结构化的图像。
3. Method
为了降低训练扩散模型以实现高分辨率图像合成的计算需求，我们观察到尽管扩散模型允许通过对相应损失项进行欠采样来忽略感知上不相关的细节[30]，但它们仍然需要在像素空间中进行昂贵的函数评估，这导致了巨大的计算时间和能源资源需求。
我们提出通过引入压缩学习阶段与生成学习阶段的显式分离来规避这一缺点（见图2）。为了实现这一点，我们利用一个自编码模型，该模型学习一个与图像空间在感知上等效但计算复杂度显著降低的空间。
这种方法提供了几个优势：（i）通过离开高维图像空间，我们获得了计算效率更高的DMs，因为采样是在低维空间上进行的。（ii）我们利用了DMs从其UNet架构[71]继承的归纳偏置，这使得它们对于具有空间结构的数据特别有效，因此减轻了之前方法[23, 66]所需的激进的、降低质量的压缩水平的需求。（iii）最后，我们获得了通用的压缩模型，其潜在空间可以用来训练多个生成模型，也可以用于其他下游应用，如单图像CLIP引导的合成[25]。
3.1. Perceptual Image Compression
我们的感知压缩模型基于之前的工作[23]，由一个自编码器组成，该自编码器通过感知损失[106]和基于patch的[33]对抗性目标[20,23,103]的组合进行训练。这确保了重建被限制在图像流形上，通过强制执行局部现实性并避免仅依赖像素空间损失（如 L 2 L_2 L2或 L 1 L_1 L1目标）引入的模糊。
更准确地说，给定一个在RGB空间中的图像 x ∈ R H × W × 3 x \in \mathbb{R}^{H \times W \times 3} x∈RH×W×3，编码器 E \mathcal{E} E将 x x x编码成潜在表示 z = E ( x ) z=\mathcal{E}(x) z=E(x)，解码器 D \mathcal{D} D从潜在空间重建图像，得到 x ~ = D ( z ) = D ( E ( x ) ) \tilde{x}=\mathcal{D}(z)=\mathcal{D}(\mathcal{E}(x)) x~=D(z)=D(E(x))，其中 z ∈ R h × w × c z \in \mathbb{R}^{h \times w \times c} z∈Rh×w×c。重要的是，编码器通过因子 f = H / h = W / w f=H / h=W / w f=H/h=W/w对图像进行下采样，我们研究了不同的下采样因子 f = 2 m f=2^m f=2m，其中 m ∈ N m \in \mathbb{N} m∈N。
为了避免任意高方差的潜在空间，我们尝试了两种不同的正则化方法。第一种变体， K L K L KL-reg.，对学习到的潜在空间施加轻微的KL惩罚，类似于VAE[46,69]，而VQ-reg.在解码器中使用矢量量化层[96]。这个模型可以被解释为VQGAN[23]，但是量化层被解码器吸收。因为我们后续的DM旨在与我们学习到的潜在空间 z = E ( x ) z=\mathcal{E}(x) z=E(x)的二维结构一起工作，我们可以使用相对温和的压缩率并实现非常好的重建。这与之前的工作[23,66]形成对比，后者依赖于学习空间 z z z的任意一维排序来自回归地建模其分布，从而忽略了 z z z的许多固有结构。因此，我们的压缩模型更好地保留了 x x x的细节（见表8）。完整的目标和训练细节可以在补充材料中找到。
3.2. Latent Diffusion Models
扩散模型[82]是旨在通过逐渐去噪正态分布变量来学习数据分布
p
(
x
)
p(x)
p(x)的概率模型，这对应于学习固定马尔可夫链长度
T
T
T的逆过程。对于图像合成，最成功的模型[15,30,72]依赖于
p
(
x
)
p(x)
p(x)的变分下界的加权变体，这反映了去噪得分匹配[85]。这些模型可以被解释为一系列等权重的去噪自编码器
ϵ
θ
(
x
t
,
t
)
;
t
=
1
…
T
\epsilon_\theta\left(x_t, t\right) ; t=1 \ldots T
ϵθ(xt,t);t=1…T，它们被训练用于预测其输入
x
t
x_t
xt的去噪版本，其中
x
t
x_t
xt是输入
x
x
x的噪声版本。相应的目标可以简化为（见B节）
L
D
M
=
E
x
,
ϵ
∼
N
(
0
,
1
)
,
t
[
∥
ϵ
−
ϵ
θ
(
x
t
,
t
)
∥
2
2
]
,
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],
LDM=Ex,ϵ∼N(0,1),t[∥ϵ−ϵθ(xt,t)∥22],
其中
t
t
t从
{
1
,
…
,
T
}
\{1, \ldots, T\}
{1,…,T}中均匀采样。潜在表示的生成建模 有了我们训练有素的感知压缩模型
E
\mathcal{E}
E和
D
\mathcal{D}
D，我们现在可以访问一个高效的、低维的潜在空间，在这个空间中，高频、不可感知的细节被抽象化。与高维像素空间相比，这个空间更适合基于似然的生成模型，因为它们现在可以（i）专注于数据的重要、语义位，并且（ii）在一个维度更低、计算效率更高的空间中训练。
与之前依赖于自回归、基于注意力的变换模型在高度压缩的、离散的潜在空间中的工作不同[23,66,103]，我们可以利用我们的模型提供的图像特定的归纳偏置。这包括能够主要从2D卷积层构建底层的UNet，并进一步聚焦于使用加权边界的感知上最相关的位，现在写作
L
L
D
M
:
=
E
E
(
x
)
,
ϵ
∼
N
(
0
,
1
)
,
t
[
∥
ϵ
−
ϵ
θ
(
z
t
,
t
)
∥
2
2
]
.
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] .
LLDM:=EE(x),ϵ∼N(0,1),t[∥ϵ−ϵθ(zt,t)∥22].
我们模型的神经骨架 ϵ θ ( ∘ , t ) \epsilon_\theta(\circ, t) ϵθ(∘,t)是作为一个时间条件的UNet[71]实现的。由于前向过程是固定的， z t z_t zt可以在训练期间从 E \mathcal{E} E高效获得，并且可以通过 D \mathcal{D} D单次传递将 p ( z ) p(z) p(z)中的样本解码到图像空间。
3.3. Conditioning Mechanisms
类似于其他类型的生成模型[56,83]，扩散模型原则上能够建模形式为 p ( z ∣ y ) p(z \mid y) p(z∣y)的条件分布。这可以通过一个条件去噪自编码器 ϵ θ ( z t , t , y ) \epsilon_\theta\left(z_t, t, y\right) ϵθ(zt,t,y)来实现，并为通过输入 y y y（如文本[68]、语义地图[33,61]或其他图像到图像的翻译任务[34]）控制合成过程铺平了道路。
然而，在图像合成的背景下，将DMs的生成能力与类标签[15]或输入图像的模糊变体[72]之外的其他类型的条件结合起来，到目前为止是一个探索不足的研究领域。
我们通过用交叉注意力机制[97]增强其底层的UNet骨架，将DMs转变为更灵活的条件图像生成器，这对于学习各种输入模态[35,36]的基于注意力的模型是有效的。为了预处理来自各种模态的
y
y
y（例如语言提示），我们引入了一个特定领域的编码器
τ
θ
\tau_\theta
τθ，它将
y
y
y投影到一个中间表示
τ
θ
(
y
)
∈
R
M
×
d
τ
\tau_\theta(y) \in \mathbb{R}^{M \times d_\tau}
τθ(y)∈RM×dτ，然后通过实现注意力
(
Q
,
K
,
V
)
=
softmax
(
Q
K
T
d
)
⋅
V
(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) \cdot V
(Q,K,V)=softmax(dQKT)⋅V的交叉注意力层映射到UNet的中间层，其中
Q
=
W
Q
(
i
)
⋅
φ
i
(
z
t
)
,
K
=
W
K
(
i
)
⋅
τ
θ
(
y
)
,
V
=
W
V
(
i
)
⋅
τ
θ
(
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)
Q=WQ(i)⋅φi(zt),K=WK(i)⋅τθ(y),V=WV(i)⋅τθ(y)
这里， φ i ( z t ) ∈ R N × d e i \varphi_i\left(z_t\right) \in \mathbb{R}^{N \times d_e^i} φi(zt)∈RN×dei表示实现 ϵ θ \epsilon_\theta ϵθ的UNet的（扁平化的）中间表示，而 W V ( i ) ∈ W_V^{(i)} \in WV(i)∈ R d × d ϵ i , W Q ( i ) ∈ R d × d τ & W K ( i ) ∈ R d × d τ \mathbb{R}^{d \times d_\epsilon^i}, W_Q^{(i)} \in \mathbb{R}^{d \times d_\tau} \& W_K^{(i)} \in \mathbb{R}^{d \times d_\tau} Rd×dϵi,WQ(i)∈Rd×dτ&WK(i)∈Rd×dτ是可学习的投影矩阵[36,97]。参见图3以获得视觉描述。
图 3. 我们通过拼接或更通用的交叉注意机制来调节 LDM。请参阅第 3.3 节
基于图像-条件对，我们然后通过
L
L
D
M
:
=
E
E
(
x
)
,
y
,
ϵ
∼
N
(
0
,
1
)
,
t
[
∥
ϵ
−
ϵ
θ
(
z
t
,
t
,
τ
θ
(
y
)
)
∥
2
2
]
,
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] \text {, }
LLDM:=EE(x),y,ϵ∼N(0,1),t[∥ϵ−ϵθ(zt,t,τθ(y))∥22],
学习条件LDM，其中
τ
θ
\tau_\theta
τθ和
ϵ
θ
\epsilon_\theta
ϵθ通过方程3共同优化。这种条件机制是灵活的，因为
τ
θ
\tau_\theta
τθ可以用特定领域的专家参数化，例如当
y
y
y是文本提示时的（未屏蔽的）Transformer[97]（见第4.3.1节）
