返回列表
扩散模型 2024年11月7日 · 68 分钟

【从VAE到LDM】Variational Auto Encoder原理以及关于Latent Diffusion的思考

从 VAE 原理讲起,理解 Latent Diffusion 为何在潜空间做扩散,从而大幅降低算力开销

论文链接:High-Resolution Image Synthesis with Latent Diffusion Models 官方实现:CompVis/latent-diffusionCompVis/stable-diffusion 视频讲解:一个视频看懂VAE的原理以及关于latent diffusion的思考

前言

目前的扩散模型范式基本上都变成做Latent Diffusion,包括像DIT(diffusion Transformer),实际上也是基于Latent Diffusion的思想,它只是把里面的UNet换成了Transformer。

传统的扩散模型一般是输入一张图像 xx,然后对它进行增噪,再通过去噪采样得到一个新的图像x^\hat x。但这个过程的计算都是在Pixel Space上面进行的,采样空间太大,学习的噪声维度和图像的维度是相同的。当进行高分辨率图像生成时,需要的计算资源会急剧增加,虽然 DDIM 等工作已经对此有所改善,但效果依然有限。

因此就产生了Latent Space上进行扩散操作的想法,经过测试证明是可以的。但压缩到一个比较小的程度,比如说压缩到32*32大小,会有信息的损失。作者认为压缩分为两种:感知压缩和语义压缩。

作者认为感知压缩损失的这些信息实际上是很好学到的,因为它有很强的相关性。但是整体语义压缩损失实际上是比较重要的,在Latent Diffusion模型里,实际上是用扩散模型去学习语义。

实际上的过程:

  • 先把xx编码成一个隐藏向量zz
  • zz增噪去噪,学得z^\hat z
  • 解码得到图像x^\hat x 在这里插入图片描述

注意: Encoder和Decoder的过程实际上是在做感知方面的压缩和解码,中间的扩散模型处理语义信息Encoder和Decoder就来自于VAE,Variational Auto Encoder其实是从AutoEncoder来的,用了一些方法把自编码器变成了生成模型,这个方法在2013年提出。实际上VAE的结果是不太好的,那为什么把中间换成扩散模型效果会变好?就这说明编码器和解码器的能力是足够的,问题出在中间的部分。

所以学习Latent Diffusion之前必须了解VAE的原理

问题定义

什么是AE(Auto Encoder)

AutoEncoder自编码器,它实际上是一种自监督学习图像特征的一个方法,它的原理很简单: 在这里插入图片描述 首先有一个xix_i,把它输入到Encoder中编码成一个ziz_i,再把这个ziz_i输入到Decoder中,给它解码成x^i\hat x_i,这就是重构图像,然后再去跟真实的xix_i做重构损失的计算,最后来优化编码器和解码器。如果训练好给定一个ziz_i,就可以得到重构的比较好的xix_i,就可以认为ziz_i实际上已经是学到了xix_i中比较丰富的细节特征了。

但是AE为什么不能作为生成模型?

假设现在有一个数据集它里边有n个图像[x1,x2....xn]{[x_1,x_2....x_n]},通过Encoder可以给它编码成[z1,z2....zn]{[z_1,z_2....z_n]}。因为生成模型有一个核心的诉求:生成一个全新的图像,即数据集里没有出现过的。现在如果想要一个全新的ZZ,但肯定不能从[z1,z2....zn]{[z_1,z_2....z_n]}中来,但是也不能随便来。如果知道ZZ服从哪一种分布,就可以从这个分布中取到一个新的zz,就可以通过解码得到一张新的图片。但训练时没有对zz的分布进行约束,导致zz的分布可能是不规律的,就不能生成新的图像。

假设zz服从正态分布P(z)N(0,1)P(z)∽N(0,1) 如果假设zz服从正态分布P(z)N(0,1)P(z)∽N(0,1)的话,很难确定正态分布N(0,1)中每个点和真实数据中每个点的对应关系,即xix_iziz_i之间的对应关系

假设已经知道后验概率Pθ(zxi)P_θ(z|x_i) 此时不再去找确切的zz了,而是假设已经知道后验概率Pθ(zxi)P_θ(z|x_i),如果知道了这个分布,就可以从中取得一个新的ziz_i,然后ziz_i再输入到Decoder中,输出的x^i\hat x_i就可以跟xix_i进行重构损失的计算来优化Decoder。

思路的转变 这时思路就从想优化的auto encoder变成一个生成模型变成了需要找到Pθ(zxi)P_θ(z|x_i)这么一个后验概率,接下来就可以通过变分贝叶斯的方法找到这个分布。

一、VAE

1. 定义

变分贝叶斯:如下图,有一个分布集合,其中有一点Pθ(zxi)P_θ(z|x_i),其中有一个已知分布族,这个分布符合高斯分布,现在希望从这个分布族里找到一个最接近PP的分布并做近似。如果最开始取的是x1x_1点的分布,然后通过不断的优化,慢慢找到x2x_2这个地方。这个分布就是离PP的距离是最近的,所以最终通过优化找到这个分布。 在这里插入图片描述 从上面的高斯分布族取出一个分布qφ(zxi)q_φ(z|x_i),此时需要优化 φ ,就需要找到已知分布中离已知 P 最近的一个分布。自然就想到一个度量方法,那就是KL散度。现在的问题就变成了最小化KL(qφ(zxi)Pθ(zxi))KL(q_φ(z|x_i) || P_θ(z|x_i))。如果以KL散度作为这个距离的衡量标准的话(虽然它不是一个严格意义的距离,它不服从对称性),这个方法就叫做变分贝叶斯法。

2. 变分下界

想要最小化KL(qφ(zxi)Pθ(zxi))KL(q_φ(z|x_i) || P_θ(z|x_i)),但是这里面有一个难点就是这个我们不知道Pθ(zxi)P_θ(z|x_i),所以就很难去通过计算方法求得它,另一种方法就是变分下界方法。

根据KL散度定义展开:

=qφ(zxi)logqφ(zxi)Pθ(zxi)dz=qφ(zxi)log(qφ(zxi))dzqφ(zxi)log(Pθ(zxi))dz=Eqφ(zxi)[log(qφ(zxi))]Eqφ(zxi)[log(Pθ(z,xi))]+Eqφ(zxi)[log(Pθ(xi))]=Eqφ(zxi)[log(qφ(zxi))]Eqφ(zxi)[log(Pθ(z,xi))]+log(Pθ(xi))=ELBO+log(Pθ(xi))log(Pθ(xi))=KL+ELBO = \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
  • 第三行到第四行根据数学期望展开,后一项的Pθ(zxi)P_θ(z|x_i)展开为Pθ(z,xi)/Pθ(xi)P_θ(z,x_i)/P_θ(x_i)
  • 第四行的最后一项,是服从qφ(zxi)q_φ(z|x_i)的数学期望,但Eqφ(zxi)[log(Pθ(xi))]E q_φ(z|x_i) [log (P_θ(x_i))]里没有zz,所以这项实际上就是log(Pθ(xi))log (P_θ(x_i))
  • 第五行的Eqφ(zxi)[log(qφ(zxi))]Eqφ(zxi)[log(Pθ(z,xi))]E q_φ(z|x_i) [log (q_φ(z|x_i))] - E q_φ(z|x_i) [log (P_θ(z,x_i))]被称为变分下界即 ELBO-ELBOlog(Pθ(xi))log (P_θ(x_i))被称为边界似然
  • 第七行通过移项得

==我们的目标是最小化KLKL散度,因为在上面的推导过程中zz是随机变量,而log(Pθ(xi))log (P_θ(x_i))中不含zz,说明在这个里面它是一个定值,所以在log(Pθ(xi))=KL+ELBOlog (P_θ(x_i)) = KL + ELBO中最小化KL散度的任务就转换成了最大化ELBOELBO==

3. 最大化 ELBO

=Eqφ(zxi)[logPθ(z)Pθ(xiz)]Eqφ(zxi)[log(qφ(zxi))]=Eqφ(zxi)[logPθ(z)]+Eqφ(zxi)[logPθ(xiz)]Eqφ(zxi)[log(qφ(zxi))]=(Eqφ(zxi)[log(qφ(zxi))]Eqφ(zxi)[logPθ(z)])+Eqφ(zxi)[logPθ(xiz)]=Eqφ(zxi)[log(qφ(zxi))logPθ(z)]+Eqφ(zxi)[logPθ(xiz)]=KL(qφ(zxi)Pθ(z))+Eqφ(zxi)[logPθ(xiz)] = 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)]
  • 第二行将log(Pθ(z,xi))log (P_θ(z,x_i))展开为Pθ(z)Pθ(xiz)P_θ(z)P_θ(x_i|z)
  • 第三行再loglog展开;
  • 第四行移项并合并得到第五行;
  • 第五行中的第一项就是前面的KLKL散度;后一项logPθ(xiz)logP_θ(x_i|z),为给定zz之后,xix_i的最大似然,在生成模型中最大似然用模型来学。

推理步骤: 此时最大化 ELBOELBO的任务就变成了最小化KL(qφ(zxi)Pθ(z))KL(q_φ(z|x_i) || P_θ(z)),同时最大化Eqφ(zxi)[Pθ(xiz)]E q_φ(z|x_i)[P_θ(x_i|z)]

qφ(zxi)N(μ,σ2)q_φ(z|x_i)∽N(\mu,\sigma^2)是从高斯分布中取的随便一个分布

当时我们希望qφ(zxi)N(μ,σ2)q_φ(z|x_i)∽N(\mu,\sigma^2)和真实的后验概率越近越好

通过推导发现,qφ(zxi)N(μ,σ2)q_φ(z|x_i)∽N(\mu,\sigma^2)需要和先验概率Pθ(z)P_θ(z)越近越好,而最开始我们已经假设了Pθ(z)N(0,1)P_θ(z)∽N(0,1),这就是为什么我们希望qφ(zxi)N(μ,σ2)q_φ(z|x_i)∽N(\mu,\sigma^2)逼近标准正态分布

4. 思路整理

  1. 先希望得到Pθ(zxi)P_θ(z|x_i),目的是得到与xix_i对应的zz,这样好优化decode
  2. 用变分贝叶斯的方法得到qφ(zxi)q_φ(z|x_i)qφ(zxi)q_φ(z|x_i)就是离Pθ(zxi)P_θ(z|x_i)最近的
  3. KLKL散度优化这两个之间的距离
  4. 通过推导转化为ELBOELBO
  5. ELBOELBO又转化为KL(qφ(zxi)Pθ(z))+Eqφ(zxi)[Pθ(xiz)]-KL(q_φ(z|x_i) || P_θ(z)) + E q_φ(z|x_i)[P_θ(x_i|z)]。第一项就是前面的KLKL散度;后一项logPθ(xiz)logP_θ(x_i|z),为给定zz之后,xix_i的最大似然,在生成模型中最大似然用模型来学。所以Pθ(xiz)P_θ(x_i|z)就是decode;logPθ(xiz)logP_θ(x_i|z)前面有一个期望值EE,这个期望是服从qφ(zxi)q_φ(z|x_i)的。 ==换句话说,如果在每次编码中得到的qφ(zxi)N(μ,σ2)q_φ(z|x_i)∽N(\mu,\sigma^2)分布后,要从这个里面重新取一个zz,通过蒙特拉洛的方法,让每次取得的zz都能最大的还原xix_i,这就是为什么VAE的训练过程是这样的(下一节解释)。==

5. VAE的训练过程

VAE的训练过程如下图: 在这里插入图片描述

  • 先输入一个xx到encode
  • 得到一个μ,σ\mu,\sigma,通过这两个值得到高斯分布N(μ,σ2)N(\mu,\sigma^2)
  • 从高斯分布中采样一个zz,再把这个zz输入到decode中得到x^\hat x,然后和xx做重构损失计算
  • 期间通过KLKL散度尽量让N(μ,σ2)N(\mu,\sigma^2)和N(0,1)做近似KL(N(μ,σ2)N(0,1))KL(N(\mu,\sigma^2) | N(0,1))

==注意:==

  • N(μ,σ2)N(\mu,\sigma^2)就是前面讲的qφ(zxi)q_φ(z|x_i),对Pθ(zxi)P_θ(z|x_i)做最佳近似,N(0,1)N(0,1)是先验分布Pθ(z)P_θ(z)
  • 找到了真实的后验概率Pθ(zxi)P_θ(z|x_i)的最佳近似qφ(zxi)q_φ(z|x_i)之后,就可以从qφ(zxi)q_φ(z|x_i)中去一个zz,我们认为这个zz是从后验分布离拿到的,所以就和xix_i是对应的;
  • zz通过decode后得到的x^\hat x,就可以和真实的xx做重构损失计算,这一步就是前面推导出的Eqφ(zxi)[logPθ(xiz)]E q_φ(z|x_i)[logP_θ(x_i|z)]
  • 有的论文认为不断的从高斯分布中采样z1,z2z_1,z_2,让它们都去重构xix_i,这也是一种数据增强的方法,这就使得decode更具有泛化性。

二、Latent Space

Latent Space是什么样的呢? 在这里插入图片描述 这是在MNIST数据集上做了一个训练,图里一共有10种颜色,代表是0到9,做了对比试验。这个就是隐空间的分布,这是一个二维的分布,实际上隐空间的维度可能是很高的,这里只是做了一个投影。

  1. 只有重构损失,即AE:有区域,无规则,不服从正态分布,所以很难取出一个zz
  2. 只有KLKL散度:无区域,有规则,即正态分布;
  3. VAE:有区域,有规则;

从这个图可知:

  • 重构损失会导致最终的特征是否有区域性;
  • KLKL散度对应最终的特征是否有规则;
  • 虽然上述分布是在隐空间中的,但实际上也是一种特征图,AE就是在学特征,其中的每个向量都是特征向量,而相似的特征必然在一个维度上,将其投影到平面上是比较相近的;
  • 图3可以看到这是一个从左上角到右下角连续的过程,最后生成了数字2;
  • KLKL使得隐空间符合设好的已知分布,重构损失会学到一定的特征,而学特征就是Decode的过程,如果给定了一个zz,它能很好的还原一个xx,那么我们认为zz已经包含了xx比较丰富的特征。
  • 图2到3加了重构损失后,就是对图2做了聚类的操作,把图像特征更像的聚集到一起,但是并不影响整体的分布。

三、LDM vs VAE

这两个模型到底差在哪?VAE的效果很一般,虽然LDM也用了VAE编码器解码器的结构,但效果好很多。

LDM训练过程简单示意图: 在这里插入图片描述 VAE训练过程简单示意图: 在这里插入图片描述

问题一(为什么LDM更好):

  • 我们认为LDM和VAE的先验概率都是服从P(z)N(0,1)P(z)∽N(0,1)
  • 在训练VAE时,用了最小化KL(qP(z))KL(q || P(z))的方法,让数据尽量接近N(0,1)N(0,1)qq是后验概率P(zxi)P(z|x_i)的近似值。这里就会产生一个问题,P(zxi)P(z|x_i)不一定是真实的高斯分布,即后验概率,所以在做qqP(zxi)P(z|x_i)的近似值时就会产生偏差,生成效果就不是很好;
  • Diffusion的训练过程时不断的增噪,最后得到N(0,1)N(0,1),采样时直接选取一个x^i\hat x_i是可以直接生成图像的,所以扩散模型通过不断的增噪去噪学得了标准正态分布N(0,1)N(0,1)到真实空间P(x)P(x)的映射关系;
  • 扩散模型可以学任意分布到标准正态之间的映射,VAE是假设的P(zxi)P(z|x_i),用符合高斯分布的qq逼近它,因为P(zxi)P(z|x_i)具体是什么分布是不知道,但现在可以直接用扩散模型做P(zxi)P(z|x_i)到标准正态分布之间的映射,生成的时候就可以取标准正态分布中的εε,得到一个服从P(zxi)P(z|x_i)分布的zz,有了这个zz,就可以通过Decode解出x^i\hat x_i

问题二(VAE自己的问题): 有两张图的分布情况如下图,这两张图都采样了同一个点,这一个点都通过一个decode去学习 在这里插入图片描述 zz重构出来的xx在和左边学习还是右边学习的情况时有分歧,所以模型在训练过程中就会不断的取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 把图像生成过程从原始的图像像素空间转换到了一个隐空间,具体来说,对于一个维度为xRH×W×3\mathbf{x}\in\mathbb{R}^{H\times W\times 3}的 RGB 图像,可以使用一个 encoderE\mathcal{E} 将其转换为隐变量z=E(x)\mathbf{z}=\mathcal{E}(\mathbf{x}),也可以用一个 decoderD\mathcal{D} 将其从隐变量转换回像素空间x~=D(E(x))\tilde{\mathbf{x}}=\mathcal{D}(\mathcal{E}(\mathbf{x}))在转换时会将图像下采样,作者测试了一系列下采样倍数f{1,2,4,8,16,32}f\in\{1, 2, 4, 8, 16, 32\},发现下采样 4-16 倍的时候可以比较好地权衡效率和质量。 在这里插入图片描述

在进行图像压缩时,为了防止压缩后的空间是某个高方差的空间,需要进行正则化。作者使用了两种正则化,第一种是 KL-正则化,也就是将隐变量和标准高斯分布使用一个 KL 惩罚项进行正则化;第二种是 VQ-正则化,也就是使用一个 vector quantization 层进行正则化。

1.2 Latent Diffusion Models

实际上 latent diffusion models 和普通的扩散模型没有太大区别,只是因为从像素空间变到了隐空间,所以维度降低了。训练的优化目标也没有太大变化,普通的扩散模型优化目标为: LDM=Ex,ϵN(0,1),t[ϵϵθ(xt,t)22]L_{DM}=E_{x,ϵ∼N(0,1),t} [∣∣ϵ−ϵ_θ(x_t,t)∣∣_2^2] 而 Latent Diffusion Models 的优化目标只是套了一层 autoencoder: LLDM=EE(x),ϵN(0,1),t[ϵϵθ(xt,t)22]L_{LDM}=E_{\mathcal{E}_(x),ϵ∼N(0,1),t} [∣∣ϵ−ϵ_θ(x_t,t)∣∣_2^2] 在采样时,首先从隐空间随机采样噪声,在去噪后再用 decoder 转换到像素空间即可。

1.3 条件生成

为了进行条件生成,需要学习ϵθ(xt,t,y)\epsilon_\theta(\mathbf{x}_t,t,y),这里使用的方法是在去噪网络中加入 cross attention 层,条件通过交叉注意力注入。

如果说不带控制信号的LDM是对p(z)p(z)建模的话,那么加上控制信号的建模可以表述成p(zy)p(z∣y)。控制信号yy可以是文本、图片布局、轮廓图等。当有有一些更为细粒度的生成存在多个控制信号,即对p(zy1,,yN)p(z∣y_1,⋯,y_N)。为了使得生成过程考虑控制信息,作者在原有的Unet backbone上引入了一个交叉注意力机制,来融入控制信号。首先通过一个将yy投影到为一个中间表征τθ(y)RM×dττ_θ(y)∈R^{M×dτ},随后和DM某层的输出进行融合。

假定φi(zt)φ_i(z_t)是隐变量zz在Unet在层ii时间步为tt时的输出,下面需要将φi(zt)φ_i(z_t)τθ(y)τ_θ(y)用交叉注意力机制融合

Q=WQ(i)φi(zt);K=WK(i)τθ(y);V=WV(i)τθ(y)Q=W_Q^{(i)} φ_i(z_t); K=W_K^{(i)} τ_θ(y); V=W_V^{(i)} τ_θ(y)

此时模型的优化目标为 LLDM=EE(x),ϵN(0,1),t[ϵϵθ(xt,t,τθ(y))22]L_{LDM}=E_{\mathcal{E}_(x),ϵ∼N(0,1),t} [∣∣ϵ−ϵ_θ(x_t,t,τ_θ(y))∣∣_2^2]

2. LDM 使用示例

huggingface Diffusers 将各种 Diffusion Model Pipeline 都包装好了,使用 Diffusion model 就和使用 Transformers 一样地方便:

# !pip install diffusers transformers
from diffusers import DiffusionPipeline

model_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 的输入,那么整个采样过程大致可以表示为:

  1. 首先采用了 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]

  1. 之后进行传统的扩散模型 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-1
    latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs).prev_sample

其中 UNET 为 UNet2DConditionModel,与传统 Unet 不同在于其应用了 Cross Attention 对文字以及图片信息进行综合处理,下文会对改模块做梳理。scheduler 可以选 DDIM 或者其他算法。

  1. 最后对 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,
        ).sample
        output_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