算法可视化与交互学习平台

学习 Latent Diffusion:在 latent space 中做 DDPM 图像生成Latent Diffusion: DDPM Image Generation in Latent Space

把 No.4 的 VAE latent space 和 No.5 的 DDPM 加噪去噪接起来,理解 Latent Diffusion 如何先把图像压缩成 latent,再在 latent 空间中训练噪声预测模型,最后通过 Decoder 把去噪后的 latent 还原成图像。

Deep LearningIntermediateFree
Kernel
1

先抓住 Latent Diffusion 的核心直觉

Latent Diffusion 的一句话版本:

不要直接在像素图像上做 DDPM,而是在 VAE 压缩后的 latent space 里做 DDPM。

No.4 VAE 已经说明:Encoder 可以把图像 x 压缩成 latent z,Decoder 可以把 z 翻译回图像。No.5 DDPM 已经说明:模型可以学习在某个噪声等级 t 下预测混入的噪声 epsilon

Latent Diffusion 把这两件事合起来:

图像 x -> VAE Encoder -> latent z -> DDPM 在 z 空间加噪 / 去噪 -> VAE Decoder -> 生成图像 x_hat

所以 Latent Diffusion 不是“另一套完全不同的扩散算法”。它仍然使用 DDPM 的核心训练目标,只是把扩散发生的位置从 pixel space 搬到了 latent space。

2

把 VAE 和 DDPM 接起来

Pixel Diffusion 直接在图像像素上加噪和去噪。Latent Diffusion 先用 VAE 把图像压缩成 latent,再在 latent 上做同样的 DDPM 训练。

No.4 VAE: image x -> Encoder -> latent z -> Decoder -> image x_hat No.5 DDPM: x_0 -> x_t, then learn epsilon_theta(x_t, t) No.10 Latent Diffusion: image x_0 -> Encoder -> z_0 z_0 -> z_t, then learn epsilon_theta(z_t, t) z_0_hat -> Decoder -> image x_hat

最重要的变化只有一个:DDPM 的变量从 x 变成了 z。原来模型预测的是像素空间里的噪声,现在预测的是 latent 空间里的噪声。

这也是 Stable Diffusion 这类模型的关键工程思想:图像很大,但 latent 更小。在 latent 上做几十步或上千步去噪,计算量会小很多。

3

Pixel Diffusion vs Latent Diffusion

为了看清楚 Latent Diffusion 为什么有用,先把两条路径并排比较。

Pixel Diffusion: 在 512 x 512 x 3 的像素张量上做加噪、预测噪声、去噪。 Latent Diffusion: 先压缩到更小的 latent 张量,例如 64 x 64 x 4, 再在 latent 张量上做加噪、预测噪声、去噪。

如果把每个数字都看作一个需要处理的维度,像素空间大约有 512*512*3 = 786432 个数;而 64*64*4 = 16384 个 latent 数字只有前者的约 1/48。这就是为什么同样是 DDPM,在 latent space 中运行会更省计算。

Pixel DDPM 关心:每个像素/通道的噪声是什么。 Latent DDPM 关心:每个 latent 位置/通道的噪声是什么。 VAE Decoder 负责把去噪后的 latent 翻译回图像。 Denoiser 负责在 latent space 里逐步恢复结构。
4

第一步:用 VAE 把图像变成 latent

Latent Diffusion 的前提是已经有一个可用的 VAE。Encoder 把图像压缩为 latent,Decoder 把 latent 还原为图像。

原始干净图像
Clean image
VAE Encoder 得到的干净 latent
Clean latent produced by the VAE encoder
VAE Encoder
VAE encoder
VAE Decoder
VAE decoder

VAE 在这里负责什么

VAE 负责建立图像和 latent 之间的翻译关系。Encoder 把图像压缩成 latent,Decoder 把 latent 还原成图像。Latent Diffusion 后续的扩散过程发生在 上,而不是直接发生在 上。

最小代码

z0 = vae_encoder(x0) x_recon = vae_decoder(z0)
5

第二步:在 latent 上做 DDPM forward 加噪

Latent Diffusion 把 DDPM 的 forward 公式从像素变量 x 换成 latent 变量 z。

干净 latent
Clean latent
第 t 个噪声等级下的 latent
Noisy latent at timestep t
和 z0 同形状的标准高斯噪声张量
Standard Gaussian noise tensor with the same shape as z0
累计信号保留比例
Cumulative signal retention

和 No.5 的 DDPM 公式有什么不同

形式完全一样,只是变量从 换成了 。这意味着噪声不是加到原始图像像素上,而是加到 VAE 压缩后的 latent 表示上。

为什么这仍然是 DDPM

DDPM 的核心是定义 forward 加噪分布,并训练模型预测加入的噪声。Latent Diffusion 没有改变这个核心,只是改变了数据所在的空间。

最小代码

noise = torch.randn_like(z0) a = alpha_bar[t] zt = torch.sqrt(a) * z0 + torch.sqrt(1 - a) * noise
6

逐步演算:一个 latent 元素如何被加噪

准备一个干净 latent 元素和一份噪声
z0 = 0.75
epsilon = -1.2
alpha_bar_t = 0.49
Initial Variables
z0
0.75
epsilon
-1.2
alpha_bar_t
0.49
Step 1 Variables
z0
0.75
epsilon
-1.2
alpha_bar_t
0.49
Step 1 / 4
7

训练目标:预测 latent noise

Denoiser 输入带噪 latent z_t 和 timestep t,输出它预测的 latent 噪声。

forward 加噪时真实加入的 latent 噪声
True latent noise added in the forward process
神经网络预测的 latent 噪声
Predicted latent noise
带噪 latent
Noisy latent

训练数据从哪里来

训练时先把真实图像 编码成 ,再随机抽 timestep 和噪声,合成 。因为噪声是我们自己加进去的,所以真实 可以直接作为监督信号。

最小代码

z0 = vae_encoder(x0) noise = torch.randn_like(z0) t = torch.randint(0, T, (batch_size,)) zt = sqrt_alpha_bar[t] * z0 + sqrt_one_minus_alpha_bar[t] * noise noise_pred = denoiser(zt, t) loss = ((noise_pred - noise) ** 2).mean()
8

生成路径:从 latent noise 到图像

生成时先在 latent space 中采样纯噪声,逐步去噪得到 z0,再用 VAE Decoder 生成图像。

latent space 中的纯噪声起点
Pure noise starting point in latent space
去噪后得到的干净 latent
Denoised clean latent
Decoder 生成的图像
Generated image from the decoder

生成阶段没有原图

训练时从真实图像得到 ;生成时没有真实图像,只有从标准高斯采样出来的 。Denoiser 的任务就是一步步把 推回更像训练数据 latent 的区域。

最小代码

z = torch.randn(latent_shape) for t in reversed(range(T)): noise_pred = denoiser(z, t) z = scheduler_step(z, noise_pred, t) image = vae_decoder(z)
9

文字条件如何进入 Latent Diffusion

先理解无条件 Latent Diffusion,再理解文字控制会更稳。无条件版本只学习 ;文字条件版本把文本 embedding 也交给 denoiser:

这里 是文本编码器输出的条件向量。模型在去噪时不只是问“这里有哪些噪声”,还会问“如果目标是这段文字,哪些结构应该被保留,哪些结构应该被去掉”。

无条件 Latent Diffusion: denoiser(z_t, t) -> predicted noise 文字条件 Latent Diffusion: text -> text encoder -> c_text denoiser(z_t, t, c_text) -> predicted noise guided by text

Stable Diffusion 可以看作大型文字条件 Latent Diffusion:VAE 负责图像与 latent 的转换,文本编码器负责把 prompt 变成条件,U-Net 或类似 denoiser 负责在 latent 中逐步去噪。

10

Toy Latent Diffusion 实验:先压缩,再在 latent 里去噪

这个实验仍然是 Toy,因为图像只有 16×16、网络很小、数据也很简单;但它不再用“预测噪声质量”做模拟。每次运行都会真实训练两个模型: 1. TinyVAE:把图像 x 编码成 latent z,再从 z 解码回图像。 2. Latent DDPM denoiser:在 latent space 中按 DDPM 公式加噪,训练网络预测真实加入的 epsilon。 训练完成后,实验会从纯高斯 latent 噪声 z_T 开始,执行 DDIM 风格的反向采样,得到 z_0_hat,再用 VAE Decoder 还原成图像。默认参数下,生成结果通常能恢复出目标图像 7 到 8 分以上的主体结构;如果训练步数太少、latent 分辨率太低,分数会下降。训练过程中会实时推送 loss 曲线和当前 step,先显示训练进度,完成后再显示完整生成 storyboard。 为了让“从纯噪声恢复原图结构”看得清楚,训练集只包含当前选择图案的中心样本与轻微扰动版本。模型不是偷看最后的原图,而是在这个很小的数据分布上学习如何从噪声回到该类图案的 latent。

Parameter Panel
7 Params
11

拆解 Toy Latent Diffusion:公式如何落到代码

卡片 10 里的实验是一个很小的 Latent Diffusion,但它的算法链条是真实的。它不是直接把噪声图调亮调暗,也不是用某个质量滑块伪造去噪效果,而是按下面的路径实际训练和采样:

1. 生成简单图像数据 x_data 2. 训练 TinyVAE: image x -> latent z -> reconstruction x_recon 3. 固定 VAE Encoder, 得到训练 latent z_data 4. 在 latent 上执行 DDPM forward: z0 -> zt 5. 训练 denoiser 预测真实噪声 epsilon 6. 从纯 latent 噪声开始 reverse sampling: z_T -> z_0_hat 7. 用 VAE Decoder 把 z_0_hat 还原为图像

理解这张实验卡时,最重要的一句话是:DDPM 仍然做噪声预测,只是变量从像素 x 换成了 latent z。

第一步:TinyVAE 学会压缩图像

真实 Latent Diffusion 通常先有一个训练好的 VAE。Toy 实验为了自包含,会先训练一个很小的 TinyVAE。Encoder 输出 latent 的均值 mu,Decoder 学会从 latent 还原图像。

recon, mu, logvar = _vae(batch)
foreground_weight = 1.0 + 8.0 * batch
recon_loss = ((recon - batch).pow(2) * foreground_weight).mean()
kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
loss = recon_loss + 0.0001 * kl_loss

这里的 foreground_weight 是教学实验里很关键的一点:它是前景像素的重建损失权重。普通 MSE 会把每个像素一视同仁:

recon_loss = ((recon - batch) ** 2).mean()

但 16×16 小图像里背景像素通常远多于前景像素。以黑底图案为例,如果大部分像素都是 0,模型即使把图像重建得偏黑,也可能拿到不算太差的平均误差;这会让 TinyVAE 误以为“全黑也不错”,从而丢掉眼睛、嘴巴、边界这些真正有信息量的结构。

所以实验代码把每个像素的误差乘上一个由原图亮度决定的权重:

foreground_weight = 1.0 + 8.0 * batch
recon_loss = ((recon - batch).pow(2) * foreground_weight).mean()
batch = 0.00 的黑色背景 -> weight = 1.0 batch = 1.00 的白色前景 -> weight = 9.0 batch = 0.45 的灰色脸部/结构 -> weight = 4.6

也就是说,前景像素错了会被罚得更重,背景像素仍然会被学习,但不会在 loss 里淹没主体。这样 TinyVAE 的 Decoder 才更容易保住图案轮廓,后面的 latent DDPM 也才有一个有意义的 latent 空间可以学习。

foreground_weight 不是 Latent Diffusion 的核心公式,而是这个 Toy 实验里的训练技巧。真正的 Latent Diffusion 核心仍然是:VAE 负责 x ↔ z 的压缩与解码,DDPM 在 z 空间里学习噪声预测。

第二步:把图像数据变成 latent 数据

VAE 训练好后,实验不再直接对像素做 DDPM,而是把所有训练图像编码成 latent,再做一次标准化。标准化后的 _z_data 更接近扩散模型习惯处理的零均值、单位尺度空间。

_all_mu, _ = _vae.encode(_x_data)
_z_mean = _all_mu.mean(dim=(0, 2, 3), keepdim=True)
_z_std = torch.clamp(_all_mu.std(dim=(0, 2, 3), keepdim=True), min=0.08)
_z_data = (_all_mu - _z_mean) / _z_std

所以实验里的 latent 不是普通平均池化结果,而是由训练后的 Encoder 产生的表示。latent_res 控制 latent 的空间分辨率,latent_channels = 4 控制每个位置有多少个 latent 通道。

第三步:在 latent 上做 DDPM forward 加噪

这一步和 No.5 的 DDPM 公式完全同构。差别只有一个:No.5 用的是像素图像 x0,这里用的是 latent clean

clean = _z_data[idx]
t = torch.randint(0, T, (32,))
noise = torch.randn_like(clean)
a = _alpha_bar[t].view(-1, 1, 1, 1)
zt = torch.sqrt(a) * clean + torch.sqrt(1 - a) * noise

noise 就是训练监督信号里的真实 epsilon。因为噪声是我们自己加进去的,所以训练时模型知道正确答案是什么。

第四步:denoiser 学会预测 latent noise

Denoiser 输入带噪 latent zt 和 timestep t,输出它认为混进去的噪声 pred_noise。训练目标不是直接预测图像,也不是直接预测干净 latent,而是预测那份真实加入的 epsilon。

pred_noise = _denoiser(zt, t)
loss = F.mse_loss(pred_noise, noise)

实验里的 _LatentDenoiser 是一个很小的卷积网络。输入通道包括三类信息:带噪 latent、归一化 timestep、latent 网格坐标。坐标通道让小网络知道“当前位置在 latent 网格的哪里”,对 2×2、4×4、8×8 的小 latent 尤其有帮助。

第五步:从纯 latent 噪声开始生成

生成时没有原图,也没有 z0。实验从一份纯高斯 latent 噪声开始:

_pure_start = torch.randn_like(_z_target)
_sample = _pure_start.clone()

然后每一步让 denoiser 预测当前 latent 里的噪声,再用 DDIM 风格的更新公式往更干净的 timestep 走。

eps_pred = _denoiser(_sample, t_batch)
z0_hat = (_sample - torch.sqrt(1 - a_t) * eps_pred) / torch.sqrt(a_t)
_sample = torch.sqrt(a_prev) * z0_hat + torch.sqrt(1 - a_prev) * eps_pred

这就是“从噪声里生成图像”的核心机制:模型每一步都先判断哪些成分像噪声,再把当前 latent 推向更像训练数据 latent 的区域。Toy 实验的训练集只包含当前图案及轻微扰动,所以从纯噪声回来的结果会倾向于这个图案,而不是任意图像。

第六步:Decoder 把去噪后的 latent 翻译回图像

因为 DDPM 发生在标准化后的 latent 空间,最后解码前要先把 latent 还原回 VAE 原来的尺度,再交给 Decoder。

_generated = _vae.decode(_sample * _z_std + _z_mean)

因此 Plot Output 里的 final generated 不是直接由 denoiser 画出来的。Denoiser 只负责把纯噪声 latent 推回一个合理的 z0;真正把 latent 翻译成可见图像的是 VAE Decoder。

把公式和代码变量对齐

x0 -> _target 或 batch 中的 16×16 图像 E_phi(x0) -> _vae.encode(...) z0 -> _target_mu / _all_mu / clean z_data -> 标准化后的训练 latent alpha_bar_t -> _alpha_bar[t] epsilon -> noise = torch.randn_like(clean) zt -> torch.sqrt(a) * clean + torch.sqrt(1 - a) * noise epsilon_theta -> _denoiser(zt, t) loss -> F.mse_loss(pred_noise, noise) z_T -> _pure_start = torch.randn_like(_z_target) z0_hat -> (_sample - sqrt(1 - a_t) * eps_pred) / sqrt(a_t) x_hat -> _vae.decode(_sample * _z_std + _z_mean)

看懂这张对齐表,就能把 Latent Diffusion 的抽象公式读回代码:VAE 决定图像和 latent 如何互译,DDPM 决定 latent 如何被加噪和去噪,denoiser 学的是每个 timestep 下的 latent noise。

AI
问问 LLM:把 Latent Diffusion 讲回 VAE + DDPM