【论文精读】Latent Diffusion

摘要

       Diffusion models（DMs）被证明在复杂自然场景的高分辨率图像合成能力优于以往的GAN或autoregressive （AR）transformer。作为基于似然的模型，其没有GAN的模式崩溃和训练不稳定问题，通过参数共享，其可以模拟自然图像的高度复杂分布；另外也不需要AR类模型庞大的参数量，但DM在像素级的训练推理仍然需要大量的计算量，故本文的重点在于在不损害DM性能的情况下减少计算量，以优化算法效率。

       基于似然估计的模型的学习大致分为两个阶段：首先为感知压缩阶段，该阶段可以去除图像的高频细节，也能学习到少量的语义变化；随后生成模型会学习数据的语义和概念组成（语义压缩）。故本文目标是找到一个在感知上与数据空间等效但计算量更小的空间，并在其中训练用于高分辨率图像合成的DM。

       故本文提出的方法在训练过程会分成两个不同的阶段：首先会训练一个自编码器，该编码器会提供了一个在感知上等效于数据空间的低维表示空间，随后在该潜空间中训练DM。该模型称为Latent Diffusion Models（LDM）。

       上图显示了提出的模型在不同阶段的浅表示失真率和单维度语义信息量的权衡图，在感知压缩阶段，自编码器编码的浅表示具有高度的信息密度，同时保持了较低的图像失真率，保留了图像的整体结构。而在生成阶段，LDM对图像信息压缩率较小，故可以建模图像细节的能力，但同时失真率较高，会损失图像的整体结构信息，但这部分信息刚好可以由感知编码器弥补。故整体上，本文方法具备提取重建图像整体结构与细节信息的能力。

       这种方法的优点在于，只需要训练一次通用自编码阶段，随后就可以将其重复用于多个DM任务，例如各种图生图和文生图任务。为了适配这些任务，本文设计了一个架构，将transformer连接到DM的UNet骨干网络，以实现任意类型的基于token的适配机制。整体改进有：

• 编码器会对数据进行压缩，且因为DM本身可以为数据提供极好的归纳偏差，故在该压缩空间中，DM仍然可以实现可靠和详细的数据重建（如上图，自编码器计算的潜空间只需进行4倍下采样就可以实现超越以往生成模型的性能），且可以有效地应用于百万像素图像的高分辨率合成
• 在多个任务（无条件图像合成、图像修复、超分辨率）和数据集上实现了最先进的性能，且显著降低了训练和推理的计算量
• 因为本文方法是分阶段训练的，故与以往同时训练自编码器和似然估计的方法相比，所提出方法不需要自编码器重建和似然估计生成能力的加权平衡。这确保了编码器的重建数据的可靠性，且降低了对潜空间的正则化要求
• 设计了一种基于交叉注意力的通用条件适配机制，实现了多模态训练。故本方法可以用来训练类别条件模型（Class-Conditional）、文本到图像和布局到图像的模型

框架

       DM允许通过对相应的损失项进行欠采样来忽略感知上不相关的图像细节，但其仍然需要在像素空间进行函数评估，计算复杂度过高。为了降低DM高分辨率图像生成的计算复杂度，本文模型采用两阶段训练，第一阶段利用自编码学习一个在感知上等效于图像空间的潜空间，得到的浅表示为压缩后的图像高密度信息，随后在该潜空间中采用DM进行图像生成，该方法的优点如下：

• 通过避开高维图像空间，使DM计算效率提高
• 利用了从UNet架构中继承的DM的归纳偏差，保证其对图像空间结构信息的特征提取能力
• 获得的通用的自编码器压缩模型，其潜空间可用于训练多个生成模型，包括各种下游应用，如采用CLIP的文本生成或图像生成

Perceptual Image Compression

       感知压缩模型为由感知损失和基于patch的对抗损失组合训练的自编码器，这确保了重建能通过局部真实性被限制在图像流形上，避免了仅仅依靠像素空间损失（$L_1$或$L_2$损失）而导致的模糊性。

       具体，给定图像$x\in {\mathbb{R}}^{H\times W\times 3}$，编码器$E$会将$x$编码为潜表示$z=E(x)$，解码器$D$再从潜表示中重建图像$\hat{x}=D(z)=D(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：在解码器内使用向量量化（vector quantization）层

       该模型类似于VQGAN，但量化层被解码器吸收。潜表示$z=E(x)$会进一步由后续的DM处理。

Latent Diffusion Models

Diffusion Models

       DM是通过对正态分布变量逐步去噪来学习数据分布 $p(x)$的概率模型，相当于学习长度为$T$的固定马尔可夫链的逆过程。可以解释为一个等权重的去噪自编码器序列${ϵ}_{\theta }(x_t,t);t=1\dots T$，其被训练以预测其输入$x_t$的去噪变体，其中$x_t$为输入$x$加噪声的版本。 相应的目标可以简化为：
$$L_{DM}={\mathbb{E}}_{x,ϵ\sim N(0,1),t}[||ϵ-{ϵ}_{\theta }(x_t,t)|{|}_2^2]$$

       其中$t$均匀采样自 { 1 , … , T } \{1,\dots ,T\} { 1,…,T}。

Generative Modeling of Latent Representations

       通过训练由$E$和$D$组成的感知压缩模型，可以获得一个有效的低维潜空间，其中图像的高频信号、难以察觉的细节被抽象掉了。与高维像素空间相比，这个空间更适合基于似然的生成模型，因为DM现在可以专注于数据中更重要的语义位，且在这个低维空间的计算效率比在像素空间中高得多。

       利用由2D卷积层构建的UNet，重新加权的bound会进一步将目标集中在感知上最相关的比特上，故DM可以提供特定于图像的归纳偏差。结合浅表示与DM的学习目标，则有：
$$L_{LDM}:={\mathbb{E}}_{E(x),ϵ\sim N(0,1),t}[||ϵ-{ϵ}_{\theta }(z_t,t)|{|}_2^2]$$

       其中LDM的神经主干${ϵ}_{\theta }(\cdot ,t)$为一个时间条件UNet，$z_t=E(x_t)$，概率模型$p(z)$的样本通过$D$解码到图像空间。

Conditioning Mechanisms

       DM能够建模形式为$p(z|y)$的条件分布，结合上式则可以用条件去噪自编码器${ϵ}_{\theta }(z_t,t,y)$实现多模态训练，其中输入$y$即为prompt隐编码，可以为文本，语义映射或图像等。

       本文通过使用交叉注意机制增强UNet主干的方式，将DM转变为更灵活的条件图像生成器。为了编码各种模态（如语言prompt）的输入$y$，引入一个特定领域的编码器${\tau }_{\theta }$将$y$投影为中间表示${\tau }_{\theta }(y)\in R^{M\times d\tau }$，然后通过交叉注意将其映射到UNet的中间层，公式为：
$$\begin{gathered} 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) \end{gathered}$$

       其中${\phi }_i(z_t)\in {\mathbb{R}}^{N\times d_e^i}$为实现了${ϵ}_{\theta }$的UNet的中间层（第$i$层）表示，$W_V^{(i)}\in {\mathbb{R}}^{d\times d_{ϵ}^i}$，$W_Q^{(i)},W_K^{(i)}\in {\mathbb{R}}^{d\times d_{\tau }}$为可学习的投影矩阵。模型的整体计算流程如下图。

       基于条件生成，学习目标进一步改进为条件LDM：