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

论文链接：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$，然后对它进行增噪，再通过去噪采样得到一个新的图像$\hat{x}$。但这个过程的计算都是在Pixel Space上面进行的，采样空间太大，学习的噪声维度和图像的维度是相同的。当进行高分辨率图像生成时，需要的计算资源会急剧增加，虽然 DDIM 等工作已经对此有所改善，但效果依然有限。

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

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

实际上的过程：

• 先把$x$编码成一个隐藏向量$z$
• 给$z$增噪去噪，学得$\hat{z}$
• 解码得到图像$\hat{x}$
  

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

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

问题定义

什么是AE（Auto Encoder）

AutoEncoder自编码器，它实际上是一种自监督学习图像特征的一个方法，它的原理很简单：

首先有一个$x_i$，把它输入到Encoder中编码成一个$z_i$，再把这个$z_i$输入到Decoder中，给它解码成${\hat{x}}_i$，这就是重构图像，然后再去跟真实的$x_i$做重构损失的计算，最后来优化编码器和解码器。如果训练好给定一个$z_i$，就可以得到重构的比较好的$x_i$，就可以认为$z_i$实际上已经是学到了$x_i$中比较丰富的细节特征了。

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

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

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

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

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

一、VAE

1. 定义

变分贝叶斯：如下图，有一个分布集合，其中有一点$P_{\theta }(z|x_i)$，其中有一个已知分布族，这个分布符合高斯分布，现在希望从这个分布族里找到一个最接近$P$的分布并做近似。如果最开始取的是$x_1$点的分布，然后通过不断的优化，慢慢找到$x_2$这个地方。这个分布就是离$P$的距离是最近的，所以最终通过优化找到这个分布。

从上面的高斯分布族取出一个分布$q_{\phi }(z|x_i)$，此时需要优化 φ ，就需要找到已知分布中离已知 P 最近的一个分布。自然就想到一个度量方法，那就是KL散度。现在的问题就变成了最小化$KL(q_{\phi }(z|x_i)||P_{\theta }(z|x_i))$。如果以KL散度作为这个距离的衡量标准的话（虽然它不是一个严格意义的距离，它不服从对称性），这个方法就叫做变分贝叶斯法。

2. 变分下界

想要最小化$KL(q_{\phi }(z|x_i)||P_{\theta }(z|x_i))$，但是这里面有一个难点就是这个我们不知道$P_{\theta }(z|x_i)$，所以就很难去通过计算方法求得它，另一种方法就是变分下界方法。

根据KL散度定义展开：
$$\begin{gathered} KL(q_{\phi }(z|x_i)||P_{\theta }(z|x_i)) \\ =\int q_{\phi }(z|x_i)log\frac{q_{\phi }(z|x_i)}{P_{\theta }(z|x_i)}dz \\ =\int q_{\phi }(z|x_i)log(q_{\phi }(z|x_i))dz-\int q_{\phi }(z|x_i)log(P_{\theta }(z|x_i))dz \\ =E{q}_{\phi }(z|x_i)[log(q_{\phi }(z|x_i))]-E{q}_{\phi }(z|x_i)[log(P_{\theta }(z,x_i))]+E{q}_{\phi }(z|x_i)[log(P_{\theta }(x_i))] \\ =E{q}_{\phi }(z|x_i)[log(q_{\phi }(z|x_i))]-E{q}_{\phi }(z|x_i)[log(P_{\theta }(z,x_i))]+log(P_{\theta }(x_i)) \\ =-ELBO+log(P_{\theta }(x_i)) \\ ⇨log(P_{\theta }(x_i))=KL+ELBO \end{gathered}$$

• 第三行到第四行根据数学期望展开，后一项的$P_{\theta }(z|x_i)$展开为$P_{\theta }(z,x_i)/P_{\theta }(x_i)$
• 第四行的最后一项，是服从$q_{\phi }(z|x_i)$的数学期望，但$E{q}_{\phi }(z|x_i)[log(P_{\theta }(x_i))]$里没有$z$，所以这项实际上就是$log(P_{\theta }(x_i))$
• 第五行的$E{q}_{\phi }(z|x_i)[log(q_{\phi }(z|x_i))]-E{q}_{\phi }(z|x_i)[log(P_{\theta }(z,x_i))]$被称为变分下界即 $-ELBO$，$log(P_{\theta }(x_i))$被称为边界似然
• 第七行通过移项得

我们的目标是最小化$KL$散度，因为在上面的推导过程中$z$是随机变量，而$log(P_{\theta }(x_i))$中不含$z$，说明在这个里面它是一个定值，所以在$log(P_{\theta }(x_i))=KL+ELBO$中最小化KL散度的任务就转换成了最大化$ELBO$

3. 最大化 ELBO

$$\begin{gathered} ELBO=E{q}_{\phi }(z|x_i)[log(P_{\theta }(z,x_i))]-E{q}_{\phi }(z|x_i)[log(q_{\phi }(z|x_i))] \\ =E{q}_{\phi }(z|x_i)[log{P}_{\theta }(z)P_{\theta }(x_i|z)]-E{q}_{\phi }(z|x_i)[log(q_{\phi }(z|x_i))] \\ =E{q}_{\phi }(z|x_i)[log{P}_{\theta }(z)]+E{q}_{\phi }(z|x_i)[log{P}_{\theta }(x_i|z)]-E{q}_{\phi }(z|x_i)[log(q_{\phi }(z|x_i))] \\ =-(E{q}_{\phi }(z|x_i)[log(q_{\phi }(z|x_i))]-E{q}_{\phi }(z|x_i)[log{P}_{\theta }(z)])+E{q}_{\phi }(z|x_i)[log{P}_{\theta }(x_i|z)] \\ =-E{q}_{\phi }(z|x_i)[\frac{log(q_{\phi }(z|x_i))}{log{P}_{\theta }(z)}]+E{q}_{\phi }(z|x_i)[log{P}_{\theta }(x_i|z)] \\ =-KL(q_{\phi }(z|x_i)||P_{\theta }(z))+E{q}_{\phi }(z|x_i)[log{P}_{\theta }(x_i|z)] \end{gathered}$$

• 第二行将$log(P_{\theta }(z,x_i))$展开为$P_{\theta }(z)P_{\theta }(x_i|z)$；
• 第三行再$log$展开；
• 第四行移项并合并得到第五行；
• 第五行中的第一项就是前面的$KL$散度；后一项$log{P}_{\theta }(x_i|z)$，为给定$z$之后，$x_i$的最大似然，在生成模型中最大似然用模型来学。

推理步骤：
此时最大化 $ELBO$的任务就变成了最小化$KL(q_{\phi }(z|x_i)||P_{\theta }(z))$，同时最大化$E{q}_{\phi }(z|x_i)[P_{\theta }(x_i|z)]$

$q_{\phi }(z|x_i)\backsim N(\mu ,{\sigma }^2)$是从高斯分布中取的随便一个分布

当时我们希望$q_{\phi }(z|x_i)\backsim N(\mu ,{\sigma }^2)$和真实的后验概率越近越好

通过推导发现，$q_{\phi }(z|x_i)\backsim N(\mu ,{\sigma }^2)$需要和先验概率$P_{\theta }(z)$越近越好，而最开始我们已经假设了$P_{\theta }(z)\backsim N(0,1)$，这就是为什么我们希望$q_{\phi }(z|x_i)\backsim N(\mu ,{\sigma }^2)$逼近标准正态分布

4. 思路整理

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

5. VAE的训练过程

VAE的训练过程如下图：

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

注意：

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

二、Latent Space

Latent Space是什么样的呢？

这是在MNIST数据集上做了一个训练，图里一共有10种颜色，代表是0到9，做了对比试验。这个就是隐空间的分布，这是一个二维的分布，实际上隐空间的维度可能是很高的，这里只是做了一个投影。

1. 只有重构损失，即AE：有区域，无规则，不服从正态分布，所以很难取出一个$z$；
2. 只有$KL$散度：无区域，有规则，即正态分布；
3. VAE：有区域，有规则；

从这个图可知：

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

三、LDM vs VAE

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

LDM训练过程简单示意图：

VAE训练过程简单示意图：

问题一（为什么LDM更好）：

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

问题二（VAE自己的问题）：
有两张图的分布情况如下图，这两张图都采样了同一个点，这一个点都通过一个decode去学习

$z$重构出来的$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 把图像生成过程从原始的图像像素空间转换到了一个隐空间，具体来说，对于一个维度为$\mathbf{x}\in {\mathbb{R}}^{H\times W\times 3}$的 RGB 图像，可以使用一个 encoder$\mathcal{E}$ 将其转换为隐变量$\mathbf{z}=\mathcal{E}(\mathbf{x})$，也可以用一个 decoder$\mathcal{D}$ 将其从隐变量转换回像素空间$\tilde{\mathbf{x}}=\mathcal{D}(\mathcal{E}(\mathbf{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_{DM}=E_{x,ϵ\sim N(0,1),t}[\mid \mid ϵ-{ϵ}_{\theta }(x_t,t)\mid {\mid }_2^2]$$
而 Latent Diffusion Models 的优化目标只是套了一层 autoencoder：
$$L_{LDM}=E_{{\mathcal{E}}_{(}x),ϵ\sim N(0,1),t}[\mid \mid ϵ-{ϵ}_{\theta }(x_t,t)\mid {\mid }_2^2]$$
在采样时，首先从隐空间随机采样噪声，在去噪后再用 decoder 转换到像素空间即可。

1.3 条件生成

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

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

假定${\phi }_i(z_t)$是隐变量$z$在Unet在层$i$时间步为$t$时的输出，下面需要将${\phi }_i(z_t)$与${\tau }_{\theta }(y)$用交叉注意力机制融合
$$\begin{gathered} Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d}})V \\ Q=W_Q^{(i)}{\phi }_i(z_t);K=W_K^{(i)}{\tau }_{\theta }(y);V=W_V^{(i)}{\tau }_{\theta }(y) \end{gathered}$$
此时模型的优化目标为
$$L_{LDM}=E_{{\mathcal{E}}_{(}x),ϵ\sim N(0,1),t}[\mid \mid ϵ-{ϵ}_{\theta }(x_t,t,{\tau }_{\theta }(y))\mid {\mid }_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]。

2. 之后进行传统的扩散模型 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 或者其他算法。

3. 最后对 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