【从VAE到LDM】Variational Auto Encoder原理以及关于Latent Diffusion的思考
从 VAE 原理讲起,理解 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。
传统的扩散模型一般是输入一张图像 ,然后对它进行增噪,再通过去噪采样得到一个新的图像。但这个过程的计算都是在Pixel Space上面进行的,采样空间太大,学习的噪声维度和图像的维度是相同的。当进行高分辨率图像生成时,需要的计算资源会急剧增加,虽然 DDIM 等工作已经对此有所改善,但效果依然有限。
因此就产生了Latent Space上进行扩散操作的想法,经过测试证明是可以的。但压缩到一个比较小的程度,比如说压缩到32*32大小,会有信息的损失。作者认为压缩分为两种:感知压缩和语义压缩。
作者认为感知压缩损失的这些信息实际上是很好学到的,因为它有很强的相关性。但是整体语义压缩损失实际上是比较重要的,在Latent Diffusion模型里,实际上是用扩散模型去学习语义。
实际上的过程:
- 先把编码成一个隐藏向量
- 给增噪去噪,学得
- 解码得到图像

注意:
Encoder和Decoder的过程实际上是在做感知方面的压缩和解码,中间的扩散模型处理语义信息Encoder和Decoder就来自于VAE,Variational Auto Encoder其实是从AutoEncoder来的,用了一些方法把自编码器变成了生成模型,这个方法在2013年提出。实际上VAE的结果是不太好的,那为什么把中间换成扩散模型效果会变好?就这说明编码器和解码器的能力是足够的,问题出在中间的部分。
所以学习Latent Diffusion之前必须了解VAE的原理
问题定义
什么是AE(Auto Encoder)
AutoEncoder自编码器,它实际上是一种自监督学习图像特征的一个方法,它的原理很简单:
首先有一个,把它输入到Encoder中编码成一个,再把这个输入到Decoder中,给它解码成,这就是重构图像,然后再去跟真实的做重构损失的计算,最后来优化编码器和解码器。如果训练好给定一个,就可以得到重构的比较好的,就可以认为实际上已经是学到了中比较丰富的细节特征了。
但是AE为什么不能作为生成模型?
假设现在有一个数据集它里边有n个图像,通过Encoder可以给它编码成。因为生成模型有一个核心的诉求:生成一个全新的图像,即数据集里没有出现过的。现在如果想要一个全新的,但肯定不能从中来,但是也不能随便来。如果知道服从哪一种分布,就可以从这个分布中取到一个新的,就可以通过解码得到一张新的图片。但训练时没有对的分布进行约束,导致的分布可能是不规律的,就不能生成新的图像。
假设服从正态分布 如果假设服从正态分布的话,很难确定正态分布N(0,1)中每个点和真实数据中每个点的对应关系,即和之间的对应关系
假设已经知道后验概率 此时不再去找确切的了,而是假设已经知道后验概率,如果知道了这个分布,就可以从中取得一个新的,然后再输入到Decoder中,输出的就可以跟进行重构损失的计算来优化Decoder。
思路的转变 这时思路就从想优化的auto encoder变成一个生成模型变成了需要找到这么一个后验概率,接下来就可以通过变分贝叶斯的方法找到这个分布。
一、VAE
1. 定义
变分贝叶斯:如下图,有一个分布集合,其中有一点,其中有一个已知分布族,这个分布符合高斯分布,现在希望从这个分布族里找到一个最接近的分布并做近似。如果最开始取的是点的分布,然后通过不断的优化,慢慢找到这个地方。这个分布就是离的距离是最近的,所以最终通过优化找到这个分布。
从上面的高斯分布族取出一个分布,此时需要优化 φ ,就需要找到已知分布中离已知 P 最近的一个分布。自然就想到一个度量方法,那就是KL散度。现在的问题就变成了最小化。如果以KL散度作为这个距离的衡量标准的话(虽然它不是一个严格意义的距离,它不服从对称性),这个方法就叫做变分贝叶斯法。
2. 变分下界
想要最小化,但是这里面有一个难点就是这个我们不知道,所以就很难去通过计算方法求得它,另一种方法就是变分下界方法。
根据KL散度定义展开:
- 第三行到第四行根据数学期望展开,后一项的展开为
- 第四行的最后一项,是服从的数学期望,但里没有,所以这项实际上就是
- 第五行的被称为变分下界即 ,被称为边界似然
- 第七行通过移项得
==我们的目标是最小化散度,因为在上面的推导过程中是随机变量,而中不含,说明在这个里面它是一个定值,所以在中最小化KL散度的任务就转换成了最大化==
3. 最大化 ELBO
- 第二行将展开为;
- 第三行再展开;
- 第四行移项并合并得到第五行;
- 第五行中的第一项就是前面的散度;后一项,为给定之后,的最大似然,在生成模型中最大似然用模型来学。
推理步骤:
此时最大化 的任务就变成了最小化,同时最大化
是从高斯分布中取的随便一个分布
当时我们希望和真实的后验概率越近越好
通过推导发现,需要和先验概率越近越好,而最开始我们已经假设了,这就是为什么我们希望逼近标准正态分布
4. 思路整理
- 先希望得到,目的是得到与对应的,这样好优化decode
- 用变分贝叶斯的方法得到,就是离最近的
- 用散度优化这两个之间的距离
- 通过推导转化为
- 又转化为。第一项就是前面的散度;后一项,为给定之后,的最大似然,在生成模型中最大似然用模型来学。所以就是decode;前面有一个期望值,这个期望是服从的。 ==换句话说,如果在每次编码中得到的分布后,要从这个里面重新取一个,通过蒙特拉洛的方法,让每次取得的都能最大的还原,这就是为什么VAE的训练过程是这样的(下一节解释)。==
5. VAE的训练过程
VAE的训练过程如下图:

- 先输入一个到encode
- 得到一个,通过这两个值得到高斯分布
- 从高斯分布中采样一个,再把这个输入到decode中得到,然后和做重构损失计算
- 期间通过散度尽量让和N(0,1)做近似
==注意:==
- 就是前面讲的,对做最佳近似,是先验分布
- 找到了真实的后验概率的最佳近似之后,就可以从中去一个,我们认为这个是从后验分布离拿到的,所以就和是对应的;
- 通过decode后得到的,就可以和真实的做重构损失计算,这一步就是前面推导出的
- 有的论文认为不断的从高斯分布中采样,让它们都去重构,这也是一种数据增强的方法,这就使得decode更具有泛化性。
二、Latent Space
Latent Space是什么样的呢?
这是在MNIST数据集上做了一个训练,图里一共有10种颜色,代表是0到9,做了对比试验。这个就是隐空间的分布,这是一个二维的分布,实际上隐空间的维度可能是很高的,这里只是做了一个投影。
- 只有重构损失,即AE:有区域,无规则,不服从正态分布,所以很难取出一个;
- 只有散度:无区域,有规则,即正态分布;
- VAE:有区域,有规则;
从这个图可知:
- 重构损失会导致最终的特征是否有区域性;
- 散度对应最终的特征是否有规则;
- 虽然上述分布是在隐空间中的,但实际上也是一种特征图,AE就是在学特征,其中的每个向量都是特征向量,而相似的特征必然在一个维度上,将其投影到平面上是比较相近的;
- 图3可以看到这是一个从左上角到右下角连续的过程,最后生成了数字2;
- 使得隐空间符合设好的已知分布,重构损失会学到一定的特征,而学特征就是Decode的过程,如果给定了一个,它能很好的还原一个,那么我们认为已经包含了比较丰富的特征。
- 图2到3加了重构损失后,就是对图2做了聚类的操作,把图像特征更像的聚集到一起,但是并不影响整体的分布。
三、LDM vs VAE
这两个模型到底差在哪?VAE的效果很一般,虽然LDM也用了VAE编码器解码器的结构,但效果好很多。
LDM训练过程简单示意图:
VAE训练过程简单示意图:

问题一(为什么LDM更好):
- 我们认为LDM和VAE的先验概率都是服从
- 在训练VAE时,用了最小化的方法,让数据尽量接近,是后验概率的近似值。这里就会产生一个问题,不一定是真实的高斯分布,即后验概率,所以在做到的近似值时就会产生偏差,生成效果就不是很好;
- Diffusion的训练过程时不断的增噪,最后得到,采样时直接选取一个是可以直接生成图像的,所以扩散模型通过不断的增噪去噪学得了标准正态分布到真实空间的映射关系;
- 扩散模型可以学任意分布到标准正态之间的映射,VAE是假设的,用符合高斯分布的逼近它,因为具体是什么分布是不知道,但现在可以直接用扩散模型做到标准正态分布之间的映射,生成的时候就可以取标准正态分布中的,得到一个服从分布的,有了这个,就可以通过Decode解出
问题二(VAE自己的问题):
有两张图的分布情况如下图,这两张图都采样了同一个点,这一个点都通过一个decode去学习
重构出来的在和左边学习还是右边学习的情况时有分歧,所以模型在训练过程中就会不断的取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 把图像生成过程从原始的图像像素空间转换到了一个隐空间,具体来说,对于一个维度为的 RGB 图像,可以使用一个 encoder 将其转换为隐变量,也可以用一个 decoder 将其从隐变量转换回像素空间在转换时会将图像下采样,作者测试了一系列下采样倍数,发现下采样 4-16 倍的时候可以比较好地权衡效率和质量。

在进行图像压缩时,为了防止压缩后的空间是某个高方差的空间,需要进行正则化。作者使用了两种正则化,第一种是 KL-正则化,也就是将隐变量和标准高斯分布使用一个 KL 惩罚项进行正则化;第二种是 VQ-正则化,也就是使用一个 vector quantization 层进行正则化。
1.2 Latent Diffusion Models
实际上 latent diffusion models 和普通的扩散模型没有太大区别,只是因为从像素空间变到了隐空间,所以维度降低了。训练的优化目标也没有太大变化,普通的扩散模型优化目标为: 而 Latent Diffusion Models 的优化目标只是套了一层 autoencoder: 在采样时,首先从隐空间随机采样噪声,在去噪后再用 decoder 转换到像素空间即可。
1.3 条件生成
为了进行条件生成,需要学习,这里使用的方法是在去噪网络中加入 cross attention 层,条件通过交叉注意力注入。
如果说不带控制信号的LDM是对建模的话,那么加上控制信号的建模可以表述成。控制信号可以是文本、图片布局、轮廓图等。当有有一些更为细粒度的生成存在多个控制信号,即对。为了使得生成过程考虑控制信息,作者在原有的Unet backbone上引入了一个交叉注意力机制,来融入控制信号。首先通过一个将投影到为一个中间表征,随后和DM某层的输出进行融合。
假定是隐变量在Unet在层时间步为时的输出,下面需要将与用交叉注意力机制融合
此时模型的优化目标为
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 的输入,那么整个采样过程大致可以表示为:
- 首先采用了 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-1
latents = 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,
).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