前言

学习 MIT S.184 课程,本篇文章记录课程第四讲:潜空间和神经网络架构,记录下个人学习笔记,和大家一起分享交流😄

Website:https://diffusion.csail.mit.edu/

Course Notes:An Introduction to Flow Matching and Diffusion Models

1. Class Overview

今天进入第四讲,我们开始讨论如何真正构建大规模图像生成器。

在前面几讲中,我们已经完成了生成模型的核心理论框架:

Lecture 1 主要介绍了 Flow Models 和 Diffusion Models,也就是从 ODE / SDE 的角度理解生成模型。

Lecture 2 讨论了 Flow Matching,重点回答了如何训练一个从噪声分布流向数据分布的向量场。

Lecture 3 讨论了 Score Matching 和 Guidance,说明了如何用 score 函数理解扩散模型,以及如何通过 prompt 对生成过程进行条件控制。

到这里为止,我们已经知道了生成模型的基本数学结构:概率路径、向量场、score 函数、训练损失和采样算法。

但是,如果我们真的想构建一个大规模图像生成器,还需要解决两个非常实际的问题:

  1. 如何让高维图像数据变得更容易建模;
  2. 如何设计合适的神经网络架构来参数化前面讲到的向量场或 score 函数。

这正是本节课的主题:

Latent spaces + Network architectures

2. Latent Spaces

我们先从 潜空间(latent space) 开始。

潜空间的核心思想是:不要直接在原始像素空间中建模图像,而是先把图像压缩成一个更低维、更容易建模的表示,然后在这个压缩后的空间中训练生成模型。

假设我们要生成下面这样一张图像:

如果直接把图像看作一个向量,那么每个像素、每个颜色通道都会对应向量中的一个维度。例如,一张 RGB 图像有 3 个颜色通道,高度为 600,宽度为 1000,那么它对应的向量维度就是:

3 × 600 × 1000 = 1,800,000 3\times 600\times 1000 = 1{,}800{,}000 3×600×1000=1,800,000

也就是说,一张普通的高分辨率图像就已经是一个 180 万维的向量。

这带来了非常大的挑战。

首先,直接在如此高维的像素空间中训练模型会消耗大量 GPU 显存。图像越大,网络中间特征、梯度和优化状态都会急剧增加。

其次,学习问题本身也会变得非常困难。因为在生成模型中,我们不是只做一次前向推理,而是需要在采样过程中反复调用神经网络。例如,ODE / SDE 采样可能需要几十步甚至上百步,每一步都要在高维空间中计算一次向量场或 score。

第三,原始像素空间存在大量冗余。相邻像素通常高度相关,例如天空区域、背景区域、大片同色区域中,很多像素信息是重复的。如果直接在像素层面建模,模型会把大量计算资源花在这些低层细节上,而不是更重要的语义结构上。

这也是为什么高维图像空间对扩散模型尤其困难。

在监督学习中,图像分类模型通常是从高维图像映射到一个低维标签空间,例如 ImageNet 分类中输出 1000 个类别概率。而在生成模型中,我们需要学习的是一个向量场:

u t θ : R d → R d u_t^\theta:\mathbb{R}^d\to\mathbb{R}^d utθ:RdRd

也就是说,输入是一个 d d d 维图像向量,输出仍然是一个 d d d 维速度向量。如果 d = 1,800,000 d=1{,}800{,}000 d=1,800,000,那么模型每一步都要预测一个 180 万维的输出,并且采样时还要重复调用很多次。

因此,直接在像素空间中做大规模生成建模并不是一个理想选择。这就引出了潜空间的作用:

先把图像压缩到一个更低维的潜表示中,再在潜空间中训练生成模型。

3. Standard Autoencoders

那么,潜空间从哪里来?最常见的方法是使用 自编码器(autoencoder)

一个自编码器由两部分组成:

  • Encoder:把数据从原始数据空间映射到潜空间;
  • Decoder:把潜空间中的表示映射回原始数据空间。

假设原始图像数据位于 x ∈ R d x\in\mathbb{R}^d xRd ,其中 d d d 是原始像素空间维度。我们希望把它压缩到一个更低维的潜变量 z ∈ R k z\in\mathbb{R}^k zRk ,其中 k < d k<d k<d ,那么编码器可以写作:

μ ϕ : R d → R k \mu_\phi:\mathbb{R}^d\to\mathbb{R}^k μϕ:RdRk

它把原始数据 x x x 编码为潜变量 z = μ ϕ ( x ) z=\mu_\phi(x) z=μϕ(x)

解码器可以写作:

μ θ : R k → R d \mu_\theta:\mathbb{R}^k\to\mathbb{R}^d μθ:RkRd

它把潜变量 z z z 解码回数据空间 x ^ = μ θ ( z ) \hat{x}=\mu_\theta(z) x^=μθ(z)

因此,整个自编码器的过程就是:

x   → μ ϕ   z   → μ θ   x ^ x \, \xrightarrow{\mu_\phi} \,z\, \xrightarrow{\mu_\theta} \,\hat{x} xμϕ zμθ x^

我们的目标是让重建结果 x ^ \hat{x} x^ 尽可能接近原始输入 x x x

这对应标准的重建损失:

L R e c o n ( ϕ , θ ) = E x ∼ p d a t a [ ∥ μ θ ( μ ϕ ( x ) ) − x ∥ 2 ] \mathcal{L}_{\mathrm{Recon}}(\phi,\theta) = \mathbb{E}_{x\sim p_{\mathrm{data}}} \left[ \left\| \mu_\theta(\mu_\phi(x))-x \right\|^2 \right] LRecon(ϕ,θ)=Expdata[μθ(μϕ(x))x2]

这个损失的含义非常直接:从数据分布中采样一个图像 x x x ,先用编码器压缩成潜变量,再用解码器还原回图像,然后计算重建图像和原始图像之间的误差,如果这个损失足够小,就说明编码器和解码器能够保留图像中的主要信息。

标准自编码器解决了两个问题:

第一,它确实可以压缩数据。只要潜空间维度 k k k 小于原始空间维度 d d d,编码器就会把图像压缩到一个更低维的表示中。

第二,它可以通过重建损失保证信息不会丢失太多。也就是说,经过编码和解码之后,图像仍然能被较好地恢复。

但这还不够。因为在我们的任务中,自编码器本身不是最终目标。我们真正想做的是先把图像编码到潜空间中,然后在潜空间中训练扩散模型或流模型。因此,我们必须关心另一个问题:

数据被编码到潜空间后,潜变量的分布是什么样子?

定义潜分布:

x ∼ p d a t a , z = μ ϕ ( x ) ⇒ z ∼ p l a t e n t x\sim p_{\mathrm{data}}, \qquad z=\mu_\phi(x) \quad\Rightarrow\quad z\sim p_{\mathrm{latent}} xpdata,z=μϕ(x)zplatent

也就是说,我们从真实数据分布中采样一个图像 x x x ,再通过编码器得到潜变量 z z z 。所有这些 z z z 构成的分布就是潜分布 p l a t e n t p_{\mathrm{latent}} platent

问题在于,标准自编码器的重建损失只关心 “能不能重建”,并不关心潜变量分布是否容易建模。换句话说,标准自编码器可能会学到一个很奇怪、很粗糙、很不规则的潜空间分布。

如果潜分布像图中那样非常不规则,那么虽然编码和解码都能工作,但后续在这个潜空间中训练生成模型会变得很困难。这就是标准自编码器的问题:

它可以压缩和重建,但无法保证潜空间分布是 “好学的”。

这里的 “好学” 不是一个严格数学定义,而是指潜分布应该相对平滑、连续、结构规整,使得后续的扩散模型或流模型可以比较容易地学习它。如果潜分布很破碎、很扭曲,那么我们虽然降低了维度,却可能让生成建模问题变得更难。

因此,我们希望构造一种自编码器,它不仅能重建图像,还能让潜空间分布尽可能接近一个简单分布,例如标准高斯分布。

这就引出了 变分自编码器(Variational Autoencoder, VAE)

4. Variational Autoencoders

在其他机器学习课上,大家可能听说过变分自编码器(VAE),但在这里,我们不会把它当作独立的通用模型来讨论,而是仅从构建潜扩散模型的角度来介绍它。

变分自编码器的核心思想是:把确定性编码变成随机编码

在标准自编码器中,给定一个输入 x x x ,编码器输出一个确定的潜变量:

z = μ ϕ ( x ) z=\mu_\phi(x) z=μϕ(x)

而在 VAE 中,编码器不再输出一个确定点,而是输出一个概率分布:

q ϕ ( z ∣ x ) q_\phi(z|x) qϕ(zx)

通常,我们令这个分布是高斯分布:

q ϕ ( z ∣ x ) = N ( z ; μ ϕ ( x ) , diag ⁡ ( σ ϕ 2 ( x ) ) ) q_\phi(z|x) = \mathcal{N} \left( z; \mu_\phi(x), \operatorname{diag}(\sigma_\phi^2(x)) \right) qϕ(zx)=N(z;μϕ(x),diag(σϕ2(x)))

这里:

  • μ ϕ ( x ) \mu_\phi(x) μϕ(x) 是编码器输出的均值;
  • σ ϕ 2 ( x ) \sigma_\phi^2(x) σϕ2(x) 是编码器输出的方差;
  • diag ⁡ ( σ ϕ 2 ( x ) ) \operatorname{diag}(\sigma_\phi^2(x)) diag(σϕ2(x)) 表示每个潜变量维度有自己的方差。

也就是说,给定一个数据点 x x x ,编码器不再只告诉我们 “它的潜变量是什么”,而是告诉我们 “它的潜变量大概分布在哪里”。

编码过程变成从这个分布中采样:

z ∼ q ϕ ( ⋅ ∣ x ) z\sim q_\phi(\cdot|x) zqϕ(x)

这就是随机编码。

对应地,解码器也可以写成一个条件概率分布:

p θ ( x ∣ z ) = N ( x ; μ θ ( z ) , σ θ 2 I d ) p_\theta(x|z) = \mathcal{N} \left( x; \mu_\theta(z), \sigma_\theta^2 I_d \right) pθ(xz)=N(x;μθ(z),σθ2Id)

这里:

  • μ θ ( z ) \mu_\theta(z) μθ(z) 是解码器输出的重建图像均值;
  • σ θ 2 \sigma_\theta^2 σθ2 通常可以看作固定的解码方差;
  • I d I_d Id 表示各维度独立同方差。

因此,VAE 的编码和解码过程可以写成:

z ∼ q ϕ ( ⋅ ∣ x ) , x ∼ p θ ( ⋅ ∣ z ) z\sim q_\phi(\cdot|x), \qquad x\sim p_\theta(\cdot|z) zqϕ(x),xpθ(z)

与标准自编码器相比,VAE 的关键变化是:

潜变量不再是确定点,而是从一个由编码器给出的分布中采样得到。

这样做的目的,是为了给潜空间分布加入可控的概率结构。

OK,那么我们来谈谈如何训练这个东西。

首先,VAE 仍然需要重建能力。也就是说,我们希望一个数据点 x x x 被编码为潜变量 z z z 后,解码器能够从 z z z 中恢复出原始 x x x 。由于 VAE 是概率模型,重建损失可以用负对数似然来定义:

L V A E - R e c o n ( ϕ , θ ) = − E x ∼ p d a t a , z ∼ q ϕ ( z ∣ x ) [ log ⁡ p θ ( x ∣ z ) ] \mathcal{L}_{\mathrm{VAE\text{-}Recon}}(\phi,\theta) = - \mathbb{E}_{x\sim p_{\mathrm{data}},z\sim q_\phi(z|x)} \left[ \log p_\theta(x|z) \right] LVAE-Recon(ϕ,θ)=Expdata,zqϕ(zx)[logpθ(xz)]

这个损失的含义是:先从数据集中采样一个图像 x x x ,再从编码器分布 q ϕ ( z ∣ x ) q_\phi(z|x) qϕ(zx) 中采样潜变量 z z z ,然后希望解码器分布 p θ ( x ∣ z ) p_\theta(x|z) pθ(xz) 给原始图像 x x x 较高概率。换句话说,给定潜变量 z z z 后,原始图像 x x x 应该是一个高似然样本。

由于我们假设:

p θ ( x ∣ z ) = N ( x ; μ θ ( z ) , σ θ 2 I d ) p_\theta(x|z) = \mathcal{N} \left( x; \mu_\theta(z), \sigma_\theta^2I_d \right) pθ(xz)=N(x;μθ(z),σθ2Id)

所以负对数似然可以化简为均方误差形式。忽略与参数无关的常数项,有:

− log ⁡ p θ ( x ∣ z ) = 1 2 σ θ 2 ∥ x − μ θ ( z ) ∥ 2 + c o n s t -\log p_\theta(x|z) = \frac{1}{2\sigma_\theta^2} \left\| x-\mu_\theta(z) \right\|^2 + \mathrm{const} logpθ(xz)=2σθ21xμθ(z)2+const

因此:

L V A E - R e c o n ( ϕ , θ ) = E x ∼ p d a t a , z ∼ q ϕ ( z ∣ x ) [ 1 2 σ θ 2 ∥ x − μ θ ( z ) ∥ 2 ] \mathcal{L}_{\mathrm{VAE\text{-}Recon}}(\phi,\theta) = \mathbb{E}_{x\sim p_{\mathrm{data}},z\sim q_\phi(z|x)} \left[ \frac{1}{2\sigma_\theta^2} \left\| x-\mu_\theta(z) \right\|^2 \right] LVAE-Recon(ϕ,θ)=Expdata,zqϕ(zx)[2σθ21xμθ(z)2]

这说明 VAE 的重建损失本质上仍然类似标准自编码器的重建误差,只不过编码阶段从确定性映射变成了随机采样。

仅有重建损失仍然不够,因为它只要求能够重建图像,并不保证潜空间分布容易建模。我们希望编码后的潜变量分布接近一个简单先验分布,通常选择标准高斯:

p p r i o r ( z ) = N ( 0 , I k ) p_{\mathrm{prior}}(z)=\mathcal{N}(0,I_k) pprior(z)=N(0,Ik)

为了衡量两个分布之间的差异,我们引入 KL 散度(Kullback-Leibler divergence)

对于两个概率密度 q ( x ) q(x) q(x) p ( x ) p(x) p(x) ,KL 散度定义为:

D K L ( q   ∥   p ) = ∫ q ( x ) log ⁡ q ( x ) p ( x ) d x = E X ∼ q [ log ⁡ q ( X ) p ( X ) ] D_{\mathrm{KL}}(q\,\|\,p) = \int q(x)\log\frac{q(x)}{p(x)}dx = \mathbb{E}_{X\sim q} \left[ \log\frac{q(X)}{p(X)} \right] DKL(qp)=q(x)logp(x)q(x)dx=EXq[logp(X)q(X)]

KL 散度可以理解为衡量分布 q q q 与分布 p p p 的差异程度。

它有两个重要性质:

  • D K L ( q ∣ p ) ≥ 0 D_{\mathrm{KL}}(q|p)\geq 0 DKL(qp)0
  • D K L ( q ∣ p ) = 0 ⟺ q = p D_{\mathrm{KL}}(q|p)=0 \quad\Longleftrightarrow\quad q=p DKL(qp)=0q=p

需要注意,KL 散度严格来说不是距离,因为它通常不对称即 D K L ( q ∣ p ) ≠ D K L ( p ∣ q ) D_{\mathrm{KL}}(q|p) \neq D_{\mathrm{KL}}(p|q) DKL(qp)=DKL(pq) ,但它仍然是衡量两个分布差异的常用工具。

在 VAE 中,我们希望每个数据点 x x x 对应的编码分布 q ϕ ( z ∣ x ) q_\phi(z|x) qϕ(zx) 接近标准高斯先验:

p p r i o r ( z ) = N ( 0 , I k ) p_{\mathrm{prior}}(z)=\mathcal{N}(0,I_k) pprior(z)=N(0,Ik)

因此定义先验正则项:

L V A E - P r i o r ( ϕ ) = E x ∼ p d a t a [ D K L ( q ϕ ( ⋅ ∣ x )   ∥   p p r i o r ) ] \mathcal{L}_{\mathrm{VAE\text{-}Prior}}(\phi) = \mathbb{E}_{x\sim p_{\mathrm{data}}} \left[ D_{\mathrm{KL}} \left( q_\phi(\cdot|x) \,\| \, p_{\mathrm{prior}} \right) \right] LVAE-Prior(ϕ)=Expdata[DKL(qϕ(x)pprior)]

这个损失的含义是:对于每个数据点 x x x ,编码器给出的潜变量分布都应该尽量接近标准高斯。这样做可以让潜空间更加规整,避免标准自编码器中潜分布过于混乱的问题。

接下来,我们把 KL 散度具体化到两个高斯分布之间。

假设:

q ( x ) = N ( x ; μ q , σ q 2 I d ) p ( x ) = N ( x ; μ p , σ p 2 I d ) q(x)=\mathcal{N}(x;\mu_q,\sigma_q^2I_d) \\[6pt] p(x)=\mathcal{N}(x;\mu_p,\sigma_p^2I_d) q(x)=N(x;μq,σq2Id)p(x)=N(x;μp,σp2Id)

那么它们之间的 KL 散度可以写成:

D K L ( q ∣ p ) = 1 2 ( K ( σ q 2 σ p 2 ) + ∥ μ q − μ p ∥ 2 σ p 2 ) D_{\mathrm{KL}}(q|p) = \frac{1}{2} \left( \mathcal{K} \left( \frac{\sigma_q^2}{\sigma_p^2} \right) + \frac{\|\mu_q-\mu_p\|^2}{\sigma_p^2} \right) DKL(qp)=21(K(σp2σq2)+σp2μqμp2)

其中:

K ( α ) = ∑ i ( α i − log ⁡ α i − 1 ) \mathcal{K}(\alpha) = \sum_i \left( \alpha_i-\log\alpha_i-1 \right) K(α)=i(αilogαi1)

如果是一维情况,则是:

K ( α ) = α − log ⁡ α − 1 \mathcal{K}(\alpha)=\alpha-\log\alpha-1 K(α)=αlogα1

这个函数在 α = 1 \alpha=1 α=1 时取得最小值 0 0 0 。也就是说,当两个高斯分布的方差相同时,这一项最小。

因此,KL 散度中有两部分:

  • 第一部分: K ( σ q 2 σ p 2 ) \mathcal{K} \left( \frac{\sigma_q^2}{\sigma_p^2} \right) K(σp2σq2) ,衡量两个分布的方差差异。
  • 第二部分: ∥ μ q − μ p ∥ 2 σ p 2 \frac{\|\mu_q-\mu_p\|^2}{\sigma_p^2} σp2μqμp2 ,衡量两个分布的均值差异。

这和直觉一致:两个高斯分布越相似,它们的均值和方差就应该越接近。

在 VAE 中,我们令先验分布为标准高斯:

p p r i o r ( z ) = N ( 0 , I k ) p_{\mathrm{prior}}(z)=\mathcal{N}(0,I_k) pprior(z)=N(0,Ik)

编码分布为:

q ϕ ( z ∣ x ) = N ( z ; μ ϕ ( x ) , diag ⁡ ( σ ϕ 2 ( x ) ) ) q_\phi(z|x) = \mathcal{N} \left( z; \mu_\phi(x), \operatorname{diag}(\sigma_\phi^2(x)) \right) qϕ(zx)=N(z;μϕ(x),diag(σϕ2(x)))

代入高斯 KL 公式,可以得到:

L V A E - P r i o r ( ϕ ) = E x ∼ p d a t a [ D K L ( q ϕ ( ⋅ ∣ x )    ∥    N ( 0 , I k ) ) ] = E x ∼ p d a t a [ 1 2 K ( σ ϕ 2 ( x ) ) + 1 2 ∥ μ ϕ ( x ) ∥ 2 ] \mathcal{L}_{\mathrm{VAE\text{-}Prior}}(\phi) = \mathbb{E}_{x\sim p_{\mathrm{data}}} \left[ D_{\mathrm{KL}} \left( q_\phi(\cdot|x) \; \| \; \mathcal{N}(0,I_k) \right) \right] = \mathbb{E}_{x\sim p_{\mathrm{data}}} \left[ \frac{1}{2} \mathcal{K} \left( \sigma_\phi^2(x) \right) + \frac{1}{2} \left\| \mu_\phi(x) \right\|^2 \right] LVAE-Prior(ϕ)=Expdata[DKL(qϕ(x)N(0,Ik))]=Expdata[21K(σϕ2(x))+21μϕ(x)2]

这里的含义非常清楚:

  • K ( σ ϕ 2 ( x ) ) \mathcal{K}(\sigma_\phi^2(x)) K(σϕ2(x)) 会鼓励编码分布的方差接近 1 1 1
  • ∥ μ ϕ ( x ) ∥ 2 \|\mu_\phi(x)\|^2 μϕ(x)2 会鼓励编码分布的均值接近 0 0 0

也就是说,KL 项会让每个 q ϕ ( z ∣ x ) q_\phi(z|x) qϕ(zx) 尽量接近标准高斯分布。

现在我们有两个目标。

第一个目标是 重建:编码后的潜变量应该保留足够信息,使解码器能够恢复原始图像。这由重建损失控制:

L V A E - R e c o n ( ϕ , θ ) \mathcal{L}_{\mathrm{VAE\text{-}Recon}}(\phi,\theta) LVAE-Recon(ϕ,θ)

第二个目标是 规整潜空间:编码分布应该接近标准高斯,使后续生成模型更容易学习。这由先验损失控制:

L V A E - P r i o r ( ϕ ) \mathcal{L}_{\mathrm{VAE\text{-}Prior}}(\phi) LVAE-Prior(ϕ)

但是这两个目标存在张力。

如果只优化重建损失,模型可能得到一个很不规则的潜空间分布,后续扩散模型难以学习。如果只优化先验损失,最简单的解是:

μ ϕ ( x ) = 0 , σ ϕ ( x ) = 1 \mu_\phi(x)=0, \qquad \sigma_\phi(x)=1 μϕ(x)=0,σϕ(x)=1

也就是说,编码器完全忽略输入 x x x ,始终输出标准高斯分布。这样潜空间当然很规整,但它不包含任何图像信息,无法重建原始图像。

因此,我们需要把两个损失结合起来。引入一个非负权重 β \beta β ,定义 VAE 总损失:

L V A E ( ϕ , θ ) = L V A E - R e c o n ( ϕ , θ ) + β L V A E - P r i o r ( ϕ ) \mathcal{L}_{\mathrm{VAE}}(\phi,\theta) = \mathcal{L}_{\mathrm{VAE\text{-}Recon}}(\phi,\theta) + \beta \mathcal{L}_{\mathrm{VAE\text{-}Prior}}(\phi) LVAE(ϕ,θ)=LVAE-Recon(ϕ,θ)+βLVAE-Prior(ϕ)

展开后可以写成:

L V A E ( ϕ , θ ) = L V A E - R e c o n ( ϕ , θ ) + β L V A E - P r i o r ( ϕ ) = E x ∼ p d a t a , z ∼ q ϕ ( z ∣ x ) [ 1 2 σ θ 2 ( z ) ∥ x − μ θ ( z ) ∥ 2 ⏟ recon. error + β 2 K ( σ ϕ 2 ( x ) ) ⏟ make latent variance = 1 + β 2 ∥ μ ϕ ( x ) ∥ 2 ⏟ make latent mean = 0 ] \begin{align*} \mathcal L_{\mathrm{VAE}}(\phi,\theta) &= \mathcal L_{\mathrm{VAE\text{-}Recon}}(\phi,\theta) + \beta \mathcal L_{\mathrm{VAE\text{-}Prior}}(\phi) \\ &= \mathbb E_{x\sim p_{\mathrm{data}}, z\sim q_\phi(z|x)} \Bigg[ \underbrace{ \frac{1}{2\sigma_\theta^2(z)} \|x-\mu_\theta(z)\|^2 }_{\text{recon. error}} + \underbrace{ \frac\beta2 \mathcal K \left( \sigma_\phi^2(x) \right) }_{\text{make latent variance}=1} + \underbrace{ \frac\beta2 \|\mu_\phi(x)\|^2 }_{\text{make latent mean}=0} \Bigg] \end{align*} LVAE(ϕ,θ)=LVAE-Recon(ϕ,θ)+βLVAE-Prior(ϕ)=Expdata,zqϕ(zx)[recon. error 2σθ2(z)1xμθ(z)2+make latent variance=1 2βK(σϕ2(x))+make latent mean=0 2βμϕ(x)2]

这个公式就是 VAE 的核心训练目标。

其中, β \beta β 控制重建质量和潜空间规整程度之间的权衡。如果 β \beta β 太小,模型更关注重建,潜空间可能不够规整;如果 β \beta β 太大,模型更关注让潜变量接近标准高斯,但可能损失太多图像信息,导致重建质量下降。因此, β \beta β 的选择通常是一个经验问题,需要根据数据、模型结构和后续生成任务进行调节。


Q:这里说的 “潜空间好学” 到底是什么意思?

A:它不是一个特别严格的数学概念,而是指编码后的潜分布应该更适合后续生成模型学习。比如,潜分布最好相对平滑、连续、接近简单分布,而不是非常破碎、扭曲或者高度不规则。

如果潜空间分布很混乱,那么即使自编码器能重建图像,后续在潜空间中训练扩散模型仍然会很困难。因此,VAE 中的 KL 项就是为了让潜空间更规整。

Q:为什么不只优化 KL 项,让潜空间完全变成标准高斯?

A:因为这样会让编码器忽略输入数据。

如果只最小化:

D K L ( q ϕ ( z ∣ x )   ∥   N ( 0 , I ) ) D_{\mathrm{KL}} \left( q_\phi(z|x)\,\|\,\mathcal{N}(0,I) \right) DKL(qϕ(zx)N(0,I))

最简单的做法就是令:

q ϕ ( z ∣ x ) = N ( 0 , I ) q_\phi(z|x)=\mathcal{N}(0,I) qϕ(zx)=N(0,I)

对所有 x x x 都一样。这样潜空间确实非常接近标准高斯,但它不包含任何关于输入图像的信息,解码器无法重建原始图像。

所以 VAE 必须同时优化重建损失和先验损失。

Q: β \beta β 应该怎么选?

A β \beta β 控制重建质量和潜空间规整程度之间的权衡。

如果 β \beta β 较小,模型更重视重建,图像还原质量可能更好,但潜空间可能更不规则。

如果 β \beta β 较大,潜空间会更接近标准高斯,但图像信息可能被压缩过度,导致重建变差。

因此, β \beta β 没有一个通用最优值,通常需要根据数据集、模型容量和后续生成模型的训练效果来经验选择。

Q:VAE 会不会出现类似模式崩溃的问题?

A:有可能。VAE 中常见的相关问题包括 posterior collapse,也就是编码器输出的潜变量不再携带足够的输入信息,而是退化到接近先验分布。

这通常发生在 KL 项太强、解码器太强或者训练设置不合适时。因此,实际模型中常常会使用 KL annealing、loss clipping、调整 β \beta β 等工程技巧来缓解这个问题。


到这里,我们已经写出了 VAE 的训练损失。不过还有一个实际训练问题:

损失中包含从编码分布中采样:

z ∼ q ϕ ( z ∣ x ) z\sim q_\phi(z|x) zqϕ(zx)

而这个分布本身依赖于参数 ϕ \phi ϕ

q ϕ ( z ∣ x ) = N ( z ; μ ϕ ( x ) , diag ⁡ ( σ ϕ 2 ( x ) ) ) q_\phi(z|x) = \mathcal{N} \left( z; \mu_\phi(x), \operatorname{diag}(\sigma_\phi^2(x)) \right) qϕ(zx)=N(z;μϕ(x),diag(σϕ2(x)))

如果直接对采样操作求梯度,会遇到问题。因为普通采样操作本身不是一个容易反向传播的确定性函数,不能简单地把它当作普通网络层来微分。

更准确地说,并不是 “从依赖 ϕ \phi ϕ 的分布中采样就绝对不能训练”,而是不能用普通的采样方式直接得到低方差、可反向传播的梯度估计。

因此,我们需要使用 重参数化技巧(reparameterization trick)

5. Reparameterization Trick

前面我们已经定义了 VAE 的编码分布:

q ϕ ( z ∣ x ) = N ( z ; μ ϕ ( x ) , diag ⁡ ( σ ϕ 2 ( x ) ) ) q_\phi(z|x) = \mathcal{N} \left( z; \mu_\phi(x), \operatorname{diag}(\sigma_\phi^2(x)) \right) qϕ(zx)=N(z;μϕ(x),diag(σϕ2(x)))

也就是说,给定一个数据点 x x x ,编码器输出两个量:

μ ϕ ( x ) , σ ϕ 2 ( x ) \mu_\phi(x), \qquad \sigma_\phi^2(x) μϕ(x),σϕ2(x)

它们分别表示潜变量分布的均值和方差。

现在的问题是,我们需要从这个分布中采样:

z ∼ q ϕ ( ⋅ ∣ x ) z\sim q_\phi(\cdot|x) zqϕ(x)

但这个采样分布本身依赖于网络参数 ϕ \phi ϕ 。如果直接写成 “从 q ϕ ( z ∣ x ) q_\phi(z|x) qϕ(zx) 中采样”,那么采样操作不容易直接放进标准反向传播框架中求梯度。

因此,我们需要使用 重参数化技巧(reparameterization trick)

它利用高斯分布的一个简单性质:

ϵ ∼ N ( 0 , I k ) , z = μ ϕ ( x ) + σ ϕ ( x ) ⊙ ϵ ⇒ z ∼ q ϕ ( ⋅ ∣ x ) \epsilon\sim\mathcal N(0,I_k), \quad z=\mu_\phi(x)+\sigma_\phi(x) \odot \epsilon \quad\Rightarrow\quad z\sim q_\phi(\cdot|x) ϵN(0,Ik),z=μϕ(x)+σϕ(x)ϵzqϕ(x)

标准高斯经过平移和缩放后,就得到具有指定均值和方差的高斯分布。

使用重参数化之后,VAE 的损失可以写成:

L V A E ( ϕ , θ ) = E x ∼ p d a t a ( x ) , ϵ ∼ N ( 0 , I k ) [ 1 2 σ 2 ∥ x − μ θ ( μ ϕ ( x ) + σ ϕ ( x ) ⊙ ϵ ) ∥ 2 + β 2 K ( σ ϕ 2 ( x ) ) + β 2 ∥ μ ϕ ( x ) ∥ 2 ] \mathcal L_{\mathrm{VAE}}(\phi,\theta) = \mathbb E_{ x\sim p_{\mathrm{data}}(x), \epsilon\sim\mathcal N(0,I_k) } \Bigg[ \frac{1}{2\sigma^2} \left\| x- \mu_\theta \left( \mu_\phi(x)+\sigma_\phi(x) \odot \epsilon \right) \right\|^2 + \frac\beta2 \mathcal K \left( \sigma_\phi^2(x) \right) + \frac\beta2 \|\mu_\phi(x)\|^2 \Bigg] LVAE(ϕ,θ)=Expdata(x),ϵN(0,Ik)[2σ21xμθ(μϕ(x)+σϕ(x)ϵ)2+2βK(σϕ2(x))+2βμϕ(x)2]

这里需要注意一个关键点:损失函数本身当然仍然依赖于参数 ϕ \phi ϕ θ \theta θ ,但期望所依赖的随机变量分布已经不依赖于参数了

也就是说,我们现在是在下面两个分布上取期望:

x ∼ p data ( x ) , ϵ ∼ N ( 0 , I k ) x \sim p_{\text{data}}(x), \qquad \epsilon \sim \mathcal{N}(0,I_k) xpdata(x),ϵN(0,Ik)

这两个分布都不依赖于 ϕ \phi ϕ θ \theta θ 。参数只出现在被期望的函数内部,而这个函数由神经网络和普通可微运算组成,因此可以直接使用 PyTorch 等框架进行反向传播。

这就是重参数化技巧的本质:

将 “从依赖参数的分布中采样”,改写成 “从固定噪声分布中采样 + 可微变换”。


Q & A

Q:为什么我们希望潜变量分布接近正态分布?

A:不是重参数化技巧本身的要求,而是 VAE / 潜扩散模型整体设计的要求。

我们的最终目标不是单纯训练一个自编码器,而是先把图像编码到潜空间,再在潜空间中训练扩散模型或流模型。因此,我们希望编码后的潜空间分布比较规整、平滑、容易学习。

标准高斯分布 N ( 0 , I ) \mathcal{N}(0,I) N(0,I) 是一个非常自然的选择。它结构简单、单峰、各向同性,而且负对数密度是一个凸二次函数。若潜空间分布接近高斯,那么后续在潜空间中学习生成模型会更容易。

Q:既然每个 q ϕ ( z ∣ x ) q_\phi(z|x) qϕ(zx) 都是高斯分布,为什么整体潜分布还可能不是高斯?

A:因为单个样本对应的编码分布是高斯,不代表所有数据样本混合后的潜分布仍然是高斯。

整体潜分布可以写成:

q ϕ ( z ) = ∫ q ϕ ( z ∣ x ) p data ( x ) d x q_{\phi}(z) = \int q_{\phi}(z|x) p_{\text{data}}(x)dx qϕ(z)=qϕ(zx)pdata(x)dx

这相当于是很多个高斯分布的混合。每个数据点 x x x 都有自己的均值 μ ϕ ( x ) \mu_\phi(x) μϕ(x) 和方差 σ ϕ 2 ( x ) \sigma_\phi^2(x) σϕ2(x) ,这些高斯分布混合在一起后,整体分布可以非常复杂,甚至高度非高斯。

事实上,高斯混合模型本身就具有很强的表达能力,可以近似很多复杂分布。因此,虽然每个条件分布 q ϕ ( z ∣ x ) q_\phi(z|x) qϕ(zx) 是高斯,边缘潜分布 q ϕ ( z ) q_\phi(z) qϕ(z) 仍然可能非常复杂。

VAE 中的 KL 正则项正是为了约束这些条件高斯分布不要过于散乱,使整体潜空间更规整。


有了重参数化技巧之后,VAE 就可以像普通神经网络一样训练。

训练过程可以概括为:

给定一个 mini-batch:

x i i = 1 B {x_i}_{i=1}^{B} xii=1B

首先用编码器分别预测每个样本的潜变量均值和对数方差:

μ i = μ ϕ ( x i ) , log ⁡ σ i 2 = log ⁡ σ ϕ 2 ( x i ) \mu_i=\mu_\phi(x_i), \qquad \log\sigma_i^2=\log\sigma_\phi^2(x_i) μi=μϕ(xi),logσi2=logσϕ2(xi)

然后从标准高斯中采样噪声:

ϵ i ∼ N ( 0 , I k ) \epsilon_i\sim\mathcal{N}(0,I_k) ϵiN(0,Ik)

再通过重参数化得到潜变量:

z i = μ i + σ i ⊙ ϵ i z_i=\mu_i+\sigma_i\odot\epsilon_i zi=μi+σiϵi

接着用解码器重建图像:

x ^ i = μ θ ( z i ) \hat{x}_i=\mu_\theta(z_i) x^i=μθ(zi)

重建损失为:

L r e c o n = 1 B ∑ i = 1 B 1 2 σ 2 ∣ x i − x ^ i ∣ 2 \mathcal{L}_{\mathrm{recon}} = \frac{1}{B} \sum_{i=1}^{B} \frac{1}{2\sigma^2} \left| x_i-\hat{x}_i \right|^2 Lrecon=B1i=1B2σ21xix^i2

KL 损失为:

L K L = 1 B ∑ i = 1 B 1 2 ∑ j = 1 k ( μ i , j 2 + σ i , j 2 − log ⁡ σ i , j 2 − 1 ) \mathcal{L}_{\mathrm{KL}} = \frac{1}{B} \sum_{i=1}^{B} \frac{1}{2} \sum_{j=1}^{k} \left( \mu_{i,j}^2 + \sigma_{i,j}^2 - \log\sigma_{i,j}^2 - 1 \right) LKL=B1i=1B21j=1k(μi,j2+σi,j2logσi,j21)

总损失为:

L = L r e c o n + β L K L \mathcal{L} = \mathcal{L}_{\mathrm{recon}} + \beta \mathcal{L}_{\mathrm{KL}} L=Lrecon+βLKL

然后对编码器参数 ϕ \phi ϕ 和解码器参数 θ \theta θ 做梯度更新。

这个算法实际上就是把前面推导的 VAE 损失变成了可实现的训练流程。它的关键步骤是第 4 步:

z i = μ i + σ i ⊙ ϵ i z_i=\mu_i+\sigma_i\odot\epsilon_i zi=μi+σiϵi

也就是重参数化。如果没有这一步,采样过程就不容易直接用普通反向传播训练。

这和强化学习中的 REINFORCE 有一点类似:REINFORCE 也处理 “期望中的分布依赖参数” 的问题,但它通常使用 score-function estimator,方差较大。而 VAE 中由于高斯分布可以重参数化,所以我们可以使用更高效的路径梯度方法。


Q & A

Q:为什么原来的损失不能直接反向传播?

A:更准确地说,不是 “计算上完全不能反向传播”,而是如果直接把采样过程当作普通操作处理,就得不到正确的、低方差的梯度估计。

原来的形式中有:

z ∼ q ϕ ( z ∣ x ) z \sim q_{\phi}(z|x) zqϕ(zx)

q ϕ q_\phi qϕ 依赖于参数 ϕ \phi ϕ 。如果直接采样 z z z ,采样结果本身会随着 ϕ \phi ϕ 的变化而变化。这个依赖关系不是普通确定性计算图能够自然表达的。

重参数化技巧把它改写成:

z = μ ϕ ( x ) + σ ϕ ( x ) ⊙ ϵ , ϵ ∼ N ( 0 , I ) z=\mu_{\phi}(x)+\sigma_{\phi}(x)\odot\epsilon, \qquad \epsilon \sim \mathcal{N}(0,I) z=μϕ(x)+σϕ(x)ϵ,ϵN(0,I)

这样随机性只来自 ϵ \epsilon ϵ ,而 ϵ \epsilon ϵ 不依赖于参数。于是 z z z 变成了关于 ϕ \phi ϕ 的确定性可微函数,因此可以通过普通反向传播计算梯度。

Q:如果直接从 q ϕ ( z ∣ x ) q_\phi(z|x) qϕ(zx) 采样,就一定没有办法训练吗?

A:不是绝对不能训练。理论上,还有其他梯度估计方法,例如 score-function estimator,也就是强化学习中常见的 REINFORCE 类方法。

但是这类方法通常方差较大,训练效率较低。VAE 中使用的重参数化技巧是一种更直接、更稳定、更高效的方式。它几乎只是对高斯采样过程的一种等价改写,没有额外信息损失,也几乎没有额外计算成本。

Q:重参数化技巧会不会改变原来的分布?

A:不会。

因为:

ϵ ∼ N ( 0 , I ) \epsilon \sim \mathcal{N}(0,I) ϵN(0,I)

经过:

z = μ ϕ ( x ) + σ ϕ ( x ) ⊙ ϵ z=\mu_{\phi}(x)+\sigma_{\phi}(x)\odot\epsilon z=μϕ(x)+σϕ(x)ϵ

之后, z z z 的分布正好就是:

N ( μ ϕ ( x ) , diag ⁡ ( σ ϕ 2 ( x ) ) ) \mathcal{N} \left(\mu_\phi(x), \operatorname{diag}(\sigma_\phi^2(x)) \right) N(μϕ(x),diag(σϕ2(x)))

也就是说,重参数化只是改变了采样的表达方式,并没有改变采样分布本身。


6. Latent Diffusion Models (LDMs) - Recipe

接下来,我们把 VAE 和扩散模型结合起来,就得到 潜扩散模型(Latent Diffusion Models, LDMs)

LDM 的流程非常清晰。

第一步,准备训练数据:

x 1 , … , x N x_1,\dots,x_N x1,,xN

例如大规模图像数据集。

第二步,用已经训练好的 VAE 编码器把所有图像转换为潜变量。对于 VAE 来说,通常可以直接取编码分布的均值作为潜表示:

z i = μ ϕ ( x i ) z_i=\mu_\phi(x_i) zi=μϕ(xi)

第三步,得到潜变量数据集:

z 1 , … , z N z_1,\dots,z_N z1,,zN

这个数据集的维度远小于原始高分辨率图像。

第四步,在潜变量数据集上训练扩散模型或流模型。也就是说,现在生成模型不再直接生成像素图像,而是生成潜变量:

z ∼ p l a t e n t z\sim p_{\mathrm{latent}} zplatent

第五步,采样完成后,用 VAE 解码器把生成的潜变量映射回图像空间:

x ^ = μ θ ( z ) \hat{x}=\mu_\theta(z) x^=μθ(z)

最终返回解码后的图像作为生成结果。

所以 LDM 的整体流程可以概括为:

image → latent → diffusion in latent space → generated latent → decoded image \text{image} \rightarrow \text{latent} \rightarrow \text{diffusion in latent space} \rightarrow \text{generated latent} \rightarrow \text{decoded image} imagelatentdiffusion in latent spacegenerated latentdecoded image

这和前面讲过的扩散模型流程本质上是一样的,只是训练数据从原始图像变成了 VAE 编码后的潜变量。也就是说潜扩散模型并没有改变扩散模型的基本训练方法,而是改变了扩散模型所建模的数据空间。


Q & A

Q:使用潜空间会不会存在信息损失?

A:会存在信息损失,因为 VAE 本质上是压缩模型,压缩就不可能完全无损。

但关键问题不是 “有没有信息损失”,而是 “损失的信息是否重要”。

在图像生成任务中,如果自编码器训练得很好,编码再解码后的图像与原图之间的差异通常非常小,很多时候肉眼几乎察觉不到。因此,对生成任务来说,这种信息损失通常是可以接受的。

但是,这也说明自编码器非常重要。潜扩散模型的效果不仅取决于扩散模型本身,也高度依赖 VAE 的质量。如果 VAE 丢失了重要结构、纹理或语义信息,那么后续扩散模型即使学得很好,也只能在一个已经损失信息的潜空间中生成样本。


潜扩散模型的核心优势是显著降低计算成本。

例如,在 Stable Diffusion 中,原始图像形状可能是 [ 3 , 256 , 256 ] [3,256,256] [3,256,256] ,经过 VAE 编码后变成 [ 4 , 32 , 32 ] [4,32,32] [4,32,32] 。原始图像维度是 3 × 256 × 256 = 196,608 3\times256\times256=196{,}608 3×256×256=196,608 ,潜变量维度是 4 × 32 × 32 = 4,096 4\times32\times32=4{,}096 4×32×32=4,096 ,压缩比例大约是 48 × 48\times 48× ,这意味着扩散模型不需要在 196,608 维像素空间中建模,而是在 4,096 维潜空间中建模,计算量和显存压力都会大幅降低。

再比如 FLUX 2.0 中,图像形状可能从 [ 3 , 1024 , 1024 ] [3,1024,1024] [3,1024,1024] 压缩到 [ 32 , 64 , 64 ] [32,64,64] [32,64,64] ,这同样显著降低了生成模型需要处理的空间维度。这也是为什么现代图像和视频生成模型几乎都会在潜空间中生成,而不是直接在像素空间中生成。

当然,潜空间生成也有一个前提:VAE 必须足够好。如果自编码器压缩时丢失了重要信息,那么扩散模型即使在潜空间中生成得很好,最终解码回图像时也可能出现质量损失。

因此,实际的大规模图像生成系统中,自编码器本身也是非常关键的组件。

7. Neural Network Architectures

前面我们讨论了“在哪个空间中生成”,也就是从像素空间转向潜空间。

接下来要讨论的是另一个关键问题:

如何用神经网络参数化向量场?

在前面的理论中,我们一直写:

u t θ ( x ∣ y ) u_t^\theta(x|y) utθ(xy)

这里:

  • θ \theta θ 是神经网络参数;
  • t t t 是时间;
  • x x x 是当前的带噪样本,在 LDM 中通常是潜图像;
  • y y y 是 prompt;
  • 输出 u t θ ( x ∣ y ) u_t^\theta(x|y) utθ(xy) 是一个和 x x x 形状相同的向量场。

也就是说,网络需要完成的任务是:

输入时间、潜图像和文本条件,输出一个与潜图像同形状的速度场或噪声预测。

例如,如果潜图像的形状是 [ 4 , 32 , 32 ] [4,32,32] [4,32,32] ,那么网络输出也应该是 [ 4 , 32 , 32 ] [4,32,32] [4,32,32] ,这就要求网络既能处理图像结构,又能处理时间信息,还能处理文本 prompt。

对于现代图像生成模型,一个常见选择是使用 Transformer 类架构,尤其是 Diffusion Transformer(DiT)

在进入 DiT 之前,需要先讨论三个输入分别如何编码:

1. 时间如何编码;

2. 文本 prompt 如何编码;

3. 图像或潜变量如何编码成 Transformer 可以处理的形式。

我们一个个来看,先看时间:

时间 t t t 是一个一维标量。但是神经网络中的其他输入通常是高维向量或向量序列。如果直接把一个标量 t t t 输入网络,它的表达能力可能不够。因此,我们通常会把时间编码成一个高维向量。

一种常见做法是使用正弦位置编码:

T i m e E m b ( t ) = 1 d [ cos ⁡ ( 2 π w 1 t ) , … , cos ⁡ ( 2 π w d / 2 t ) , sin ⁡ ( 2 π w 1 t ) , … , sin ⁡ ( 2 π w d / 2 t ) ] T \mathrm{TimeEmb}(t) = \frac{1}{\sqrt{d}} \left[ \cos(2\pi w_1t), \dots, \cos(2\pi w_{d/2}t), \sin(2\pi w_1t), \dots, \sin(2\pi w_{d/2}t) \right]^T TimeEmb(t)=d 1[cos(2πw1t),,cos(2πwd/2t),sin(2πw1t),,sin(2πwd/2t)]T

其中频率 w i w_i wi 通常按几何级数选取:

w i = w min ⁡ ( w max ⁡ w min ⁡ ) i − 1 d / 2 − 1 , i = 1 , … , d / 2 w_i = w_{\min} \left( \frac{w_{\max}}{w_{\min}} \right)^{ \frac{i-1}{d/2-1} }, \qquad i=1,\dots,d/2 wi=wmin(wminwmax)d/21i1,i=1,,d/2

这样,一个标量时间 t t t 就被映射成了一个 d d d 维向量。

这种编码的具体形式并不是最重要的。关键是:

时间需要被提升到高维表示中,使网络能够更容易感知不同时间步之间的差异。

在扩散模型中,时间非常重要,因为不同时间步对应不同噪声水平,网络在不同时间下需要做的事情也不同。

接下来是 prompt 的编码:

例如 prompt 是:

A dog running on grass in a park at sunshine in an Italian city.

神经网络不能直接处理原始字符串,因此需要先把文本转换成向量表示。

现代文本到图像生成模型通常使用预训练语言模型或视觉语言模型来编码文本,例如:

  • CLIP text encoder;
  • T5 encoder;
  • 其他大语言模型 embedding。

这些模型会把原始 prompt 转换成一组文本 token 的向量表示。

如果 prompt 被分成 S S S 个 token,每个 token 的 embedding 维度是 k k k,那么 prompt embedding 可以写成:

P r o m p t E m b ( y r a w ) ∈ R S × k \mathrm{PromptEmb}(y_{\mathrm{raw}}) \in \mathbb{R}^{S\times k} PromptEmb(yraw)RS×k

也就是说,prompt 最终变成了一个长度为 S S S 的向量序列。

这一点非常重要,因为 Transformer 天然适合处理序列数据。文本 prompt 本来就是序列,经过 embedding 后就可以作为后续交叉注意力模块的输入。

接下来需要处理图像或潜图像:

图像通常具有形状:

x ∈ R C × H × W x\in\mathbb{R}^{C\times H\times W} xRC×H×W

它不是一个普通的一维序列,而是具有通道、高度和宽度的三维张量。

为了让 Transformer 处理图像,我们通常会把图像切分成多个 patch,然后把每个 patch 展平成一个向量。这个过程称为 patchify

x ~ = P a t c h i f y ( x ) ∈ R L × k \tilde{x} = \mathrm{Patchify}(x) \in\mathbb{R}^{L\times k} x~=Patchify(x)RL×k

其中:

  • L L L 是 patch 数量,也就是序列长度;
  • k k k 是每个 patch embedding 的维度。

这样,图像就被转换成了一个向量序列,可以输入 Transformer。

需要注意的是,在潜扩散模型中,我们处理的往往不是原始像素图像,而是 VAE 编码后的潜变量。但是潜变量通常仍然保留类似图像的形状,例如:

[ 4 , 32 , 32 ] [4,32,32] [4,32,32]

它仍然有通道、高度和宽度。因此,我们仍然可以对潜变量进行 patchify,把它转换成向量序列。

有了时间 embedding、prompt embedding 和图像 patch embedding 后,就可以构建扩散模型中的 Transformer 架构。

这种架构通常称为 Diffusion Transformer(DiT)

DiT 的核心思想是:

使用 Transformer 来参数化扩散模型中的向量场或噪声预测网络。

也就是说,我们不再使用传统 U-Net,而是使用 Transformer 来处理潜图像 token,并结合时间和文本条件输出预测结果。

DiT 的整体流程可以概括为三步:

第一步,构造输入。

时间被编码成向量:

t ~ = T i m e E m b ( t ) ∈ R k \tilde{t} = \mathrm{TimeEmb}(t) \in \mathbb{R}^{k} t~=TimeEmb(t)Rk

prompt 被编码成文本 token 序列:

y ~ = P r o m p t E m b ( y ) ∈ R S × k \tilde{y} = \mathrm{PromptEmb}(y) \in \mathbb{R}^{S\times k} y~=PromptEmb(y)RS×k

图像或潜图像被 patchify 成图像 token 序列:

x ~ 0 = P a t c h E m b ( x ) ∈ R N × k \tilde{x}_0=\mathrm{PatchEmb}(x)\in\mathbb{R}^{N\times k} x~0=PatchEmb(x)RN×k

这里 N N N 表示图像 patch token 的数量。

第二步,经过多个 DiTBlock。

每个 DiTBlock 接收当前图像 token、时间 embedding 和 prompt embedding:

x ~ i + 1 = D i T B l o c k ( x ~ i , t ~ , y ~ ) ∈ R N × k \tilde{x}_{i+1} = \mathrm{DiTBlock} ( \tilde{x}_i, \tilde{t}, \tilde{y} ) \in \mathbb{R}^{N\times k} x~i+1=DiTBlock(x~i,t~,y~)RN×k

经过多层迭代后,得到最终图像 token:

x ~ N \tilde{x}_N x~N

第三步,unpatchify。

Transformer 输出仍然是一个 token 序列,需要把它重新还原成图像或潜图像的形状:

u = U n p a t c h i f y ( x ~ N W ) ∈ R C × H × W u = \mathrm{Unpatchify} ( \tilde{x}_N W ) \in \mathbb{R}^{C\times H\times W} u=Unpatchify(x~NW)RC×H×W

最终输出 u u u 就是模型预测的向量场、噪声或 velocity,形状与输入 x x x 相同。

为了理解 DiTBlock,需要先回顾 Transformer 中最核心的操作:scaled dot-product attention

给定 queries、keys 和 values:

Q ∈ R N × d h , K ∈ R M × d h , V ∈ R M × d h Q\in\mathbb{R}^{N\times d_h}, \qquad K\in\mathbb{R}^{M\times d_h}, \qquad V\in\mathbb{R}^{M\times d_h} QRN×dh,KRM×dh,VRM×dh

注意力操作定义为:

A t t n ( Q , K , V ) = s o f t m a x ( Q K T d h ) V ∈ R N × d h \mathrm{Attn}(Q,K,V) = \mathrm{softmax} \left( \frac{QK^T}{\sqrt{d_h}} \right) V \in \mathbb{R}^{N\times d_h} Attn(Q,K,V)=softmax(dh QKT)VRN×dh

其中 softmax 按行进行。

这个公式可以这样理解:

  • Q K T QK^T QKT 计算 query 与 key 之间的相似度;
  • 除以 d h \sqrt{d_h} dh 是为了避免内积尺度过大;
  • softmax 将相似度转成权重;
  • 最后用这些权重对 value 做加权求和。

因此,attention 本质上是一个:

把一个向量序列映射到另一个向量序列的操作。

它的三个输入分别是:

  1. Queries;
  2. Keys;
  3. Values。

在不同场景中, Q , K , V Q,K,V Q,K,V 可以来自同一个序列,也可以来自不同序列。

DiTBlock 的作用是把图像 token、文本 token 和时间信息结合起来。

一个典型 DiTBlock 可以包含三类机制。

第一类是 self-attention,用于处理图像自身信息。

在图像 self-attention 中:

Q = image tokens , K = image tokens , V = image tokens Q=\text{image tokens}, \qquad K=\text{image tokens}, \qquad V=\text{image tokens} Q=image tokens,K=image tokens,V=image tokens

也就是说,图像 token 之间彼此交互,让模型能够建模图像不同区域之间的关系。

第二类是 cross-attention,用于引入文本 prompt 信息。

在文本 cross-attention 中:

Q = image tokens , K = text embeddings , V = text embeddings Q=\text{image tokens}, \qquad K=\text{text embeddings}, \qquad V=\text{text embeddings} Q=image tokens,K=text embeddings,V=text embeddings

这表示图像 token 会去 “关注” 文本 token,从而把 prompt 信息注入图像生成过程。

例如,prompt 中包含 “dog”、“grass”、“sunshine” 等语义时,图像 token 可以通过 cross-attention 获得这些语义条件,从而生成符合文本描述的内容。

第三类是 adaptive layer normalization,用于引入时间信息。

时间 t t t 通常不会像图像和文本一样直接作为 token 序列输入,而是通过自适应归一化影响网络。具体来说,layer normalization 后的缩放和偏移参数可以由时间 embedding 决定。

直观地说:

时间 embedding 会调节网络在不同噪声水平下的行为。

因为扩散模型在不同时间步面临的任务不同:早期可能需要处理强噪声,后期则需要恢复细节。因此时间条件必须影响整个网络的计算。

在实际模型中,自注意力、交叉注意力、前馈网络、时间调制等模块通常通过残差连接组合起来。

因此,一个 DiTBlock 可以粗略理解为:

1. 用 self-attention 处理图像 token 自身关系;

2. 用 cross-attention 将文本 prompt 信息注入图像 token;

3. 用 adaptive layer normalization 将时间信息注入网络;

4. 通过残差连接和 MLP 更新 token 表示。

当然,真实模型中会有很多工程细节,例如 multi-head attention、MLP block、LayerNorm 位置、残差缩放、条件注入方式等。不同模型会有不同变体。

但这里我们的核心目的是建立简单的整体图景:

DiT 是一种把图像 latent、时间和 prompt 都变成可处理向量表示,然后通过 Transformer block 学习生成向量场的架构。

正如课程中强调的,理解 Transformer 最好的方式不是只看公式,而是亲手实现。Lab 03 会进一步帮助理解这些模块。

8. Large-scale Diffusion Models

前面我们已经介绍了构建大规模图像生成模型所需的两个关键组件:

第一是 潜空间建模。通过 VAE 或类似自编码器,把高维像素图像压缩到更低维的 latent space 中,再在潜空间中训练生成模型。

第二是 神经网络架构。通过 DiT 或其变体,用 Transformer 来参数化扩散模型或流匹配模型中的向量场。

现在我们可以把这些组件放到真实大模型中看一看。

Stable Diffusion 3 是一个典型的现代文本到图像生成模型:

从这张 slide 中可以看到,Stable Diffusion 3 的整体设计几乎完整对应了前面几讲的内容。

首先,它使用的是 Flow Matching,并且采用类似 “straight line” 的调度方式,也就是 CondOT path。这类路径可以写成:

x t = t z + ( 1 − t ) ϵ x_t=t z+(1-t)\epsilon xt=tz+(1t)ϵ

其中:

  • ϵ \epsilon ϵ 表示初始噪声;
  • z z z 表示数据样本;
  • t = 0 t=0 t=0 时是噪声;
  • t = 1 t=1 t=1 时是数据。

对应的目标速度场是:

u t t a r g e t ( x ∣ z ) = z − ϵ u_t^{\mathrm{target}}(x|z)=z-\epsilon uttarget(xz)=zϵ

也就是说,模型要学习如何把噪声沿着一条相对直接的路径推向数据。

其次,Stable Diffusion 3 使用了 classifier-free guidance。在训练时,模型会以一定概率丢弃 prompt,让同一个模型同时学习有条件生成和无条件生成。采样时再通过 guidance scale 放大 prompt 相关的方向,从而增强图像与文本提示之间的匹配程度,CFG 的权重大约在 w ∈ [ 2.0 , 5.0 ] w\in[2.0,5.0] w[2.0,5.0] 这个范围内。

再次,Stable Diffusion 3 是在 latent space 中做 Flow Matching,而不是直接在像素空间中做。这一点非常关键。它使用预训练 VAE 把图像压缩为 latent,然后在 latent 上训练流匹配模型。采样得到 latent 后,再通过 VAE decoder 解码成图像。

所以它的整体流程可以写成:

image → latent → flow matching in latent space → generated latent → decoded image \text{image} \rightarrow \text{latent} \rightarrow \text{flow matching in latent space} \rightarrow \text{generated latent} \rightarrow \text{decoded image} imagelatentflow matching in latent spacegenerated latentdecoded image

此外,Stable Diffusion 3 的模型规模非常大,slide 中提到其参数量约为 8  billion 8\text{ billion} 8 billion ,采样时大约使用 50 50 50 个模拟步或采样步。

这也说明,即便有潜空间压缩,现代图像生成模型依然是非常庞大的系统工程。它不仅依赖算法本身,还依赖大规模数据、大规模模型和大规模训练基础设施。

Stable Diffusion 3 的核心网络架构是 MM-DiT,也就是 Multimodal Diffusion Transformer。

普通 DiT 的基本思想是:把图像或 latent image 切成 patch token,然后用 Transformer 处理这些 token,并通过时间 embedding 和条件信息调节网络。

Stable Diffusion 3 在这个基础上进一步扩展为多模态结构。

它的文本条件来自两类 text encoder:

  1. CLIP:提供较粗粒度的图文语义对齐信息;
  2. T5-XXL:提供更强的序列级语言理解能力。

也就是说,模型并不是只用一个简单的文本向量来表示 prompt,而是结合了不同类型的文本 embedding。

在 MM-DiT 中,文本 token 和图像 token 会在网络中被联合处理。相比传统的 “图像 token 通过 cross-attention 去读取文本 token” 的方式,MM-DiT 更强调让文本和图像在整个网络内部持续交互。

可以把它理解为:

Stable Diffusion 3 不只是把 prompt 当作一个额外条件,而是把文本和图像都放进多模态 Transformer 结构中共同建模。

这样做的目的,是让模型在生成过程中更充分地利用文本信息,尤其是在复杂 prompt、多对象关系、文字生成和空间布局等场景中提升表现。

接下来是视频生成模型的例子:Meta MovieGen。

MovieGen 和 Stable Diffusion 3 在整体思想上非常相似。

它同样使用:

  • Flow Matching;
  • straight-line scheduler / CondOT path;
  • classifier-free guidance;
  • latent space generation;
  • Diffusion Transformer 类架构。

但视频生成相比图像生成有一个额外难点:视频多了时间维度

这里要特别区分两种 “时间”。

第一种是生成模型中的时间 t t t ,也就是 Flow Matching 或 Diffusion 中从噪声到数据的连续时间变量。这个时间是生成过程的时间。

第二种是视频本身的时间维度,也就是 frame dimension。它表示视频中的第 1 帧、第 2 帧、第 3 帧等。

这两个时间不是同一个概念。

对于图像来说,latent 通常是 z ∈ R C × H × W z\in\mathbb{R}^{C\times H\times W} zRC×H×W ,而对于视频来说,latent 可能变成 z ∈ R C × T × H × W z\in\mathbb{R}^{C\times T\times H\times W} zRC×T×H×W ,其中 T T T 是视频帧数。因此,视频生成模型不仅要生成每一帧的空间内容,还要保证帧与帧之间的时间一致性。这包括:

  • 物体运动要连续;
  • 场景变化要合理;
  • 角色外观不能闪烁;
  • camera motion 要自然;
  • 长视频中前后语义要保持一致。

这使得视频生成比图像生成更困难。

从网络结构上看,如果图像模型做的是 2D patchify,那么视频模型往往需要扩展到 3D patchify,也就是同时在时间、高度和宽度三个方向上切 patch。

图像 patch 可以理解为 H × W H\times W H×W 上的小块,而视频 patch 则可以理解为 T × H × W T\times H\times W T×H×W 上的小块。因此,视频版 DiT 必须处理时空 token,而不只是空间 token。

MovieGen 的规模也非常大。slide 中提到它约有 30  billion 30\text{ billion} 30 billion 参数,并使用了 6144 6144 6144 块 H100 GPU 进行训练。这说明大规模视频生成不仅是算法问题,也是极其昂贵的工程和算力问题。

9. A Guide to the Diffusion Literature

最后,我们简要介绍下扩散模型文献中常见的几种不同表述方式。这些表述方式看起来不同,但在很多核心场景下,它们其实描述的是相似的问题。

扩散模型文献中首先容易混淆的是 时间约定

在这门课的 Flow Matching 约定中: t = 0 t=0 t=0 表示噪声分布 p i n i t p_{\mathrm{init}} pinit ,而 t = 1 t=1 t=1 表示数据分布 p d a t a p_{\mathrm{data}} pdata 。也就是说,采样过程是 noise → data \text{noise}\rightarrow \text{data} noisedata ,这可以称为 flow time convention

但在很多 score-based diffusion 或传统扩散模型文献中,时间约定通常相反: t = 0 t=0 t=0 表示数据,而当 t → ∞ t\rightarrow\infty t t t t 变大时,数据逐渐变成噪声。也就是说,forward diffusion 是 data → noise \text{data}\rightarrow \text{noise} datanoise ,而采样时再反过来 noise → data \text{noise}\rightarrow \text{data} noisedata ,这可以称为 diffusion time convention

此外,还有一些模型采用离散时间,而不是连续时间。例如 DDPM 通常使用离散 Markov chain 来描述加噪和去噪过程;DDIM 则可以看作与 probability flow ODE 有密切关系的确定性采样过程。

所以,当阅读扩散模型论文时,一定要先弄清楚作者采用的是哪种时间约定。否则很容易把 t = 0 t=0 t=0 t = 1 t=1 t=1 、噪声端和数据端弄反。

扩散模型的另一个核心问题是:

如何把数据逐渐破坏成噪声?

在这门课中,我们主要使用 probability path 的方式描述这个过程。

例如高斯概率路径:

p t ( x ∣ z ) = N ( α t z , β t 2 I d ) p_t(x|z) = \mathcal{N} \left( \alpha_t z, \beta_t^2 I_d \right) pt(xz)=N(αtz,βt2Id)

它表示:给定一个数据点 z z z ,在时间 t t t 上构造一个以 α t z \alpha_t z αtz 为均值、以 β t 2 I d \beta_t^2 I_d βt2Id 为协方差的高斯分布。

如果使用噪声采样形式,可以写成:

x t = α t z + β t ϵ , ϵ ∼ N ( 0 , I d ) x_t=\alpha_t z+\beta_t\epsilon, \qquad \epsilon\sim\mathcal{N}(0,I_d) xt=αtz+βtϵ,ϵN(0,Id)

除了 probability path,还有另外两种常见表述。

第一种是 stochastic interpolant。它直接定义一个插值函数,把噪声和数据连接起来:

I t ( ϵ , z ) I_t(\epsilon,z) It(ϵ,z)

第二种是 forward diffusion process。它通过一个 SDE 描述如何逐步给数据加噪:

d X t = a t ( X t ) d t + σ t d W t dX_t=a_t(X_t)dt+\sigma_t dW_t dXt=at(Xt)dt+σtdWt

这三种形式看起来不同:

  1. probability path;
  2. stochastic interpolant;
  3. forward diffusion SDE;

但在高斯设定下,它们通常可以相互转换,描述的是同一类从数据到噪声或从噪声到数据的路径构造方式。因此,不同论文可能使用不同语言,但核心问题往往是一样的:

如何构造一个连接噪声分布和数据分布的路径。

很多介绍扩散模型的文章都会展示一张经典图:forward SDE 把数据变成噪声,reverse SDE 再把噪声变回数据。

forward SDE 可以写成:

d x = f ( x , t ) d t + g ( t ) d w dx=f(x,t)dt+g(t)dw dx=f(x,t)dt+g(t)dw

它表示从数据出发,随着时间推进不断加入噪声,最终得到接近纯噪声的样本。

对应的 reverse SDE 可以写成:

d x = [ f ( x , t ) − g 2 ( t ) ∇ x log ⁡ p t ( x ) ] d t + g ( t ) d w ˉ dx= \left[ f(x,t) - g^2(t)\nabla_x\log p_t(x) \right]dt + g(t)d\bar{w} dx=[f(x,t)g2(t)xlogpt(x)]dt+g(t)dwˉ

其中 ∇ x log ⁡ p t ( x ) \nabla_x\log p_t(x) xlogpt(x) 就是 score function。

这个公式说明,如果我们知道每个时刻的 score function,就可以构造反向 SDE,从噪声一步步采样回数据。

这也是早期 score-based diffusion models 的核心视角:

学习 score function,然后用它构造反向去噪过程。

而在这门课中,我们已经看到,至少对于高斯概率路径,score function 和 vector field 之间可以互相转换。也就是说,学习 score 和学习向量场并不是完全割裂的两件事。在很多设定下,它们只是同一个生成动力学的两种参数化方式。

我们在课程中主要采用 Flow Matching 的视角,是因为它非常简洁。

Flow Matching 的优势在于:

第一,它可以只讨论 ODE 和向量场,不必一开始就引入完整的 SDE 理论。

第二,它的训练目标非常直接:学习一个速度场,使样本沿着概率路径从初始分布移动到目标数据分布。

第三,它允许我们从任意初始分布 p i n i t p_{\mathrm{init}} pinit 变换到任意目标分布 p d a t a p_{\mathrm{data}} pdata ,也就是说,初始分布不一定非要是标准高斯。

在前面的例子中,我们经常使用高斯噪声,是因为它简单、方便、易于分析。但 Flow Matching 框架本身并不要求初始分布必须是高斯。

因此,Flow Matching 可以被理解为一种更一般的分布桥接方法:

p i n i t → p d a t a p_{\mathrm{init}} \rightarrow p_{\mathrm{data}} pinitpdata

Flow Matching 的一个重要特点是,它可以用于连接任意分布。

这不仅适用于从纯噪声生成图像,也适用于很多更一般的任务,例如:

  • 无音频视频 → \rightarrow 有音频视频;
  • 低分辨率图像 → \rightarrow 高分辨率图像;
  • 未扰动细胞状态 → \rightarrow 扰动后细胞状态;
  • 局部缺失图像 → \rightarrow 完整图像;
  • 草图 → \rightarrow 真实图像;
  • 某种模态 → \rightarrow 另一种模态。

从这个角度看,生成模型不只是 “从噪声生成图片”,而是可以被理解为:

学习一个从源分布到目标分布的动态变换过程。

这也是 Flow Matching 框架非常吸引人的地方。它把生成建模、图像修复、超分辨率、条件生成、多模态转换等任务都放在了同一个数学框架下。


Q & A

Q:无分类器引导(CFG)中丢弃 Prompt 的概率如何选择?

在 CFG 训练中,模型会以一定概率把条件 y y y 替换为空条件 ∅ \varnothing 。这样模型既能学习 u t θ ( x ∣ y ) u_t^\theta(x|y) utθ(xy) 也能学习 u t θ ( x ∣ ∅ ) u_t^\theta(x|\varnothing) utθ(x) 。采样时再构造加权向量场:

u t θ , w ( x ∣ y ) = ( 1 − w ) u t θ ( x ∣ ∅ ) + w u t θ ( x ∣ y ) u_t^{\theta,w}(x|y) = (1-w)u_t^\theta(x|\varnothing) + w u_t^\theta(x|y) utθ,w(xy)=(1w)utθ(x)+wutθ(xy)

如果训练时总是丢弃 prompt,那么模型几乎学不到有条件生成。如果训练时从不丢弃 prompt,那么模型又学不到无条件分支 u t θ ( x ∣ ∅ ) u_t^\theta(x|\varnothing) utθ(x) ,这样采样时 CFG 就无法正常工作。

因此,丢弃概率需要在两者之间折中。经验上,很多模型会选择大约 20% 的概率丢弃条件。这不是严格理论推导出的最优值,而是一个实践中常用且有效的经验设置。


下次课会有一节关于离散扩散模型的补充内容。我们平时使用的 ChatGPT 等大语言模型,其实也能基于扩散模型的原理来构建,这类方法就叫离散扩散模型,届时我们会详细讲解。

结语

本讲围绕 Latent Spaces 与 Neural Network Architectures,系统性地补全了现代大规模生成模型从理论到工程落地的关键拼图 。

首先,在 latent space 部分,我们明确了高维像素空间在生成建模中的根本困难:维度过高、计算成本极大、冗余信息严重,使得直接学习向量场 u t θ : R d → R d u_t^\theta:\mathbb{R}^d \to \mathbb{R}^d utθ:RdRd 在工程上几乎不可扩展。因此,引入 VAE / autoencoder 的核心目的,不只是压缩数据,而是将生成建模问题转移到一个更低维、更结构化的潜空间中。

随后,我们分析了标准 autoencoder 与 VAE 的差异:Autoencoder 关注重建能力,但不约束 latent 分布结构;VAE 通过 KL 正则项,将潜变量分布约束到标准高斯附近,从而使 latent space 更 “规整”、更适合后续 flow / diffusion modeling

接着,在 VAE reparameterization trick 中,我们将随机采样转化为可微路径: z = μ ϕ ( x ) + σ ϕ ( x ) ⊙ ϵ z = \mu_\phi(x) + \sigma_\phi(x)\odot \epsilon z=μϕ(x)+σϕ(x)ϵ ,从而使得整个生成过程可以在标准反向传播框架中稳定训练。这一步本质上将 “随机变量依赖参数” 的问题转化为 “确定性函数 + 外部噪声”,是现代生成模型训练的基础技巧之一。

最后,在 Neural Network Architectures 部分,我们讨论了如何用 Transformer 实现生成模型的参数化,从而得到 DiT(Diffusion Transformer)范式。

下一讲我们将重点讨论离散扩散模型相关的内容,敬请期待🤗

参考

更多推荐