在这篇文章中,我们来精读2022年Stable Diffusion的论文:High-Resolution Image Synthesis with Latent Diffusion Models,论文发表之后,Stable Diffusion就开始火了起来。

阅读这篇文章你可以了解到
文章目录
1 摘要与引言
论文摘要的大意如下:扩散模型的生成效果很好,但是,在像素空间上训练和推理扩散模型的计算开销都很大。为了在不降低质量与易用性的前提下用较少的计算资源训练扩散模型,我们在一个预训练过的自编码器的隐空间上使用扩散模型。相较以往的工作,在这种表示下训练扩散模型首次在减少计算复杂度和维持图像细节间达到几近最优的平衡点,极大地提升了视觉保真度。通过向模型架构中引入交叉注意力层,我们把扩散模型变成了强大而灵活的带约束图像生成器,它支持常见的约束,如文字、边界框,且能够以纯卷积方式实现高分辨率的图像合成。我们的隐扩散模型(latent diffusion model, LDM) 在使用比像素扩散模型少得多的计算资源的前提下,在各项图像合成任务上取得最优成果或顶尖成果。
整理一下。论文提出了一种叫LDM的图像生成模型。论文想解决的问题是减少像素空间扩散模型的运算开销。为此,LDM借助了VQVAE「先压缩、再生成」的想法,把扩散模型用在AE的隐空间上,在几乎不降低生成质量的前提下减少了计算量。另外,LDM还支持带约束图像合成及纯卷积图像超分辨率。所以我们接下来读文章时只需要关注LDM的两个创新点:
- 1)LDM的AE是怎么设计以达到压缩比例与质量的平衡的。
- 2)LDM怎么实现带约束的图像合成。
引言基本是摘要的扩写。首先,引言大致介绍了图像合成任务的背景,提及了扩散模型近期的突出表现。随后,引言介绍了本文想解决的主要问题:扩散模型的训练和推理太耗时了,需要在不降低效果的前提下减少扩散模型的运算量。最后,引言揭示了本工作的解决方法:使用类似VQGAN的两阶段图像生成方法。
1.1 改进DDPM扩散模型的动机
引言的前两部分没有什么关键信息,而最后一部分介绍了本工作改进扩散模型的动机,值得一读。如下图所示,DDPM的论文展示了从不同去噪时刻的同一个噪声图像开始的不同生成结果,比如 X 750 X_{750} X750指从 t = 750 t=750 t=750时刻的去噪图像开始,多次以不同随机数执行DDPM的反向过程,生成的多幅图像。LDM作者认为,DDPM的这一实验表明,扩散模型的图像生成分两个阶段:先是对语义进行压缩,再是对图像的感知细节压缩。正因此,随机对早期的噪声图像去噪,生成图像的内容会更多样;而随机对后期的噪声图像去噪,生成图像只是在细节上有所不同。LDM的作者认为,扩散模型的大量计算都浪费在了生成整幅图像的细节上,不如只让扩散模型描述比较关键的语义压缩部分,而让自编码器(AE)负责感知细节压缩部分。
引言在结尾总结了本工作的贡献:
- 1)相比之前按序列处理图像的纯Transformer的方法,扩散模型能更好地处理二维数据。因此,LDM生成隐空间图像时不需要那么重的压缩比例(比如DIV2K数据集上,LDM只需要将图像下采样4倍,而之前的纯Transformer方法要下采样8倍或16倍),图像在压缩时能有更高的保真度,整套方法能更高效地生成高分辨率图像。
- 2)在大幅降低计算开销的前提下在多项图像生成任务上取得了顶尖成果。
- 3)相比于之前同时训练图像压缩模型和图像生成模型的方法,该方法分步训练两个模型,训练起来更加简单。
- 4)对于有着稠密约束的任务(如超分辨率、补全、语义生成),该方法的模型能换成一个纯卷积版本的,且能生成边长为1024的图像。
- 5)该工作设计了一种通用的约束机制,该机制基于交叉注意力,支持多模态训练。作者训练了多种带约束的模型。
- 6)作者把工作开源了,并提供了预训练模型。
我们来整理一下这些贡献。读论文时,可以忽略第6条。第2条是成果,与方法设计无关。第1、3条主要描述了提出两阶段图像生成建模方法的贡献。第4条是把方法拓展到稠密约束任务的贡献。第5条是提出了新约束机制的贡献。所以,在学习论文的方法时,我们还是主要关注摘要里就提过的那两个创新点。在读完引言后,我们可以把阅读目标再细化一下:
- LDM的AE是怎么设计以达到压缩比例与质量的平衡的。与纯基于Transformer的VQGAN相比,它有什么不同。
- LDM怎么用交叉注意力机制实现带约束的图像生成。
2 方法
在方法章节中,作者先是大致介绍了使用LDM这种两阶段图像生成架构的优点,再分三部分详细介绍了论文的实现细节:图像压缩AE的实现、LDM的实现、约束的实现。开头的介绍和AE的实现相对比较重要,我们放在一起详细阅读;相对于DDPM,LDM几乎没有做任何修改,只是把要拟合的图片从真实图片换成了压缩图片,这一部分我们会快速浏览一遍;而添加约束的方法有所创新,我们会详细阅读一遍。
2.1 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
×
16
16 \times 16
16×16)。而本身用于二维图像生成的LDM能更好地利用二维信息,因此可以处理更大的压缩图像(
64
×
64
64 \times 64
64×64)。这样,LDM的压缩模型的压缩程度不必那么重,其保真度会比VQGAN高。
看完了3.1节,我们来回头看第3节开头介绍了LDM的三项优点:
- 1)通过规避在高维图像空间上训练扩散模型,作者开发出了一个可在低维空间上采样而计算效率大幅提升的扩散模型;
- 2)作者发掘了扩散模型中来自U-Net架构的归纳偏置(inductive bias),使得它们能高效地处理有空间结构的数据(比如二维图像),避免像之前基于Transformer的方法一样使用激进、有损质量的压缩比例;
- 3)本工作的压缩模型是通用的,它的隐空间能用来训练多种图像生成模型
第一个优点是相对于DDPM。第二个是优点是相对于使用Transformer的VQGAN,我们在上一段已经分析过了。第三个优点是相对于之前那些换一个任务就需要换一个压缩模型的两阶段图像生成模型。
归纳偏置可以简单理解为某个学习算法对一类数据的优势。比如CNN结构适合处理图像数据。
2.2 隐扩散模型(LDM)
在DDPM中,一个参数为
θ
\theta
θ的神经网络
ϵ
θ
\epsilon_\theta
ϵθ会根据当前时刻
t
t
t的带噪图片
X
t
X_t
Xt预测本时刻的噪声
ϵ
θ
(
X
t
,
t
)
\epsilon_\theta(X_t,t)
ϵθ(Xt,t)。网络的学习目标是让预测的噪声和真实的噪声一致。
L
D
M
=
∣
∣
ϵ
−
ϵ
θ
(
X
t
,
t
)
∣
∣
2
2
L_{DM}=||\epsilon-\epsilon_\theta(X_t,t)||^2_2
LDM=∣∣ϵ−ϵθ(Xt,t)∣∣22
LDM的原理和DDPM完全一样,只不过拟合的图片从像素空间上的真实图片 X 0 X_0 X0变成了隐空间上的压缩图片 Z 0 Z_0 Z0,每一轮的带噪图片X_t由变成了隐空间上的带噪图片 Z t Z_t Zt。在训练时,相比DDPM,只需要多对 X 0 X_0 X0用一次编码器变成 Z 0 Z_0 Z0即可。
L D M = ∣ ∣ ϵ − ϵ θ ( Z t , t ) ∣ ∣ 2 2 ; Z t = e n c o d e r ( X t ) L_{DM}=||\epsilon-\epsilon_\theta(Z_t,t)||^2_2;Z_t=encoder(X_t) LDM=∣∣ϵ−ϵθ(Zt,t)∣∣22;Zt=encoder(Xt)
2.3 约束机制
让模型支持带约束图像生成,其实就是想办法把额外的约束信息输入进扩散模型中。显然,最简单的添加约束的方法就是把额外的信息和扩散模型原本的输入 Z t Z_t Zt拼接起来。如果约束是一个值,就把相同的值拼接到 Z t Z_t Zt的每一个像素上;如果约束本身具有空间结构(如语音分割图片),就可以把约束重采样至和 Z t Z_t Zt一样的分辨率,再逐像素拼接。除了直接的拼接外,作者在LDM中还使用了另一种融合约束信息的方法。
DDPM中含有自注意力层。自注意力操作其实基于注意力操作 A t t e n t i o n ( Q , K , V ) Attention(Q,K,V) Attention(Q,K,V),它可以解释成一个数据库中存储了许多数据 V V V,数据的索引(键)是 K K K,现在要用查询 Q Q Q查询数据库里的数据并返回查询结果。注意力操作有几种用法,第一种用法是交叉注意力 C o r o s s A t t n ( A , B ) = A t t e n t i o n ( W Q A , W K B , W V B ) CorossAttn(A,B)=Attention(W_QA,W_KB,W_VB) CorossAttn(A,B)=Attention(WQA,WKB,WVB),可以理解成数据 A A A, B B B 做了一次信息融合;第二种用法是自注意力 S e l f A t t n ( A ) = A t t e n t i o n ( W Q A , , W K A , W V A ) SelfAttn(A)=Attention(W_QA,,W_KA,W_VA) SelfAttn(A)=Attention(WQA,,WKA,WVA),可以理解成数据自己做了一次特征提取。
既然交叉注意力操作可以融合两类信息,何不把DDPM的自注意力层换成交叉注意力层,把 K K K, V V V 换成来自约束的信息,以实现带约束图像生成呢?如下图所示,通过把用编码器 τ θ \tau_\theta τθ编码过的约束信息输入进扩散模型交叉注意力层的 K K K, V V V ,LDM实现了带约束图像生成。
根据论文中实验的设计,对于作用于全局的约束,如文本描述,使用交叉注意力较好;对于有空间信息的约束,如语义分割图片,则用拼接的方式较好。

3 Stable Diffusion的Unet结构组成
Stable Diffusion代码实现中的另一个重点是去噪网络U-Net的实现。接下俩,我们来逐步学习Stable Diffusion中的U-Net是怎么从最经典的纯卷积U-Net逐渐发展而来的。
3.1 原始Unet结构
最早的U-Net的结构如下图所示:

可以看出,U-Net的结构有以下特点:
- 整体上看,U-Net由若干个大层组成。特征在每一大层会被下采样成尺寸更小的特征,再被上采样回原尺寸的特征。整个网络构成一个U形结构。
- 下采样后,特征的通道数会变多。一般情况下,每次下采样后图像尺寸减半,通道数翻倍。上采样过程则反之。
- 为了防止信息在下采样的过程中丢失,U-Net每一大层在下采样前的输出会作为额外输入拼接到每一大层上采样前的输入上。这种数据连接方式类似于ResNet中的「短路连接」。
3.2 DDPM的Unet结构
DDPM则使用了一种改进版的U-Net。改进主要有两点:
- 原来的卷积层被替换成了ResNet中的残差卷积模块。每一大层有若干个这样的子模块。对于较深的大层,残差卷积模块后面还会接一个自注意力模块。
- 原来模型每一大层只有一个短路连接。现在每个大层下采样部分的每个子模块的输出都会额外输入到其对称的上采样部分的子模块上。直观上来看,就是短路连接更多了一点,输入信息更不容易在下采样过程中丢失。

3.3 Stable Diffusion的Unet结构
最后,LDM提出了一种给U-Net添加额外约束信息的方法:把U-Net中的自注意力模块换成交叉注意力模块。具体来说,DDPM的U-Net的自注意力模块被换成了标准的Transformer模块。约束信息 C C C可以作为Cross Attention的K, V输入进模块中。
Stable Diffusion的U-Net还在结构上有另一处修改,该U-Net的每一大层都有Transformer块,而不是只有较深的大层有。

至此,我们已经学完了Stable Diffusion的采样原理和U-Net结构。
这里回答一个疑问,为啥你使用SD模型的时候下载的都是一个cpkt,然后加载之后会有三个模型出来,分别是Unet扩散模型、CLIP模型、VAE模型,是因为这些模型的变量被维护成一个类了,这点可以从源代码中得到答案
4 实验
4.1 感知压缩程度的折衷
论文首先讨论了图像压缩模型在不同的下采样比例 f f f下的实验结果,其中 f ∈ [ 1 , 2 , 4 , 8 , 16 , 32 ] f \in {[1,2,4,8,16,32]} f∈[1,2,4,8,16,32]。这些实验分两部分,第一部分是训练速度上的实验,第二部分是采样速度与效果上的实验。
在ImageNet上以不同下采样比例 f f f训练一定步数后LDM的采样指标对比结果如下图所示。其中,FID指标越低越好,Inception Score越高越好。结果显示,无论下采样比例 f f f是过大还是过小都会降低训练速度。作者分析了 f f f较小或较大时训练速度慢的原因:
- f f f过小时扩散模型把过多的精力放在了本应由压缩模型负责的感知压缩上;
- 过大时图像信息在压缩中损失过多。LDM在{4-16}的表现相对好一些。

在实验的第二部分中,作者比较了不同采样比例
f
f
f的LDM在CelebA-HQ(下图左侧)和ImageNet(下图右侧)上的采样速度和采样效果。下图中,横坐标为吞吐量,越靠右表示采样速度越快。同一个模型的不同实验结果表示使用不同DDIM采样步数时的实验结果,每一条线上的结果从右到左分别是DDIM采样步数取{10,20,50,100,200}的采样结果(DDIM步数越少,采样速度越快,生成图片质量越低)。对于CelebA-HQ上的实验,若采样步数较多,则还是LDM-{4-8}效果较好,只有在采样步数较少时压缩比更高的LDM才有优势。而对于ImageNet上的实验,
f
f
f太小或太大的结果都很差,整体上还是LDM-{4-8}的结果较好。

综上,根据实验,作者认为
f
f
f取适中的或比较妥当。下采样比例
f
=
8
f=8
f=8也正是Stable Diffusion采用的配置。所以SD生成512x512尺寸的图像的效果是最优的,因为潜空间的图像大小是64,64x8=512。
4.2 带约束图像生成
这一节里,作者展示了LDM的文生图能力。论文中的LDM用了一个从头训练的基于Transformer的文本编码器,与后续使用CLIP的Stable Diffusion差别较大。这一部分的结果没那么重要,大致看一看就好。
本文的文生图模型是一个在LAION-400M数据集上训练的KL约束LDM。它的文本编码器是一个Transformer,编码后的特征会以交叉注意力的形式传入LDM。采样时,LDM使用了Classifier-Free Guidance。
CFG 是什么?
在实际操作中,当使用 Classifier-Free Guidance
时,生成过程会考虑来自有条件和无条件模型的两种预测。通过调节两者的混合比例(通常称为 guidance scale
或者说引导强度),可以控制生成结果与条件信息之间的匹配程度。
3. 为什么使用 CFG?
- 高效性:不需要额外训练分类器,减少了计算成本。
- 灵活性:可以通过调整权重 www 来控制生成结果的质量与条件的符合度。
- w=1:表示直接使用条件生成。
- w>1:增强条件的影响力,但可能导致样本多样性降低。
- w<1:弱化条件的影响,增加生成样本的多样性。
- 效果好:实验证明 CFG 能显著提高生成模型在扩散任务中的性能,使生成结果更加逼真并且符合条件。
LDM与其他模型的文生图效果对比如下图所示。虽然这个版本的LDM并没有显著优于其他模型,但它的参数量是最少的。

剩余的带约束图像合成任务都可以看成是图像转图像任务,比如
- 1)图像超分辨率是低质量图像到高质量图像的转换
- 2)语义生成是把语义分割图像转换成一幅合成图像。
要添加这些约束,只需要把这些任务的输入图片和LDM原本的输入拼接起来即可。比如
- 1)对于图像超分辨率,可以把输入图片直接与隐空间图片拼接,解码后图片会被自然上采样倍;
- 2)对于语义生成,可以把下采样倍的语义分割图与拼接。论文用这些任务上的实验证明了LDM的泛用性。
参考资料
https://mp.weixin.qq.com/s/TBvbtP4JRsQNx-O62mHWvQ
https://arxiv.org/abs/2006.11239
https://arxiv.org/abs/2112.10752
推荐阅读
《VQ-VAE》:Stable Diffusion设计的架构源泉
《VQGAN》:VQVAE+Transformer它的潜力直接促成了Stable Diffusion的诞生
《DDPM》论文精读:开创了Stable Diffusion扩散AI绘画新时代
《Transformer》论文精读:只用注意力机制,Transformer奠定了大模型网络架构的基石

2972

被折叠的 条评论
为什么被折叠?



