论文链接:High-Resolution Image Synthesis with Latent Diffusion Models
官方实现:CompVis/latent-diffusion、CompVis/stable-diffusion
视频讲解:一个视频看懂VAE的原理以及关于latent diffusion的思考
前言
目前的扩散模型范式基本上都变成做Latent Diffusion,包括像DIT(diffusion Transformer),实际上也是基于Latent Diffusion的思想,它只是把里面的UNet换成了Transformer。
传统的扩散模型一般是输入一张图像 x x x,然后对它进行增噪,再通过去噪采样得到一个新的图像 x ^ \hat x x^。但这个过程的计算都是在Pixel Space上面进行的,采样空间太大,学习的噪声维度和图像的维度是相同的。当进行高分辨率图像生成时,需要的计算资源会急剧增加,虽然 DDIM 等工作已经对此有所改善,但效果依然有限。
因此就产生了Latent Space上进行扩散操作的想法,经过测试证明是可以的。但压缩到一个比较小的程度,比如说压缩到32*32大小,会有信息的损失。作者认为压缩分为两种:感知压缩和语义压缩。
作者认为感知压缩损失的这些信息实际上是很好学到的,因为它有很强的相关性。但是整体语义压缩损失实际上是比较重要的,在Latent Diffusion模型里,实际上是用扩散模型去学习语义。
实际上的过程:
- 先把 x x x编码成一个隐藏向量 z z z
- 给 z z z增噪去噪,学得 z ^ \hat z z^
- 解码得到图像 x ^ \hat x x^
注意:
Encoder和Decoder的过程实际上是在做感知方面的压缩和解码,中间的扩散模型处理语义信息Encoder和Decoder就来自于VAE,Variational Auto Encoder其实是从AutoEncoder来的,用了一些方法把自编码器变成了生成模型,这个方法在2013年提出。实际上VAE的结果是不太好的,那为什么把中间换成扩散模型效果会变好?就这说明编码器和解码器的能力是足够的,问题出在中间的部分。
所以学习Latent Diffusion之前必须了解VAE的原理
问题定义
什么是AE(Auto Encoder)
AutoEncoder自编码器,它实际上是一种自监督学习图像特征的一个方法,它的原理很简单:
首先有一个 x i x_i xi,把它输入到Encoder中编码成一个 z i z_i zi,再把这个 z i z_i zi输入到Decoder中,给它解码成 x ^ i \hat x_i x^i,这就是重构图像,然后再去跟真实的 x i x_i xi做重构损失的计算,最后来优化编码器和解码器。如果训练好给定一个 z i z_i zi,就可以得到重构的比较好的 x i x_i xi,就可以认为 z i z_i zi实际上已经是学到了 x i x_i xi中比较丰富的细节特征了。
但是AE为什么不能作为生成模型?
假设现在有一个数据集它里边有n个图像 [ x 1 , x 2 . . . . x n ] {[x_1,x_2....x_n]} [x1,x2....xn],通过Encoder可以给它编码成 [ z 1 , z 2 . . . . z n ] {[z_1,z_2....z_n]} [z1,z2....zn]。因为生成模型有一个核心的诉求:生成一个全新的图像,即数据集里没有出现过的。现在如果想要一个全新的 Z Z Z,但肯定不能从 [ z 1 , z 2 . . . . z n ] {[z_1,z_2....z_n]} [z1,z2....zn]中来,但是也不能随便来。如果知道 Z Z Z服从哪一种分布,就可以从这个分布中取到一个新的 z z z,就可以通过解码得到一张新的图片。但训练时没有对 z z z的分布进行约束,导致 z z z的分布可能是不规律的,就不能生成新的图像。
假设 z z z服从正态分布 P ( z ) ∽ N ( 0 , 1 ) P(z)∽N(0,1) P(z)∽N(0,1)
如果假设 z z z服从正态分布 P ( z ) ∽ N ( 0 , 1 ) P(z)∽N(0,1) P(z)∽N(0,1)的话,很难确定正态分布N(0,1)中每个点和真实数据中每个点的对应关系,即 x i x_i xi和 z i z_i zi之间的对应关系
假设已经知道后验概率 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi)
此时不再去找确切的 z z z了,而是假设已经知道后验概率 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi),如果知道了这个分布,就可以从中取得一个新的 z i z_i zi,然后 z i z_i zi再输入到Decoder中,输出的 x ^ i \hat x_i x^i就可以跟 x i x_i xi进行重构损失的计算来优化Decoder。
思路的转变
这时思路就从想优化的auto encoder变成一个生成模型变成了需要找到 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi)这么一个后验概率,接下来就可以通过变分贝叶斯的方法找到这个分布。
一、VAE
1. 定义
变分贝叶斯:如下图,有一个分布集合,其中有一点 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi),其中有一个已知分布族,这个分布符合高斯分布,现在希望从这个分布族里找到一个最接近 P P P的分布并做近似。如果最开始取的是 x 1 x_1 x1点的分布,然后通过不断的优化,慢慢找到 x 2 x_2 x2这个地方。这个分布就是离 P P P的距离是最近的,所以最终通过优化找到这个分布。
从上面的高斯分布族取出一个分布 q φ ( z ∣ x i ) q_φ(z|x_i) qφ(z∣xi),此时需要优化 φ ,就需要找到已知分布中离已知 P 最近的一个分布。自然就想到一个度量方法,那就是KL散度。现在的问题就变成了最小化 K L ( q φ ( z ∣ x i ) ∣ ∣ P θ ( z ∣ x i ) ) KL(q_φ(z|x_i) || P_θ(z|x_i)) KL(qφ(z∣xi)∣∣Pθ(z∣xi))。如果以KL散度作为这个距离的衡量标准的话(虽然它不是一个严格意义的距离,它不服从对称性),这个方法就叫做变分贝叶斯法。
2. 变分下界
想要最小化 K L ( q φ ( z ∣ x i ) ∣ ∣ P θ ( z ∣ x i ) ) KL(q_φ(z|x_i) || P_θ(z|x_i)) KL(qφ(z∣xi)∣∣Pθ(z∣xi)),但是这里面有一个难点就是这个我们不知道 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi),所以就很难去通过计算方法求得它,另一种方法就是变分下界方法。
根据KL散度定义展开:
K L ( q φ ( z ∣ x i ) ∣ ∣ P θ ( z ∣ x i ) ) = ∫ q φ ( z ∣ x i ) l o g q φ ( z ∣ x i ) P θ ( z ∣ x i ) d z = ∫ q φ ( z ∣ x i ) l o g ( q φ ( z ∣ x i ) ) d z − ∫ q φ ( z ∣ x i ) l o g ( P θ ( z ∣ x i ) ) d z = E q φ ( z ∣ x i ) [ l o g ( q φ ( z ∣ x i ) ) ] − E q φ ( z ∣ x i ) [ l o g ( P θ ( z , x i ) ) ] + E q φ ( z ∣ x i ) [ l o g ( P θ ( x i ) ) ] = E q φ ( z ∣ x i ) [ l o g ( q φ ( z ∣ x i ) ) ] − E q φ ( z ∣ x i ) [ l o g ( P θ ( z , x i ) ) ] + l o g ( P θ ( x i ) ) = − E L B O + l o g ( P θ ( x i ) ) ⇨ l o g ( P θ ( x i ) ) = K L + E L B O KL(q_φ(z|x_i) || P_θ(z|x_i)) \\ = \int q_φ(z|x_i) log \frac{q_φ(z|x_i)}{P_θ(z|x_i)} dz \\ = \int q_φ(z|x_i) log (q_φ(z|x_i))dz - \int q_φ(z|x_i) log (P_θ(z|x_i))dz \\ = E q_φ(z|x_i) [log (q_φ(z|x_i))] - E q_φ(z|x_i) [log (P_θ(z,x_i))] +E q_φ(z|x_i) [log (P_θ(x_i))] \\ = E q_φ(z|x_i) [log (q_φ(z|x_i))] - E q_φ(z|x_i) [log (P_θ(z,x_i))] +log (P_θ(x_i)) \\ = -ELBO + log (P_θ(x_i)) \\ ⇨ log (P_θ(x_i)) = KL + ELBO KL(qφ(z∣xi)∣∣Pθ(z∣xi))=∫qφ(z∣xi)logPθ(z∣xi)qφ(z∣xi)dz=∫qφ(z∣xi)log(qφ(z∣xi))dz−∫qφ(z∣xi)log(Pθ(z∣xi))dz=Eqφ(z∣xi)[log(qφ(z∣xi))]−Eqφ(z∣xi)[log(Pθ(z,xi))]+Eqφ(z∣xi)[log(Pθ(xi))]=Eqφ(z∣xi)[log(qφ(z∣xi))]−Eqφ(z∣xi)[log(Pθ(z,xi))]+log(Pθ(xi))=−ELBO+log(Pθ(xi))⇨log(Pθ(xi))=KL+ELBO
- 第三行到第四行根据数学期望展开,后一项的 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi)展开为 P θ ( z , x i ) / P θ ( x i ) P_θ(z,x_i)/P_θ(x_i) Pθ(z,xi)/Pθ(xi)
- 第四行的最后一项,是服从 q φ ( z ∣ x i ) q_φ(z|x_i) qφ(z∣xi)的数学期望,但 E q φ ( z ∣ x i ) [ l o g ( P θ ( x i ) ) ] E q_φ(z|x_i) [log (P_θ(x_i))] Eqφ(z∣xi)[log(Pθ(xi))]里没有 z z z,所以这项实际上就是 l o g ( P θ ( x i ) ) log (P_θ(x_i)) log(Pθ(xi))
- 第五行的 E q φ ( z ∣ x i ) [ l o g ( q φ ( z ∣ x i ) ) ] − E q φ ( z ∣ x i ) [ l o g ( P θ ( z , x i ) ) ] E q_φ(z|x_i) [log (q_φ(z|x_i))] - E q_φ(z|x_i) [log (P_θ(z,x_i))] Eqφ(z∣xi)[log(qφ(z∣xi))]−Eqφ(z∣xi)[log(Pθ(z,xi))]被称为变分下界即 − E L B O -ELBO −ELBO, l o g ( P θ ( x i ) ) log (P_θ(x_i)) log(Pθ(xi))被称为边界似然
- 第七行通过移项得
我们的目标是最小化 K L KL KL散度,因为在上面的推导过程中 z z z是随机变量,而 l o g ( P θ ( x i ) ) log (P_θ(x_i)) log(Pθ(xi))中不含 z z z,说明在这个里面它是一个定值,所以在 l o g ( P θ ( x i ) ) = K L + E L B O log (P_θ(x_i)) = KL + ELBO log(Pθ(xi))=KL+ELBO中最小化KL散度的任务就转换成了最大化 E L B O ELBO ELBO
3. 最大化 ELBO
E L B O = E q φ ( z ∣ x i ) [ l o g ( P θ ( z , x i ) ) ] − E q φ ( z ∣ x i ) [ l o g ( q φ ( z ∣ x i ) ) ] = E q φ ( z ∣ x i ) [ l o g P θ ( z ) P θ ( x i ∣ z ) ] − E q φ ( z ∣ x i ) [ l o g ( q φ ( z ∣ x i ) ) ] = E q φ ( z ∣ x i ) [ l o g P θ ( z ) ] + E q φ ( z ∣ x i ) [ l o g P θ ( x i ∣ z ) ] − E q φ ( z ∣ x i ) [ l o g ( q φ ( z ∣ x i ) ) ] = − ( E q φ ( z ∣ x i ) [ l o g ( q φ ( z ∣ x i ) ) ] − E q φ ( z ∣ x i ) [ l o g P θ ( z ) ] ) + E q φ ( z ∣ x i ) [ l o g P θ ( x i ∣ z ) ] = − E q φ ( z ∣ x i ) [ l o g ( q φ ( z ∣ x i ) ) l o g P θ ( z ) ] + E q φ ( z ∣ x i ) [ l o g P θ ( x i ∣ z ) ] = − K L ( q φ ( z ∣ x i ) ∣ ∣ P θ ( z ) ) + E q φ ( z ∣ x i ) [ l o g P θ ( x i ∣ z ) ] ELBO = E q_φ(z|x_i) [log (P_θ(z,x_i))] - E q_φ(z|x_i) [log (q_φ(z|x_i))] \\ = E q_φ(z|x_i) [logP_θ(z)P_θ(x_i|z)] - E q_φ(z|x_i) [log (q_φ(z|x_i))] \\ = E q_φ(z|x_i) [logP_θ(z)] + E q_φ(z|x_i)[logP_θ(x_i|z)] - E q_φ(z|x_i) [log (q_φ(z|x_i))] \\ = -(E q_φ(z|x_i) [log (q_φ(z|x_i))] - E q_φ(z|x_i) [logP_θ(z)]) + E q_φ(z|x_i)[logP_θ(x_i|z)] \\ = -E q_φ(z|x_i) [\frac{log (q_φ(z|x_i))}{logP_θ(z)}] + E q_φ(z|x_i)[logP_θ(x_i|z)] \\ = -KL(q_φ(z|x_i) || P_θ(z)) + E q_φ(z|x_i)[logP_θ(x_i|z)] ELBO=Eqφ(z∣xi)[log(Pθ(z,xi))]−Eqφ(z∣xi)[log(qφ(z∣xi))]=Eqφ(z∣xi)[logPθ(z)Pθ(xi∣z)]−Eqφ(z∣xi)[log(qφ(z∣xi))]=Eqφ(z∣xi)[logPθ(z)]+Eqφ(z∣xi)[logPθ(xi∣z)]−Eqφ(z∣xi)[log(qφ(z∣xi))]=−(Eqφ(z∣xi)[log(qφ(z∣xi))]−Eqφ(z∣xi)[logPθ(z)])+Eqφ(z∣xi)[logPθ(xi∣z)]=−Eqφ(z∣xi)[logPθ(z)log(qφ(z∣xi))]+Eqφ(z∣xi)[logPθ(xi∣z)]=−KL(qφ(z∣xi)∣∣Pθ(z))+Eqφ(z∣xi)[logPθ(xi∣z)]
- 第二行将 l o g ( P θ ( z , x i ) ) log (P_θ(z,x_i)) log(Pθ(z,xi))展开为 P θ ( z ) P θ ( x i ∣ z ) P_θ(z)P_θ(x_i|z) Pθ(z)Pθ(xi∣z);
- 第三行再 l o g log log展开;
- 第四行移项并合并得到第五行;
- 第五行中的第一项就是前面的 K L KL KL散度;后一项 l o g P θ ( x i ∣ z ) logP_θ(x_i|z) logPθ(xi∣z),为给定 z z z之后, x i x_i xi的最大似然,在生成模型中最大似然用模型来学。
推理步骤:
此时最大化 E L B O ELBO ELBO的任务就变成了最小化 K L ( q φ ( z ∣ x i ) ∣ ∣ P θ ( z ) ) KL(q_φ(z|x_i) || P_θ(z)) KL(qφ(z∣xi)∣∣Pθ(z)),同时最大化 E q φ ( z ∣ x i ) [ P θ ( x i ∣ z ) ] E q_φ(z|x_i)[P_θ(x_i|z)] Eqφ(z∣xi)[Pθ(xi∣z)]
q φ ( z ∣ x i ) ∽ N ( μ , σ 2 ) q_φ(z|x_i)∽N(\mu,\sigma^2) qφ(z∣xi)∽N(μ,σ2)是从高斯分布中取的随便一个分布
当时我们希望 q φ ( z ∣ x i ) ∽ N ( μ , σ 2 ) q_φ(z|x_i)∽N(\mu,\sigma^2) qφ(z∣xi)∽N(μ,σ2)和真实的后验概率越近越好
通过推导发现, q φ ( z ∣ x i ) ∽ N ( μ , σ 2 ) q_φ(z|x_i)∽N(\mu,\sigma^2) qφ(z∣xi)∽N(μ,σ2)需要和先验概率 P θ ( z ) P_θ(z) Pθ(z)越近越好,而最开始我们已经假设了 P θ ( z ) ∽ N ( 0 , 1 ) P_θ(z)∽N(0,1) Pθ(z)∽N(0,1),这就是为什么我们希望 q φ ( z ∣ x i ) ∽ N ( μ , σ 2 ) q_φ(z|x_i)∽N(\mu,\sigma^2) qφ(z∣xi)∽N(μ,σ2)逼近标准正态分布
4. 思路整理
- 先希望得到 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi),目的是得到与 x i x_i xi对应的 z z z,这样好优化decode
- 用变分贝叶斯的方法得到 q φ ( z ∣ x i ) q_φ(z|x_i) qφ(z∣xi), q φ ( z ∣ x i ) q_φ(z|x_i) qφ(z∣xi)就是离 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi)最近的
- 用 K L KL KL散度优化这两个之间的距离
- 通过推导转化为 E L B O ELBO ELBO
- E L B O ELBO ELBO又转化为 − K L ( q φ ( z ∣ x i ) ∣ ∣ P θ ( z ) ) + E q φ ( z ∣ x i ) [ P θ ( x i ∣ z ) ] -KL(q_φ(z|x_i) || P_θ(z)) + E q_φ(z|x_i)[P_θ(x_i|z)] −KL(qφ(z∣xi)∣∣Pθ(z))+Eqφ(z∣xi)[Pθ(xi∣z)]。第一项就是前面的 K L KL KL散度;后一项 l o g P θ ( x i ∣ z ) logP_θ(x_i|z) logPθ(xi∣z),为给定 z z z之后, x i x_i xi的最大似然,在生成模型中最大似然用模型来学。所以 P θ ( x i ∣ z ) P_θ(x_i|z) Pθ(xi∣z)就是decode; l o g P θ ( x i ∣ z ) logP_θ(x_i|z) logPθ(xi∣z)前面有一个期望值 E E E,这个期望是服从 q φ ( z ∣ x i ) q_φ(z|x_i) qφ(z∣xi)的。
换句话说,如果在每次编码中得到的 q φ ( z ∣ x i ) ∽ N ( μ , σ 2 ) q_φ(z|x_i)∽N(\mu,\sigma^2) qφ(z∣xi)∽N(μ,σ2)分布后,要从这个里面重新取一个 z z z,通过蒙特拉洛的方法,让每次取得的 z z z都能最大的还原 x i x_i xi,这就是为什么VAE的训练过程是这样的(下一节解释)。
5. VAE的训练过程
VAE的训练过程如下图:
- 先输入一个 x x x到encode
- 得到一个 μ , σ \mu,\sigma μ,σ,通过这两个值得到高斯分布 N ( μ , σ 2 ) N(\mu,\sigma^2) N(μ,σ2)
- 从高斯分布中采样一个 z z z,再把这个 z z z输入到decode中得到 x ^ \hat x x^,然后和 x x x做重构损失计算
- 期间通过 K L KL KL散度尽量让 N ( μ , σ 2 ) N(\mu,\sigma^2) N(μ,σ2)和N(0,1)做近似 K L ( N ( μ , σ 2 ) ∣ N ( 0 , 1 ) ) KL(N(\mu,\sigma^2) | N(0,1)) KL(N(μ,σ2)∣N(0,1))
注意:
- N ( μ , σ 2 ) N(\mu,\sigma^2) N(μ,σ2)就是前面讲的 q φ ( z ∣ x i ) q_φ(z|x_i) qφ(z∣xi),对 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi)做最佳近似, N ( 0 , 1 ) N(0,1) N(0,1)是先验分布 P θ ( z ) P_θ(z) Pθ(z)
- 找到了真实的后验概率 P θ ( z ∣ x i ) P_θ(z|x_i) Pθ(z∣xi)的最佳近似 q φ ( z ∣ x i ) q_φ(z|x_i) qφ(z∣xi)之后,就可以从 q φ ( z ∣ x i ) q_φ(z|x_i) qφ(z∣xi)中去一个 z z z,我们认为这个 z z z是从后验分布离拿到的,所以就和 x i x_i xi是对应的;
- z z z通过decode后得到的 x ^ \hat x x^,就可以和真实的 x x x做重构损失计算,这一步就是前面推导出的 E q φ ( z ∣ x i ) [ l o g P θ ( x i ∣ z ) ] E q_φ(z|x_i)[logP_θ(x_i|z)] Eqφ(z∣xi)[logPθ(xi∣z)]
- 有的论文认为不断的从高斯分布中采样 z 1 , z 2 z_1,z_2 z1,z2,让它们都去重构 x i x_i xi,这也是一种数据增强的方法,这就使得decode更具有泛化性。
二、Latent Space
Latent Space是什么样的呢?
这是在MNIST数据集上做了一个训练,图里一共有10种颜色,代表是0到9,做了对比试验。这个就是隐空间的分布,这是一个二维的分布,实际上隐空间的维度可能是很高的,这里只是做了一个投影。
- 只有重构损失,即AE:有区域,无规则,不服从正态分布,所以很难取出一个 z z z;
- 只有 K L KL KL散度:无区域,有规则,即正态分布;
- VAE:有区域,有规则;
从这个图可知:
- 重构损失会导致最终的特征是否有区域性;
- K L KL KL散度对应最终的特征是否有规则;
- 虽然上述分布是在隐空间中的,但实际上也是一种特征图,AE就是在学特征,其中的每个向量都是特征向量,而相似的特征必然在一个维度上,将其投影到平面上是比较相近的;
- 图3可以看到这是一个从左上角到右下角连续的过程,最后生成了数字2;
- K L KL KL使得隐空间符合设好的已知分布,重构损失会学到一定的特征,而学特征就是Decode的过程,如果给定了一个 z z z,它能很好的还原一个 x x x,那么我们认为 z z z已经包含了 x x x比较丰富的特征。
- 图2到3加了重构损失后,就是对图2做了聚类的操作,把图像特征更像的聚集到一起,但是并不影响整体的分布。
三、LDM vs VAE
这两个模型到底差在哪?VAE的效果很一般,虽然LDM也用了VAE编码器解码器的结构,但效果好很多。
LDM训练过程简单示意图:
VAE训练过程简单示意图:
问题一(为什么LDM更好):
- 我们认为LDM和VAE的先验概率都是服从 P ( z ) ∽ N ( 0 , 1 ) P(z)∽N(0,1) P(z)∽N(0,1)
- 在训练VAE时,用了最小化 K L ( q ∣ ∣ P ( z ) ) KL(q || P(z)) KL(q∣∣P(z))的方法,让数据尽量接近 N ( 0 , 1 ) N(0,1) N(0,1), q q q是后验概率 P ( z ∣ x i ) P(z|x_i) P(z∣xi)的近似值。这里就会产生一个问题, P ( z ∣ x i ) P(z|x_i) P(z∣xi)不一定是真实的高斯分布,即后验概率,所以在做 q q q到 P ( z ∣ x i ) P(z|x_i) P(z∣xi)的近似值时就会产生偏差,生成效果就不是很好;
- Diffusion的训练过程时不断的增噪,最后得到 N ( 0 , 1 ) N(0,1) N(0,1),采样时直接选取一个 x ^ i \hat x_i x^i是可以直接生成图像的,所以扩散模型通过不断的增噪去噪学得了标准正态分布 N ( 0 , 1 ) N(0,1) N(0,1)到真实空间 P ( x ) P(x) P(x)的映射关系;
- 扩散模型可以学任意分布到标准正态之间的映射,VAE是假设的 P ( z ∣ x i ) P(z|x_i) P(z∣xi),用符合高斯分布的 q q q逼近它,因为 P ( z ∣ x i ) P(z|x_i) P(z∣xi)具体是什么分布是不知道,但现在可以直接用扩散模型做 P ( z ∣ x i ) P(z|x_i) P(z∣xi)到标准正态分布之间的映射,生成的时候就可以取标准正态分布中的 ε ε ε,得到一个服从 P ( z ∣ x i ) P(z|x_i) P(z∣xi)分布的 z z z,有了这个 z z z,就可以通过Decode解出 x ^ i \hat x_i x^i
问题二(VAE自己的问题):
有两张图的分布情况如下图,这两张图都采样了同一个点,这一个点都通过一个decode去学习
z z z重构出来的 x x x在和左边学习还是右边学习的情况时有分歧,所以模型在训练过程中就会不断的取average,所以它学习高频特征的能力很不好。
四、Latent Diffusion Models
扩散模型(Diffusion Models, DMs)通常在像素空间中操作,因此训练这些模型需要消耗大量的GPU资源,而推理过程基于序列化生成也会产生较高的计算成本。为在有限的计算资源下实现高质量的生成效果和灵活性,Latent Diffusion Models (LDM) 提出了将扩散过程转移到预训练自编码器的潜在空间中。相比于基于像素的 DMs,这种方法显著降低了计算需求,并且在空间维度上表现出更好的可扩展性。
此外,LDM 通过引入跨注意力(cross-attention)层,使模型能够更灵活地根据不同类型的条件输入进行图像生成。这一创新方法大幅提升了扩散模型的训练和推理效率,使得高质量图像生成在消费级显卡上成为可能。当前广受欢迎的 DALL-E 2 和 Stable Diffusion 就应用了 LDM 技术,进一步推动了扩散模型在图像生成领域的应用。
1. 方法介绍
本方法的整体结构如下图所示,主要分为三部分:最左侧的红框对应于感知图像压缩,中间的绿框对应 Latent Diffusion Models,右侧的白框表示生成条件,下面将分别介绍这三个部分。
1.1 感知图像压缩
LDM 把图像生成过程从原始的图像像素空间转换到了一个隐空间,具体来说,对于一个维度为 x ∈ R H × W × 3 \mathbf{x}\in\mathbb{R}^{H\times W\times 3} x∈RH×W×3的 RGB 图像,可以使用一个 encoder E \mathcal{E} E 将其转换为隐变量 z = E ( x ) \mathbf{z}=\mathcal{E}(\mathbf{x}) z=E(x),也可以用一个 decoder D \mathcal{D} D 将其从隐变量转换回像素空间 x ~ = D ( E ( x ) ) \tilde{\mathbf{x}}=\mathcal{D}(\mathcal{E}(\mathbf{x})) x~=D(E(x))在转换时会将图像下采样,作者测试了一系列下采样倍数 f ∈ { 1 , 2 , 4 , 8 , 16 , 32 } f\in\{1, 2, 4, 8, 16, 32\} f∈{1,2,4,8,16,32},发现下采样 4-16 倍的时候可以比较好地权衡效率和质量。
在进行图像压缩时,为了防止压缩后的空间是某个高方差的空间,需要进行正则化。作者使用了两种正则化,第一种是 KL-正则化,也就是将隐变量和标准高斯分布使用一个 KL 惩罚项进行正则化;第二种是 VQ-正则化,也就是使用一个 vector quantization 层进行正则化。
1.2 Latent Diffusion Models
实际上 latent diffusion models 和普通的扩散模型没有太大区别,只是因为从像素空间变到了隐空间,所以维度降低了。训练的优化目标也没有太大变化,普通的扩散模型优化目标为:
L D M = E x , ϵ ∼ N ( 0 , 1 ) , t [ ∣ ∣ ϵ − ϵ θ ( x t , t ) ∣ ∣ 2 2 ] L_{DM}=E_{x,ϵ∼N(0,1),t} [∣∣ϵ−ϵ_θ(x_t,t)∣∣_2^2] LDM=Ex,ϵ∼N(0,1),t[∣∣ϵ−ϵθ(xt,t)∣∣22]
而 Latent Diffusion Models 的优化目标只是套了一层 autoencoder:
L L D M = E E ( x ) , ϵ ∼ N ( 0 , 1 ) , t [ ∣ ∣ ϵ − ϵ θ ( x t , t ) ∣ ∣ 2 2 ] L_{LDM}=E_{\mathcal{E}_(x),ϵ∼N(0,1),t} [∣∣ϵ−ϵ_θ(x_t,t)∣∣_2^2] LLDM=EE(x),ϵ∼N(0,1),t[∣∣ϵ−ϵθ(xt,t)∣∣22]
在采样时,首先从隐空间随机采样噪声,在去噪后再用 decoder 转换到像素空间即可。
1.3 条件生成
为了进行条件生成,需要学习 ϵ θ ( x t , t , y ) \epsilon_\theta(\mathbf{x}_t,t,y) ϵθ(xt,t,y),这里使用的方法是在去噪网络中加入 cross attention 层,条件通过交叉注意力注入。
如果说不带控制信号的LDM是对 p ( z ) p(z) p(z)建模的话,那么加上控制信号的建模可以表述成 p ( z ∣ y ) p(z∣y) p(z∣y)。控制信号 y y y可以是文本、图片布局、轮廓图等。当有有一些更为细粒度的生成存在多个控制信号,即对 p ( z ∣ y 1 , ⋯ , y N ) p(z∣y_1,⋯,y_N) p(z∣y1,⋯,yN)。为了使得生成过程考虑控制信息,作者在原有的Unet backbone上引入了一个交叉注意力机制,来融入控制信号。首先通过一个将 y y y投影到为一个中间表征 τ θ ( y ) ∈ R M × d τ τ_θ(y)∈R^{M×dτ} τθ(y)∈RM×dτ,随后和DM某层的输出进行融合。
假定 φ i ( z t ) φ_i(z_t) φi(zt)是隐变量 z z z在Unet在层 i i i时间步为 t t t时的输出,下面需要将 φ i ( z t ) φ_i(z_t) φi(zt)与 τ θ ( y ) τ_θ(y) τθ(y)用交叉注意力机制融合
A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d ) V Q = W Q ( i ) φ i ( z t ) ; K = W K ( i ) τ θ ( y ) ; V = W V ( i ) τ θ ( y ) Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt d}) V \\ Q=W_Q^{(i)} φ_i(z_t); K=W_K^{(i)} τ_θ(y); V=W_V^{(i)} τ_θ(y) Attention(Q,K,V)=softmax(dQKT)VQ=WQ(i)φi(zt);K=WK(i)τθ(y);V=WV(i)τθ(y)
此时模型的优化目标为
L L D M = E E ( x ) , ϵ ∼ N ( 0 , 1 ) , t [ ∣ ∣ ϵ − ϵ θ ( x t , t , τ θ ( y ) ) ∣ ∣ 2 2 ] L_{LDM}=E_{\mathcal{E}_(x),ϵ∼N(0,1),t} [∣∣ϵ−ϵ_θ(x_t,t,τ_θ(y))∣∣_2^2] LLDM=EE(x),ϵ∼N(0,1),t[∣∣ϵ−ϵθ(xt,t,τθ(y))∣∣22]
2. LDM 使用示例
huggingface Diffusers 将各种 Diffusion Model Pipeline 都包装好了,使用 Diffusion model 就和使用 Transformers 一样地方便:
# !pip install diffusers transformers
from diffusers import DiffusionPipelinemodel_id = "CompVis/ldm-text2im-large-256"# load model and scheduler
ldm = DiffusionPipeline.from_pretrained(model_id)# run pipeline in inference (sample random noise and denoise)
prompt = "A painting of a squirrel eating a burger"
images = ldm([prompt], num_inference_steps=50, eta=0.3, guidance_scale=6).images# save images
for idx, image in enumerate(images):image.save(f"squirrel-{idx}.png")
3. LDM Pipeline
LDM 的 pipeline 可以简化表示为:Pipeline(prompt, num_inference_steps, latents)
。我们暂时考虑没有 negative prompt 和 初始 latent 的输入,那么整个采样过程大致可以表示为:
- 首先采用了 BERT 架构模型对 prompt 进行处理,生成
text_hidden_state
;同时生成随机噪声latents
。
text_hidden_state = LDMBERT(prompt) # shape=[bs, len_seq, d_model] = [1, 77, 1280]
latents = randn_tensor(latents_shape)
对于 "CompVis/ldm-text2im-large-256"
,其中使用了 LDMBert
, 参考 huggignface 的 LDMBert 实现,LDMBert
与传统 BERT 架构相似,规模不同,LDMBert 采用 32 层, hidden_size 为 1280,属实比 bert-base 大上不少。同时文本被 padding 到了固定的 77 长度,以此来保证文字的 hidden state 格式为 [batch_size, 77, 1280]
。
- 之后进行传统的扩散模型 backward process:
for t in self.progress_bar(self.scheduler.timesteps):noise_pred = self.unet(latents_input, t, encoder_hidden_states=context).sample# compute the previous noisy sample x_t -> x_t-1latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs).prev_sample
其中 UNET 为 UNet2DConditionModel
,与传统 Unet 不同在于其应用了 Cross Attention 对文字以及图片信息进行综合处理,下文会对改模块做梳理。scheduler 可以选 DDIM 或者其他算法。
- 最后对 latent hidden state 进行 decode,生成图片:
latents = 1 / self.vqvae.config.scaling_factor * latents
image = self.vqvae.decode(latents).sample
4. LDM 中的 UNET
backward process 中的 self.unet(...)
,即 UNET2DCondition(sample, timestep, encoder_hidden_state)
前向推导可以看成五部分,(以下以 CompVis/ldm-text2im-large-256
为例介绍):
- 准备 time steps :Timesteps 编码信息是 diffusion 中 predict noise residual 模型的标配:
# 经过两次映射得到 timesteps 对应的 embedding
t_emb = self.time_proj(timesteps)
emb = self.time_embedding(t_emb, timestep_cond)
- pre-process: LDM 只用了一个 2D 卷积对输入的 hidden state 进行处理
sample = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding)(sample)
- down sampling :down sampling 包括了 3 个 CrossAttnDownBlock2D, 和 1 个
DownBlock2D
。
# down sampling 大致前向推导
down_block_res_samples = (sample,)
for downsample_block in self.down_blocks:sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale)# 用于 UNET 的残差链接down_block_res_samples += res_samples
其中每个 CrossAttnDownBlock2D 大概前向过程为:
# CrossAttnDownBlock2D
def forward(self, hidden_states, temb, encoder_hidden_states=None)output_states = ()for resnet, attn in zip(self.resnets, self.attentions):hidden_states = resnet(hidden_states, temb)hidden_states = attn(hidden_states,encoder_hidden_states=encoder_hidden_states,cross_attention_kwargs=cross_attention_kwargs,).sampleoutput_states += (hidden_states,)# downsampler = Conv2D hidden_states = downsampler(hidden_states)output_states += (hidden_states,)return hidden_states, output_states
在 CompVis/ldm-text2im-large-256
中,每个 CrossAttnDownBlock2D
包含了 2 个 attn(Transformer2DModel)
以及 2 个 resnet (ResnetBlock2D)
。
文字与图像的交互就发生在 Transformer2DModel
当中。每个 Transformer2DModel 先对输入的图像数据进行预处理,将图片格式从如 (batch_size, channel, width, height)
或 (batch_size, num_image_vectors)
转换为 (batch_size, len_seq, hidden_size)
,而后将 hidden_states
传入 1 层传统 Transformer layer(非 bert 或 GPT 类型),先对图像hidden_states
进行 self-attention,而后结合 encoder_hidden_states
进行 cross attention 处理。
- mid processing:
sample = MidBlock2DCrossAttn()(sample, emb,encoder_hidden_states)
在 CompVis/ldm-text2im-large-256
中,upsampling 和 down sampling 之间采用 MidBlock2DCrossAttn 连接,MidBlock2DCrossAttn
包括了 1 个 1 层的 Transformer2DModel
以及 1 个 resnet ResnetBlock2D
。
- upsampling :upsampling 采用的模块 UpBlocks 包括了
("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
,各个模块的架构与 down sampling 中的模块相似。
# upsample_block
for i, upsample_block in enumerate(self.up_blocks):sample = upsample_block(hidden_states=sample,temb=emb,res_hidden_states_tuple=res_samples,upsample_size=upsample_size,scale=lora_scale,)
- post-process
# GroupNorm
sample = self.conv_norm_out(sample)
# Silu
sample = self.conv_act(sample)
# Conv2d(320, 4, kernel=(3,3), s=(1,1), padding=(1,1))
sample = self.conv_out(sample)
总结起来,down sampling,midprocess,upsampling 三个步骤中都涉及到了 Transformer2DModel
,实现多模态的信息交互。
参考资料
DIFFUSION 系列笔记| Latent Diffusion Model
CompVis/ldm-text2im-large-256