在上一篇文章中，我们梳理了基于自编码器（AE）的图像生成模型的发展脉络，并引出了Stable Diffusion的核心思想。简单来说，Stable Diffusion是一个两阶段的图像生成模型，它先用一个AE压缩图像，再在压缩图像所在的隐空间上用DDPM生成图像。在这篇文章中，我们来精读Stable Diffusion的论文：High-Resolution Image Synthesis with Latent Diffusion Models。
注意：如果你从未学习过扩散模型，Stable Diffusion并不是你应该的读的第一篇论文。请参照我[上篇文章]的早期工作总结，至少在学会了DDPM后再来学习Stable Diffusion。
摘要与引言
论文摘要的大意如下：扩散模型的生成效果很好，但是，在像素空间上训练和推理扩散模型的计算开销都很大。为了在不降低质量与易用性的前提下用较少的计算资源训练扩散模型，我们在一个预训练过的自编码器的隐空间上使用扩散模型。相较以往的工作，在这种表示下训练扩散模型首次在减少计算复杂度和维持图像细节间达到几近最优的平衡点，极大地提升了视觉保真度。通过向模型架构中引入交叉注意力层，我们把扩散模型变成了强大而灵活的带约束图像生成器，它支持常见的约束，如文字、边界框，且能够以纯卷积方式实现高分辨率的图像合成。我们的隐扩散模型（latent diffusion model, LDM） 在使用比像素扩散模型少得多的计算资源的前提下，在各项图像合成任务上取得最优成果或顶尖成果。
整理一下。论文提出了一种叫LDM的图像生成模型。论文想解决的问题是减少像素空间扩散模型的运算开销。为此，LDM借助了VQVAE「先压缩、再生成」的想法，把扩散模型用在AE的隐空间上，在几乎不降低生成质量的前提下减少了计算量。另外，LDM还支持带约束图像合成及纯卷积图像超分辨率。
在上一篇回顾LDM早期工作的文章中，我们已经理解了LDM想解决的问题及解决问题的思路。因此，在读完摘要后，我们接下来读文章时只需要关注LDM的两个创新点：
- LDM的AE是怎么设计以达到压缩比例与质量的平衡的。
- LDM怎么实现带约束的图像合成。
引言基本是摘要的扩写。首先，引言大致介绍了图像合成任务的背景，提及了扩散模型近期的突出表现。随后，引言介绍了本文想解决的主要问题：扩散模型的训练和推理太耗时了，需要在不降低效果的前提下减少扩散模型的运算量。最后，引言揭示了本工作的解决方法：使用类似VQGAN的两阶段图像生成方法。
引言的前两部分没有什么关键信息，而最后一部分介绍了本工作改进扩散模型的动机，值得一读。如下图所示，DDPM的论文展示了从不同去噪时刻的同一个噪声图像开始的不同生成结果，比如$\mathbf{x}_{750}$指从时刻$t=750$的去噪图像开始，多次以不同随机数执行DDPM的反向过程，生成的多幅图像。LDM作者认为，DDPM的这一实验表明，扩散模型的图像生成分两个阶段：先是对语义进行压缩，再是对图像的感知细节压缩。正因此，随机对早期的噪声图像去噪，生成图像的内容会更多样；而随机对后期的噪声图像去噪，生成图像只是在细节上有所不同。LDM的作者认为，扩散模型的大量计算都浪费在了生成整幅图像的细节上，不如只让扩散模型描述比较关键的语义压缩部分，而让自编码器（AE）负责感知细节压缩部分。
引言在结尾总结了本工作的贡献：
- 相比之前按序列处理图像的纯Transformer的方法，扩散模型能更好地处理二维数据。因此，LDM生成隐空间图像时不需要那么重的压缩比例（比如DIV2K数据集上，LDM只需要将图像下采样4倍，而之前的纯Transformer方法要下采样8倍或16倍），图像在压缩时能有更高的保真度，整套方法能更高效地生成高分辨率图像。
- 在大幅降低计算开销的前提下在多项图像生成任务上取得了顶尖成果。
- 相比于之前同时训练图像压缩模型和图像生成模型的方法，该方法分步训练两个模型，训练起来更加简单。
- 对于有着稠密约束的任务（如超分辨率、补全、语义生成），该方法的模型能换成一个纯卷积版本的，且能生成边长为1024的图像。
- 该工作设计了一种通用的约束机制，该机制基于交叉注意力，支持多模态训练。作者训练了多种带约束的模型。
- 作者把工作开源了，并提供了预训练模型。
我们来整理一下这些贡献。读论文时，可以忽略第6条。第2条是成果，与方法设计无关。第1、3条主要描述了提出两阶段图像生成建模方法的贡献。第4条是把方法拓展到稠密约束任务的贡献。第5条是提出了新约束机制的贡献。所以，在学习论文的方法时，我们还是主要关注摘要里就提过的那两个创新点。在读完引言后，我们可以把阅读目标再细化一下：
- LDM的AE是怎么设计以达到压缩比例与质量的平衡的。与纯基于Transformer的VQGAN相比，它有什么不同。
- LDM怎么用交叉注意力机制实现带约束的图像生成。
相关工作
作者主要从两个角度回顾了早期工作：不同架构的图像生成模型与两阶段的图像合成方法。其回顾逻辑与本系列的第一篇文章类似，在此就不过多介绍了。除了介绍早期工作外，作者重申了引言中的对比结果，强调了LDM相对于扩散模型的创新和相对于两阶段图像生成模型的创新。
方法
在方法章节中，作者先是大致介绍了使用LDM这种两阶段图像生成架构的优点，再分三部分详细介绍了论文的实现细节：图像压缩AE的实现、LDM的实现、约束的实现。开头的介绍和AE的实现相对比较重要，我们放在一起详细阅读；相对于DDPM，LDM几乎没有做任何修改，只是把要拟合的图片从真实图片换成了压缩图片，这一部分我们会快速浏览一遍；而添加约束的方法有所创新，我们会详细阅读一遍。
AE与两阶段图像生成模型
我们来先读3.1节，看一看AE的具体实现方法，再回头读第3节开头介绍的两阶段图像生成模型的优点。
LDM配套的图像压缩模型（论文中称之为”感知压缩模型”）和VQGAN几乎完全一样。该压缩模型的原型是一个AE。普通的AE会用原图像和重建图像的重建误差（L1误差或者L2误差）来训练。在普通的AE的基础上，该压缩模型参考了GAN的误差设置方法，使用感知误差代替重建误差，并添加了基于patch的对抗误差。
但该图像压缩模型的输出与VQGAN有所不同。我们先回忆一下VQGAN的原理。VQGAN的输出会接到Transformer里，Transformer的输入必须是离散的。因此，VQGAN必须要额外完成两件事：1）让连续输出变成离散输出；2）用正则化方法防止过拟合。为此，VQGAN使用了VQVAE里的向量离散化操作，该操作能同时完成这两件事。
而LDM的压缩模型的输出会接入一个扩散模型里，扩散模型的输入是连续的。因此，LDM的压缩模型只需要额外完成使用正则化方法这一件事。该压缩模型不必像VQGAN一样非得用向量离散化来完成正则化。如我们在第一篇文章中讨论的，作者在LDM的压缩模型中使用了两种正则化方法：VQ正则化与KL正则化。前者来自于VQVAE，后者来自于VAE。
该压缩模型相较VQGAN有一项明显的优势。VQGAN的Transformer只能按一维序列来处理图像（通过把二维图像reshape成一维），且只能处理较小的压缩图像($16\times16$)。而本身用于二维图像生成的LDM能更好地利用二维信息，因此可以处理更大的压缩图像($64\times 64$)。这样，LDM的压缩模型的压缩程度不必那么重，其保真度会比VQGAN高。
看完了3.1节，我们来回头看第3节开头介绍了LDM的三项优点：1）通过规避在高维图像空间上训练扩散模型，作者开发出了一个因在低维空间上采样而计算效率大幅提升的扩散模型；2）作者发掘了扩散模型中来自U-Net架构的归纳偏置（inductive bias），使得它们能高效地处理有空间结构的数据（比如二维图像），避免像之前基于Transformer的方法一样使用激进、有损质量的压缩比例；3）本工作的压缩模型是通用的，它的隐空间能用来训练多种图像生成模型。第一个优点是相对于DDPM。第二个是优点是相对于使用Transformer的VQGAN，我们在上一段已经分析过了。第三个优点是相对于之前那些换一个任务就需要换一个压缩模型的两阶段图像生成模型。
归纳偏置可以简单理解为某个学习算法对一类数据的优势。比如CNN结构适合处理图像数据。
隐扩散模型（LDM）
在DDPM中，一个参数为$\theta$的神经网络$\epsilon_\theta$会根据当前时刻$t$的带噪图片$x_t$预测本时刻的噪声$\epsilon_\theta(x_t, t)$。网络的学习目标是让预测的噪声和真实的噪声$\epsilon$一致。
LDM的原理和DDPM完全一样，只不过训练图片从像素空间上的真实图片$x_0$变成了隐空间上的压缩图片$z_0$，每一轮的带噪图片由$x_t$变成了隐空间上的带噪图片$z_t$。在训练时，相比DDPM，只需要多对$x_0$用一次编码器变成$z_0$即可。
如果你在理解这部分内容时有疑问，请去阅读DDPM的相关文章。LDM的具体结构我们会在第三篇代码阅读文章中讨论。
约束机制
让模型支持带约束图像生成，其实就是想办法把额外的约束信息输入进扩散模型中。显然，最简单的添加约束的方法就是把额外的信息和扩散模型原本的输入$z_t$拼接起来。如果约束是一个值，就把相同的值拼接到$z_t$的每一个像素上；如果约束本身具有空间结构（如语音分割图片），就可以把约束重采样至和$z_t$一样的分辨率，再逐像素拼接。除了直接的拼接外，作者在LDM中还使用了另一种融合约束信息的方法。
DDPM中含有自注意力层。自注意力操作其实基于注意力操作$Attention(Q, K, V)$，它可以解释成一个数据库中存储了许多数据$V$，数据的索引（键）是$K$，现在要用查询$Q$查询数据库里的数据并返回查询结果。注意力操作有几种用法，第一种用法是交叉注意力$CrossAttn(A, B)=Attention(W_Q \cdot A, W_K \cdot B, W_V \cdot B)$，可以理解成数据$A$, $B$做了一次信息融合；第二种用法是自注意力$SelfAttn(A)=Attention(W_Q \cdot A, W_K \cdot A, W_V \cdot A)$，可以理解成数据$A$自己做了一次特征提取。
既然交叉注意力操作可以融合两类信息，何不把DDPM的自注意力层换成交叉注意力层，把$K$, $V$换成来自约束的信息，以实现带约束图像生成呢？如下图所示，通过把用编码器$\tau_\theta$编码过的约束信息输入进扩散模型交叉注意力层的$K$, $V$，LDM实现了带约束图像生成。这里的实现细节我们会在第三篇代码阅读文章中讨论。
根据论文中实验的设计，对于作用于全局的约束，如文本描述，使用交叉注意力较好；对于有空间信息的约束，如语义分割图片，则用拼接的方式较好。
实验
在这一章里，作者按照介绍方法的顺序，依次探究了图像压缩模型、无约束图像生成、带约束图像合成的实验结果。我们主要关心前两部分的实验结果。
感知压缩程度的折衷
论文首先讨论了图像压缩模型在不同的下采样比例$f$下的实验结果，其中$f\in\{1, 2, 4, 8, 16, 32\}$。这些实验分两部分，第一部分是训练速度上的实验，第二部分是采样速度与效果上的实验。
在ImageNet上以不同下采样比例$f$训练一定步数后LDM的采样指标对比结果如下图所示。其中，FID指标越低越好，Inception Score越高越好。结果显示，无论下采样比例$f$是过大还是过小都会降低训练速度。作者分析了$f$较小或较大时训练速度慢的原因：$f$过小时扩散模型把过多的精力放在了本应由压缩模型负责的感知压缩上；$f$过大时图像信息在压缩中损失过多。LDM-$\{4\text{-}16\}$的表现相对好一些。
在实验的第二部分中，作者比较了不同采样比例$f$的LDM在CelebA-HQ（下图左侧）和ImageNet（下图右侧）上的采样速度和采样效果。下图中，横坐标为吞吐量，越靠右表示采样速度越快。同一个模型的不同实验结果表示使用不同DDIM采样步数时的实验结果，每一条线上的结果从右到左分别是DDIM采样步数取$\{10, 20, 50, 100, 200\}$的采样结果（DDIM步数越少，采样速度越快，生成图片质量越低）。对于CelebA-HQ上的实验，若采样步数较多，则还是LDM-$\{4, 8\}$效果较好，只有在采样步数较少时压缩比更高的LDM才有优势。而对于ImageNet上的实验，$f$太小或太大的结果都很差，整体上还是LDM-$\{4, 8\}$的结果较好。
综上，根据实验，作者认为$f$取适中的$4$或$8$比较妥当。下采样比例$f=8$也正是Stable Diffusion采用的配置。
图像生成效果
在这一节中，作者在几个常见的数据集上对比了LDM与其他模型的无约束图像生成效果。作者主要比较了两类指标：表示采样质量的FID和表示数据分布覆盖率的精确率及召回率（Precision-and-Recall）。
在介绍具体结果之前，先对这个不太常见的精确率及召回率指标做一个解释。精确率及召回率常用于分类等有确定答案的任务中，分别表示所有被分类为正的样本中有多少是分对了的、所有真值为正的样本中有多少是被成功分类成正的。而无约束图像生成中的精确率及召回率的解释可以参加论文Improved Precision and Recall Metric for Assessing
Generative Models。如下图所示，设真实分布为蓝色，生成模型的分布为红色，则红色样本落在蓝色分布的比例为精确率，蓝色样本落在红色分布的比例为召回率。简单来说，精确率能描述采样质量，召回率能描述生成分布与真实分布的覆盖情况。
接下来，我们回头来看论文展示的无约束图像生成对比结果，如下图所示。整体上看，LDM的表现还不错。虽然在FID指标上无法超过GAN或其他扩散模型，但是在精确率和召回率上还是颇具优势。唯一没有被LDM战胜的是LSUN-Bedrooms上的ADM模型，但作者提到，相比ADM，LDM只用了一半的参数，且只需四分之一的训练资源。
带约束图像合成
这一节里，作者展示了LDM的文生图能力。论文中的LDM用了一个从头训练的基于Transformer的文本编码器，与后续使用CLIP的Stable Diffusion差别较大。这一部分的结果没那么重要，大致看一看就好。
本文的文生图模型是一个在LAION-400M数据集上训练的KL约束LDM。它的文本编码器是一个Transformer，编码后的特征会以交叉注意力的形式传入LDM。采样时，LDM使用了Classifier-Free Guidance。
Classifier-Free Guidance可以让输出图片更符合文本约束。这是一种适用于所有扩散模型的采样策略，并非要和LDM绑定，感兴趣可以去阅读相关论文。
LDM与其他模型的文生图效果对比如下图所示。虽然这个版本的LDM并没有显著优于其他模型，但它的参数量是最少的。
LDM在类别约束的图像合成上表现也很不错，超越了当时的其他模型。其结果在此略过。
剩余的带约束图像合成任务都可以看成是图像转图像任务，比如图像超分辨率是低质量图像到高质量图像的转换、语义生成是把语义分割图像转换成一幅合成图像。要添加这些约束，只需要把这些任务的输入图片和LDM原本的输入$z_t$拼接起来即可。比如对于图像超分辨率，可以把输入图片直接与隐空间图片$z_t$拼接，解码后图片会被自然上采样$f$倍；对于语义生成，可以把下采样$f$倍的语义分割图与$z_t$拼接。论文用这些任务上的实验证明了LDM的泛用性。由于这部分实验与LDM的主要知识无关，具体实验结果就不在此详细介绍了。
总结
论文末尾探讨了LDM的两大不足。首先，尽管LDM的计算需求比其他像素空间上的扩散模型要少得多，但受制于扩散模型本身的串行采样，它的采样速度还是比GAN慢上许多。其次，LDM使用了一个自编码器来压缩图像，重建图像带来的精度损失会成为某些需要精准像素值的任务的性能瓶颈。
论文最后再次总结了此方法的贡献。LDM的主要贡献其实只有两点：在不损失效果的情况下用两阶段的图像生成方法大幅提升了训练和采样效率、借助交叉注意力实现了各任务通用的约束机制。这两个贡献总结得非常精准。之后的Stable Diffusion之所以大受欢迎，第一就是因为它采样所需的计算资源不多，大众能使用消费级显卡完成图像生成，第二就是因为它强大的文字转图片生成效果。
我们再从知识学习的角度总结一下LDM。LDM的核心知识是DDPM和VQGAN。如果你能看懂之前这两篇论文，那你一下子就能明白LDM是的核心思想是什么，看论文时只需要精读交叉注意力约束机制那一段即可，其他实验内容在现在看来已经价值不大了。由于近两年有大量基于Stable Diffusion开发的工作，相比论文，阅读源代码的重要性会大很多。我们会在下一篇文章里详细学习Stable Diffusion的官方源码和最常用的Stable Diffusion第三方实现——Diffusers框架。