Stable Diffusion的解读(一)
文章目录
- Stable Diffusion的解读(一)
摘要
本周主要研究了Stable Diffusion模型的发展历程及其核心思想。Stable Diffusion结合了变分自编码器(VAE)和去噪扩散概率模型(DDPM)的优点,通过将图像生成拆分为压缩和生成两个阶段,显著提升了生成效率。模型在压缩阶段采用KL或VQ正则化,以减少过拟合现象。在生成阶段,利用DDPM生成高质量的图像。此外,Stable Diffusion通过引入交叉注意力机制,实现了多种约束条件下的图像生成。
Abstract
This week focused on the development and core concepts of the Stable Diffusion model. Stable Diffusion combines the strengths of Variational Autoencoders (VAE) and Denoising Diffusion Probabilistic Models (DDPM) by splitting image generation into compression and generation phases, significantly improving efficiency. The model employs KL or VQ regularization during compression to reduce overfitting. In the generation phase, DDPM is used to produce high-quality images. Additionally, Stable Diffusion incorporates cross-attention mechanisms to enable image generation under various constraints.
一、学习>机器学习部分
在2022年的这波AI绘画浪潮中,Stable Diffusion无疑是最受欢迎的图像生成模型。究其原因,第一,Stable Diffusion通过压缩图像尺寸显著提升了扩散模型的运行效率,使得每个用户能在自己的商业级显卡上运行模型;第二,有许多基于Stable Diffusion的应用,比如Stable Diffusion自带的文生图、图像补全,以及ControlNet、LoRA、DreamBooth等插件式应用;第三,得益于前两点,Stable Diffusion已经形成了一个庞大的用户社群,大家互相分享模型,交流心得。
不仅是大众,Stable Diffusion也吸引了大量科研人员,很多本来研究GAN的人纷纷转来研究扩散模型。然而,许多人在学习Stable Diffusion时却犯了难:又是公式扎堆的扩散模型,又是VAE,又是U-Net,这该怎么学起呀?
其实,一上来就读Stable Diffusion是很难读懂的。而如果你把之前的一些更基础的文章读懂,再回头来读Stable Diffusion,就会畅行无阻了。
1. Stable Diffusion的早期工作
1.1 从编码器谈起
包括Stable Diffusion在内,很多图像生成模型都可以看成是一种非常简单的模型——自编码器——的改进版。要谈Stable Diffusion是怎么逐渐诞生的,其实就是在谈自编码器是一步一步进化的。我们的学习就从自编码器开始。
尽管PNG、JPG等图像压缩方法已经非常成熟,但我们会想,会不会还有更好的图像压缩算法呢?图像压缩,其实就是找两个映射,一个把图片编码成压缩数据,另一个把压缩数据解码回图片。我们知道,神经网络理论上可以拟合任何映射。那我们干脆用两个神经网络来拟合两种映射,以实现一个图像压缩算法。负责编码的神经网络叫编码器(Encoder),负责解码的神经网络叫做解码器(Decoder)。
光定义了神经网络还不够,我们还需要给两个神经网络设置一个学习目标。在运行过程中,神经网络应该满足一个显然的约束:编码再解码后的重建图像应该和原图像尽可能一致,即二者的均方误差应该尽可能小。这样,我们只需要随便找一张图片,通过编码器和解码器得到重建图像,就能训练神经网络了。我们不需要给图片打上标签,整个训练过程是自监督的。所以我们说,整套模型是一个自编码器(Autoencoder,AE)。
图像压缩模型AE为什么会和图像生成扯上关系呢?你可以试着把AE的输入图像和编码器遮住,只看解码部分。把一个压缩数据解码成图像,换个角度看,不就是在根据某一数据生成图像嘛。
很可惜,AE并不是一个合格的图像生成模型。我们常说的图像生成,具体是指让程序生成各种各样的图片。为了让程序生成不同的图片,我们一般是让程序根据随机数(或是随机向量)来生成图片。而普通的AE会有过拟合现象,这导致AE的解码器只认得训练集里的图片经编码器解码出来的压缩数据,而不认得随机生成的压缩数据,进而也无法达到图像生成的要求。
所谓过拟合,就是指模型只能处理训练数据,而不能推广到一般的数据上。举一个极端的例子,如下图所示,编码器和解码器直接记忆了整个数据集,把所有图片压缩成了一个数字。也就是模型把编码器当成一个图片到数字的词典,把解码器当成一个数字到图片的词典。这样,不管数据集有多大,所有图片都可以被压缩成一个数字。这样的AE确实压缩能力很强,但它完全没用,因为它过拟合了,处理不了训练集以外的数据。
过拟合现象在普通版AE中是不可避免的。为了利用AE的解码器来生成图片,许多工作都在试图克服AE的过拟合现象。AE的改进思路很多,在这篇文章中,我们仅把AE的改进路线粗略地分成两种:解决过拟合问题以直接用AE做图像生成、用AE压缩图像间接实现图像生成。
1.2 第一条路线:VAE和DDPM
在第一条改进路线中,许多后续工作都试图用更高级的数学模型来解决AE的过拟合问题。变分自编码器(Variational Autoencoder, VAE) 就是其中的代表。
VAE对AE做了若干改动。第一,VAE让编码器的输出不再是一个确定的数据,而是一个正态分布中的一个随机数据。更具体一点,训练时,编码器会同时输出一个均值和方差。随后,模型会从这个均值和方差表达的正态分布里随机采样一个数据,作为解码器的输入。直观上看,这一改动就是在AE的基础上,让编码器多输出了一个方差,使得原AE编码器的输出发生了一点随机扰动。
这一改动可以缓解过拟合现象。这是为什么呢?我们可以这样想:原来的AE之所以会过拟合,是因为它强行记住了训练集里每一个数据的编码输出。现在,我们在VAE里让编码器不再输出一个固定值,而是随机输出一个在均值附近的值。这样的话,VAE就不能死记硬背了,必须要找出数据中的规律。
VAE的第二项改动是多添加一个学习目标,让编码器的输出和标准正态分布尽可能相似。前面我们谈过,图像生成模型一般会根据一个随机向量来生成图像。最常用的产生随机向量的方法是去标准正态分布里采样。也就是说,在用VAE生成图像时,我们会抛掉编码器,用下图所示的流程来生成图像。如果我们不约束编码器的输出分布,不让它输出一个和标准正态分布很相近的分布的话,解码器就不能很好地根据来自标准正态分布的随机向量生成图像了。
综上,VAE对AE做了两项改进:使编码器输出一个正态分布,且该分布要尽可能和标准正态分布相似。训练时,模型从编码器输出的分布里随机采样一个数据作为解码器的输入;图像采样(图像生成)时,模型从标准正态分布里随机采样一个数据作为解码器的输入。VAE的误差函数由两部分组成:原图像和重建图像的重建误差、编码器输出和标准正态分布之间的误差。VAE要最小化重建误差,最大化编码器输出与标准正态分布的相似度。
分布与分布之间的误差可以用一个叫KL散度的指标表示。所以,在上面那个误差函数公式中,负的相似度应该被替换成KL散度。VAE的这两项改动本质上都是在解决AE的过拟合问题,所以,VAE的改动可以被看成一种正则化方法。我们可以把VAE的正则化方法简称为KL正则化。
VAE确实能减轻AE的过拟合。然而,由于VAE只是让重建图像和原图像的均方误差(重建误差)尽可能小,而没有对重建图像的质量施加更多的约束,VAE的重建结果和图像生成结果都非常模糊。
以下是VAE在CelebA数据集上图像生成结果。
在众多对VAE的改进方法中,一个叫做去噪扩散概率模型(Denoising Diffusion Probabilistic Model, DDPM) 的图像生成模型脱颖而出。DDPM正是当今扩散模型的开山鼻祖。我们来看一下DDPM是怎样基于VAE对图像生成建模的。
VAE之所以效果不好,很可能是因为它的约束太少了。VAE的编码和解码都是用神经网络表示的。神经网络是一个黑盒,我们不好对神经网络的中间步骤施加约束,只好在编码器的输出(某个正态分布)和解码器的输出(重建图像)上施加约束。能不能让VAE的编码和解码过程更可控一点呢?
DDPM的设计灵感来自热力学:一个分布可以通过一系列简单的变化(如添加高斯噪声)逐渐变成另一个分布。恰好,VAE的编码器不正是想让来自训练集的图像(训练集分布)变成标准正态分布吗?既然如此,就不要用一个可学习的神经网络来表示VAE的编码器了,干脆用一些预定义好的加噪声操作来表示编码过程。可以从数学上证明,经过了多次加噪声操作后,最后的图像分布会是一个标准正态分布。
既然编码是加噪声,那解码时就应该去掉噪声。DDPM的解码器也不再是一个不可解释的神经网络,而是一个能预测若干个去噪结果的神经网络。
相比只有两个约束条件的VAE,DDPM的约束条件就多得多了。在DDPM中,第t个去噪操作应该尽可能抵消掉第t个加噪操作。
让我们来更具体地认识一下DDPM的学习目标。所谓添加噪声,就是在一个均值约等于当前图像的正态分布上采样。比如要对图像 x x x添加噪声,我们可以在 N ( 0.9 x , I ) N(0.9x,I) N(0.9x,I)这个分布里采样一张新图像。新的图像每个像素的均值是原来的0.9倍左右,且新图像会出现很多噪声。我们设 q ( x t ∣ x t − 1 ) q({x_t}|{x_{t - 1}}) q(xt∣xt−1)为第 t t t步加噪声的正态分布。经过一些数学推导,我们可以求出这一步操作的逆操作 q ( x t − 1 ∣ x t ) q({x_{t - 1}}|{x_t}) q(xt−1∣xt),这个加噪声逆操作也是一个正态分布。既然如此,我们可以设第 t t t步去噪声也为一个正态分布 p ( x t − 1 ∣ x t ) p({x_{t - 1}}|{x_t}) p(xt−1∣xt),让第 t t t步去噪声和第 t t t步加噪声的逆操作尽可能相似。
总结一下,DDPM对VAE做了如下改动:
相比于VAE,DDPM的编码过程和解码过程的定义更加明确,可以施加的约束更多。因此,它的生成效果会比VAE好很多。同时,DDPM和VAE类似,它在编码时会从分布里采样,而不是只输出一个固定值,不会出现AE的过拟合问题。
DDPM的生成效果确实很好。但是,由于DDPM始终会对同一个尺寸的数据进行操作,图像的尺寸极大地影响了DDPM的运行速度,用DDPM生成高分辨率图像需要耗费大量计算资源。因此,想要用DDPM生成高质量图像,还得经过另一条路线。
1.3 第二条路线:VQVAE
在AE的第二条改进路线中,一些工作干脆放弃使用AE做图像生成,转而利用AE的图像压缩能力,把图像生成拆成两步来做:先用AE的编码器把图像压缩成更小的图像,再用另一个图像生成模型生成小图像,并用AE的解码器把小图像重建回真实图像。
为什么会有这么奇怪的图像生成方法呢?这得从另一类图像生成模型讲起。在机器翻译模型Transformer横空出世后的一段时间里,有很多工作都想把Transformer用在图像生成上。但是,原本用来生成文本的Transformer无法直接应用在图像上。在自然语言处理(NLP)中,一个句子可以用若干个单词表示。而每个单词又是用一个整数表示。所以,Transformer生成句子时,实际上是在生成若干个离散的整数,也就是生成一个离散向量。而在图像生成模型中,每个像素的颜色值是一个连续的浮点数。想把Transformer直接用在图像生成上,就得想办法把图像用离散向量表示。我们知道,AE可以把图像编码成一个连续向量。能不能做一些修改,让AE把图像编码成一个离散向量呢?
Vector Quantised-Variational AutoEncoder (VQVAE) 就是一个能把图像编码成离散向量的AE(虽然作者在取名时用了VAE)。我们来简单看一下VQVAE是怎样把图像编码成离散向量的。
假设我们有了一个能编码出离散向量的AE。
由于神经网络不能很好地处理离散数据,我们要引入NLP里的通常做法,加一个把离散向量映射成连续向量的嵌入层。
现在我们再回头讨论怎么让编码器输出一个离散向量。我们可以让AE的解码器保持不变,还是输出一个连续向量,再通过一个「向量离散化」操作,把连续向量变成离散向量。这个操作会把编码器的输出对齐到嵌入层的向量上,其原理类似于把0.99和1.01离散化成1,只不过它是对向量整体考虑,而不是对每一个数单独考虑。向量离散化操作的具体原理我们不在此处细究。
忽略掉实现细节,我们可以认为VQVAE能够把图像压缩成离散向量。更准确地说,VQVAE能把图像等比例压缩成离散的「小图像」。压缩成二维图像而不是一维向量,能够保留原图像的一些空间特性,为之后第二步图像生成铺路。
整理一下,VQVAE是一个能把图像压缩成离散小图像的AE。为了用VQVAE生成图像,需要执行一个两阶段的图像生成流程:
- 训练时,先训练一个图像压缩模型(VQVAE),再训练一个生成压缩图像的模型(比如Transformer)
- 生成时,先用第二个模型生成出一个压缩图像,再用第一个模型的解码器把压缩图像复原成真实图像
之所以要执行两阶段的图像生成流程,而不是只用第二个模型生成大图像,有两个原因。第一个原因是前面提到的,Transformer等生成模型只支持生成离散图像,需要用另一个模型把连续的颜色值变成离散值以兼容这些模型。第二个原因是为了减少模型的运算量。以Transformer为例,Transformer的运算次数大致与像素数的平方成正比,拿Transformer生成高分辨率图像的运算开销是不可接受的。而如果用一个AE把图像压缩一下的话,用Transformer就可行了。
VQVAE给后续工作带来了三条启发:第一,可以用AE把图像压缩成离散向量;第二,如果一个图像生成模型生成高分辨率的图像的计算代价太高,可以先用AE把图像压缩,再生成压缩图像。这两条启发对应上一段提到的使用VQVAE的两条动机。
而第三条启发就比较有意思了。在讨论VQVAE的过程中,我们完全没有考虑过拟合的事。这是因为经过了向量离散化操作后,解码器的输入已经不再是编码器的输出,而是嵌入层里的向量了。这种做法杜绝了AE的死记硬背,缓解了过拟合现象。这样,我们可以换一个角度看待VQVAE:编码器还是AE的编码器,编码器的输出是连续向量,后续的向量离散化操作和嵌入层全部都是解码器的一部分。从这个角度看,VQVAE其实提出了一个由向量离散化和嵌入层组成的正则化模块。这个模块和VAE的KL散度约束一样,都解决了AE的过拟合问题。我们把VQVAE的正则化方法叫做VQ正则化。
VQVAE论文提出的图像生成方法效果一般。和普通的AE一样,VQVAE在训练时只用了重建误差来约束图像质量,重建图像的细节依然很模糊。且VQVAE配套的第二阶段图像生成模型不是较为强力的Transformer,而是一个基于CNN的图像生成模型。
后续的VQGAN论文对VQVAE进行了改进。对于一阶段的图像压缩模型,VQGAN在VQVAE的基础上引入了生成对抗网络(GAN)中一些监督误差,提高了图像压缩模型的重建质量;对于两阶段的图像生成模型,该方法使用了Transformer。凭借这些改动,VQGAN方法能够生成高质量的高清图片。并且,通过把额外的约束条件(如语义分割图像、文字)输入进Transformer,VQGAN方法能够实现带约束的图像生成。以下是VQGAN方法根据语义分割图像生成的高清图片。
图像生成模型可以是无约束或带约束的。无约束图像生成模型只需要输入一个随机向量,训练数据不需要任何标注,可以进行无监督训练。带约束图像生成模型会在无约束图像生成模型的基础上多加一些输入,并给每个训练图像打上描述约束的标签,执行监督训练。比如要训练文生图模型,就要给每个训练图片带上文字描述。
1.4 路线的交汇点:Stable Diffusion
看完上面这两条AE的改进路线,相信你已经能够猜出Stable Diffusion的核心思想了。让我们看看Stable Diffusion是怎么从这两条路径中汲取灵感的。
在发布了VQGAN后,德国的CompVis实验室开始探索起VQGAN的改进方法。VQGAN能把图像边长压缩16倍,而VQGAN配套的Transformer只能一次生成 16 × 16 16×16 16×16的图片。也就是说,整套方法一次只能生成 256 × 256 256×256 256×256的图片。为了生成分辨率更高的图片,VQGAN方法需要借助滑动窗口。能不能让模型一次性生成分辨率更高的图片呢?制约VQGAN方法生成分辨率的主要因素是Transformer。如果能把Transformer换成一个效率更高,能生成更高分辨率的图像的模型,不就能生成比 256 × 256 256×256 256×256更大的图片了吗?CompVis实验室开始把目光着眼于DDPM上。
于是,在发布VQGAN的一年后,CompVis实验室又发布了名为High-Resolution Image Synthesis with Latent Diffusion Models的论文,提出了一种叫做隐扩散模型(latent diffusion model, LDM) 的图像生成模型。通过与AI公司Stability AI合作,借助他们庞大的算力资源训练LDM,CompVis实验室发布了商业名为Stable Diffusion的开源文生图AI绘画模型。
LDM其实就是在VQGAN方法的基础上,把图像生成模型从Transformer换成了DDPM。或者从另一个角度说,为了让DDPM生成高分辨率图像,LDM利用了VQVAE的第二条启发:先用AE把图像压缩,再用DDPM生成压缩图像。LDM的AE一般是把图像边长压缩8倍,DDPM生成 64 × 64 64×64 64×64的压缩图像,整套LDM能生成 512 × 512 512×512 512×512的图像。
和Transformer不同,DDPM处理的图像是用连续向量表示的。因此,在LDM中使用VQGAN做图像压缩时,不一定需要向量离散化操作,只需要在AE的基础上加一点轻微的正则化就行。作者在实现LDM时讨论了两类正则化,一类是VAE的KL正则化,一类是VQ正则化(对应VQVAE的第三条启发),两种正则化都能取得不错的效果。
LDM依然可以实现带约束的图像生成。用DDPM替换掉Transformer后,额外的约束会输入进DDPM中。作者在论文中讨论了几种把约束输入进DDPM的方式。
在搞懂了早期工作后,理解Stable Diffusion的核心思想就是这么简单。让我们把Stable Diffusion的发展过程及主要结构总结一下。Stable Diffusion由两类AE的变种发展而来,一类是有强大生成能力却需要耗费大量运算资源的DDPM,一类是能够以较高保真度压缩图像的VQVAE。Stable Diffusion是一个两阶段的图像生成模型,它先用一个使用KL正则化或VQ正则化的VQGAN来实现图像压缩,再用DDPM生成压缩图像。可以把额外的约束(如文字)输入进DDPM以实现带约束图像生成。
1.5 相关论文
图像生成必读论文:
Neural Discrete Representation Learning (VQVAE): https://arxiv.org/abs/1711.00937
Taming Transformers for High-Resolution Image Synthesis (VQGAN): https://arxiv.org/abs/2012.09841
Denoising Diffusion Probabilistic Models (DDPM): https://arxiv.org/abs/2006.11239
图像生成选读论文:
Auto-Encoding Variational Bayes (VAE): https://arxiv.org/abs/1312.6114 提出VAE的文章。数学公式较多,只需要了解VAE的大致结构就好,不需要详细阅读论文。
Pixel Recurrent Neural Networks (PixelCNN): https://arxiv.org/abs/1601.06759 提出了一种拟合离散分布的图像生成模型,自回归图像生成模型的代表。这是VQVAE使用的第二阶段图像生成模型。有兴趣可以了解一下。
Deep Unsupervised Learning using Nonequilibrium Thermodynamics: https://arxiv.org/abs/1503.03585 DDPM的前作,首个提出扩散模型思想的文章。其核心原理和DDPM几乎完全一致,但是模型结构和优化目标不够先进,生成效果没有改进后的DDPM好。数学公式较多,不必细读,可以在学习DDPM时对比着阅读。
Denoising Diffusion Implicit Models (DDIM): https://arxiv.org/abs/2010.02502 一种加速DDPM采样的方法,广泛运用在包含Stable Diffusion在内的扩散模型中。推荐阅读。
Classifier-Free Diffusion Guidance: https://arxiv.org/abs/2207.12598 一种让扩散模型的输出更加贴近约束的方法,广泛运用在包含Stable Diffusion在内的扩散模型中,用于生成更符合文字描述的图片。推荐阅读。
Generative Adversarial Networks (GAN): https://arxiv.org/abs/1406.2661 以及 A Style-Based Generator Architecture for Generative Adversarial Networks (StyleGAN): https://arxiv.org/abs/1812.04948 可以了解一下GAN是怎么确保图像生成质量的,并认识CelebAHQ和FFHQ这两个常用的人脸数据集。
其他必读文章:
Deep Residual Learning for Image Recognition (ResNet): https://arxiv.org/abs/1512.03385 深度学习的经典文章。其中提出的残差连接被用到了DDPM中。
Attention Is All You Need (Transformer): https://arxiv.org/abs/1706.03762 深度学习的经典文章。其中提出的自注意力模块被用到了DDPM中。
其他选读文章:
Learning Transferable Visual Models From Natural Language Supervision (CLIP): https://arxiv.org/abs/2103.00020 提出了对齐文本和图像的方法。绝大多数文生图模型的核心。
U-Net: Convolutional Networks for Biomedical Image Segmentation (U-Net): https://arxiv.org/abs/1505.04597 一种被广泛运用的神经网络架构。DDPM的神经网络的主架构。U-Net的结构很简单,可以不用去读论文,直接看代码。
2. Latent Diffusion Models论文精读
2.1 摘要与引言
论文摘要的大意如下:扩散模型的生成效果很好,但是,在像素空间上训练和推理扩散模型的计算开销都很大。为了在不降低质量与易用性的前提下用较少的计算资源训练扩散模型,我们在一个预训练过的自编码器的隐空间上使用扩散模型。相较以往的工作,在这种表示下训练扩散模型首次在减少计算复杂度和维持图像细节间达到几近最优的平衡点,极大地提升了视觉保真度。通过向模型架构中引入交叉注意力层,我们把扩散模型变成了强大而灵活的带约束图像生成器,它支持常见的约束,如文字、边界框,且能够以纯卷积方式实现高分辨率的图像合成。我们的隐扩散模型(latent diffusion model, LDM) 在使用比像素扩散模型少得多的计算资源的前提下,在各项图像合成任务上取得最优成果或顶尖成果。
整理一下。论文提出了一种叫LDM的图像生成模型。论文想解决的问题是减少像素空间扩散模型的运算开销。为此,LDM借助了VQVAE「先压缩、再生成」的想法,把扩散模型用在AE的隐空间上,在几乎不降低生成质量的前提下减少了计算量。另外,LDM还支持带约束图像合成及纯卷积图像超分辨率。
在上一篇回顾LDM早期工作的文章中,我们已经理解了LDM想解决的问题及解决问题的思路。因此,在读完摘要后,我们接下来读文章时只需要关注LDM的两个创新点:
- LDM的AE是怎么设计以达到压缩比例与质量的平衡的。
- LDM怎么实现带约束的图像合成。
引言基本是摘要的扩写。首先,引言大致介绍了图像合成任务的背景,提及了扩散模型近期的突出表现。随后,引言介绍了本文想解决的主要问题:扩散模型的训练和推理太耗时了,需要在不降低效果的前提下减少扩散模型的运算量。最后,引言揭示了本工作的解决方法:使用类似VQGAN的两阶段图像生成方法。
引言的前两部分没有什么关键信息,而最后一部分介绍了本工作改进扩散模型的动机,值得一读。如下图所示,DDPM的论文展示了从不同去噪时刻的同一个噪声图像开始的不同生成结果,比如 x 750 x_{750} x750指从时刻 t = 750 t=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怎么用交叉注意力机制实现带约束的图像生成。
2.2 相关工作
作者主要从两个角度回顾了早期工作:不同架构的图像生成模型与两阶段的图像合成方法。其回顾逻辑与第1节Stable Diffusion的早期工作类似,在此就不过多介绍了。除了介绍早期工作外,作者重申了引言中的对比结果,强调了LDM相对于扩散模型的创新和相对于两阶段图像生成模型的创新。
2.3 方法
在方法章节中,作者先是大致介绍了使用LDM这种两阶段图像生成架构的优点,再分三部分详细介绍了论文的实现细节:图像压缩AE的实现、LDM的实现、约束的实现。开头的介绍和AE的实现相对比较重要,我们放在一起详细阅读;相对于DDPM,LDM几乎没有做任何修改,只是把要拟合的图片从真实图片换成了压缩图片,这一部分我们会快速浏览一遍;而添加约束的方法有所创新,我们会详细阅读一遍。
2.3.1 AE与两阶段图像生成模型
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)。而本身用于二维图像生成的LDM能更好地利用二维信息,因此可以处理更大的压缩图像(64×64)。这样,LDM的压缩模型的压缩程度不必那么重,其保真度会比VQGAN高。
看完了这一小节,我们来回头看第3节开头介绍了LDM的三项优点:1)通过规避在高维图像空间上训练扩散模型,作者开发出了一个因在低维空间上采样而计算效率大幅提升的扩散模型;2)作者发掘了扩散模型中来自U-Net架构的归纳偏置(inductive bias),使得它们能高效地处理有空间结构的数据(比如二维图像),避免像之前基于Transformer的方法一样使用激进、有损质量的压缩比例;3)本工作的压缩模型是通用的,它的隐空间能用来训练多种图像生成模型。第一个优点是相对于DDPM。第二个是优点是相对于使用Transformer的VQGAN,我们在上一段已经分析过了。第三个优点是相对于之前那些换一个任务就需要换一个压缩模型的两阶段图像生成模型。
归纳偏置可以简单理解为某个学习算法对一类数据的优势。比如CNN结构适合处理图像数据。
2.3.2 隐扩散模型(LDM)
在DDPM中,一个参数为 θ θ θ的神经网络 ϵ θ {\epsilon _\theta } ϵθ会根据当前时刻 t t t的带噪图片 x t x_t xt预测本时刻的噪声 ϵ θ ( x t , t ) {\epsilon _\theta }({x_t},t) ϵθ(xt,t)。网络的学习目标是让预测的噪声和真实的噪声 ϵ \epsilon ϵ一致。
L D M = E x , ϵ ∼ N ( 0 , 1 ) , t [ ∣ ∣ ϵ − ϵ θ ( x t , t ) ∣ ∣ 2 2 ] {L_{DM}} = {{\rm E}_{x,\epsilon \sim {\cal N}(0,1),t}}[||\epsilon - {\epsilon _\theta }({x_t},t)||_2^2] LDM=Ex,ϵ∼N(0,1),t[∣∣ϵ−ϵθ(xt,t)∣∣22]
LDM的原理和DDPM完全一样,只不过训练图片从像素空间上的真实图片 x 0 x_0 x0变成了隐空间上的压缩图片 z 0 z_0 z0,每一轮的带噪图片由 x t x_t xt变成了隐空间上的带噪图片 z t z_t zt。在训练时,相比DDPM,只需要多对 x 0 x_0 x0用一次编码器变成 z 0 z_0 z0即可。
L L D M = E ε ( x ) , ϵ ∼ N ( 0 , 1 ) , t [ ∣ ∣ ϵ − ϵ θ ( z t , t ) ∣ ∣ 2 2 ] {L_{LDM}} = {{\rm E}_{\varepsilon (x),\epsilon \sim {\cal N}(0,1),t}}[||\epsilon - {\epsilon _\theta }({z_t},t)||_2^2] LLDM=Eε(x),ϵ∼N(0,1),t[∣∣ϵ−ϵθ(zt,t)∣∣22]
2.3.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查询数据库里的数据并返回查询结果。注意力操作有几种用法,第一种用法是交叉注意力 C 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 ) CrossAttn(A,B)=Attention(W_Q⋅A,W_K⋅B,W_V⋅B) CrossAttn(A,B)=Attention(WQ⋅A,WK⋅B,WV⋅B),可以理解成数据 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_Q⋅A,W_K⋅A,W_V⋅A) SelfAttn(A)=Attention(WQ⋅A,WK⋅A,WV⋅A),可以理解成数据 A A A自己做了一次特征提取。
既然交叉注意力操作可以融合两类信息,何不把DDPM的自注意力层换成交叉注意力层,把 K K K, V V V换成来自约束的信息,以实现带约束图像生成呢?如下图所示,通过把用编码器 τ θ τ_θ τθ编码过的约束信息输入进扩散模型交叉注意力层的 K K K, V V V,LDM实现了带约束图像生成。
根据论文中实验的设计,对于作用于全局的约束,如文本描述,使用交叉注意力较好;对于有空间信息的约束,如语义分割图片,则用拼接的方式较好。
2.4 实验
在这一章里,作者按照介绍方法的顺序,依次探究了图像压缩模型、无约束图像生成、带约束图像合成的实验结果。我们主要关心前两部分的实验结果。
2.4.1 感知压缩程度的折衷
论文首先讨论了图像压缩模型在不同的下采样比例 f f f下的实验结果,其中 f ∈ 1 , 2 , 4 , 8 , 16 , 32 f∈{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过小时扩散模型把过多的精力放在了本应由压缩模型负责的感知压缩上; f f f过大时图像信息在压缩中损失过多。LDM-{4-16}的表现相对好一些。
在实验的第二部分中,作者比较了不同采样比例 f f f的LDM在CelebA-HQ(下图左侧)和ImageNet(下图右侧)上的采样速度和采样效果。下图中,横坐标为吞吐量,越靠右表示采样速度越快。同一个模型的不同实验结果表示使用不同DDIM采样步数时的实验结果,每一条线上的结果从右到左分别是DDIM采样步数取 10 , 20 , 50 , 100 , 200 {10,20,50,100,200} 10,20,50,100,200的采样结果(DDIM步数越少,采样速度越快,生成图片质量越低)。对于CelebA-HQ上的实验,若采样步数较多,则还是LDM-{4,8}效果较好,只有在采样步数较少时压缩比更高的LDM才有优势。而对于ImageNet上的实验, f f f太小或太大的结果都很差,整体上还是LDM-{4,8}的结果较好。
综上,根据实验,作者认为 f f f取适中的4或8比较妥当。下采样比例 f = 8 f=8 f=8也正是Stable Diffusion采用的配置。
2.4.2 图像生成效果
在这一节中,作者在几个常见的数据集上对比了LDM与其他模型的无约束图像生成效果。作者主要比较了两类指标:表示采样质量的FID和表示数据分布覆盖率的精确率及召回率(Precision-and-Recall)。
在介绍具体结果之前,先对这个不太常见的精确率及召回率指标做一个解释。精确率及召回率常用于分类等有确定答案的任务中,分别表示所有被分类为正的样本中有多少是分对了的、所有真值为正的样本中有多少是被成功分类成正的。而无约束图像生成中的精确率及召回率的解释可以参加论文Improved Precision and Recall Metric for Assessing
Generative Models。如下图所示,设真实分布为蓝色,生成模型的分布为红色,则红色样本落在蓝色分布的比例为精确率,蓝色样本落在红色分布的比例为召回率。简单来说,精确率能描述采样质量,召回率能描述生成分布与真实分布的覆盖情况。
接下来,我们回头来看论文展示的无约束图像生成对比结果,如下图所示。整体上看,LDM的表现还不错。虽然在FID指标上无法超过GAN或其他扩散模型,但是在精确率和召回率上还是颇具优势。唯一没有被LDM战胜的是LSUN-Bedrooms上的ADM模型,但作者提到,相比ADM,LDM只用了一半的参数,且只需四分之一的训练资源。
2.4.3 带约束图像合成
这一节里,作者展示了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 zt拼接起来即可。比如对于图像超分辨率,可以把输入图片直接与隐空间图片 z t z_t zt拼接,解码后图片会被自然上采样 f f f倍;对于语义生成,可以把下采样 f f f倍的语义分割图与 z t z_t zt拼接。论文用这些任务上的实验证明了LDM的泛用性。由于这部分实验与LDM的主要知识无关,具体实验结果就不在此详细介绍了。
2.5 论文总结
论文末尾探讨了LDM的两大不足。首先,尽管LDM的计算需求比其他像素空间上的扩散模型要少得多,但受制于扩散模型本身的串行采样,它的采样速度还是比GAN慢上许多。其次,LDM使用了一个自编码器来压缩图像,重建图像带来的精度损失会成为某些需要精准像素值的任务的性能瓶颈。
论文最后再次总结了此方法的贡献。LDM的主要贡献其实只有两点:在不损失效果的情况下用两阶段的图像生成方法大幅提升了训练和采样效率、借助交叉注意力实现了各任务通用的约束机制。这两个贡献总结得非常精准。之后的Stable Diffusion之所以大受欢迎,第一就是因为它采样所需的计算资源不多,大众能使用消费级显卡完成图像生成,第二就是因为它强大的文字转图片生成效果。
我们再从知识学习的角度总结一下LDM。LDM的核心知识是DDPM和VQGAN。如果你能看懂之前这两篇论文,那你一下子就能明白LDM是的核心思想是什么,看论文时只需要精读交叉注意力约束机制那一段即可,其他实验内容在现在看来已经价值不大了。由于近两年有大量基于Stable Diffusion开发的工作,相比论文,阅读源代码的重要性会大很多。
参考:
Stable Diffusion 解读(一):回顾早期工作
Stable Diffusion 解读(二):论文精读
总结
下次准备介绍实现Stable Diffusion的diffusers库。