流匹配与扩散模型-06：变分自编码器

自编码器

直接在像素空间中训练生成模型的成本会随分辨率迅速增长。以 $1024 \times 1024$ 的 RGB 图像为例，$x \in \mathbb{R}^d$ 的维度为 $d=3\times1024\times1024\approx3\times10^6$；视频还要额外乘上帧数。相邻像素通常高度相关，天空中的大片蓝色和物体内部的重复纹理占据了大量坐标，却没有提供同等数量的独立语义。流模型和扩散模型仍要在每个采样步输出与输入同维的向量场、噪声或速度，把算力反复花在这些冗余坐标上，显存与计算压力因而尤为明显。

自编码器利用“自然图像集中在较低维流形附近”这一假设，把像素中的重复信息压成更少的潜在坐标。确定性编码器和解码器分别写作

$$\mu_\phi:\mathbb{R}^d\rightarrow\mathbb{R}^k,\qquad \mu_\theta:\mathbb{R}^k\rightarrow\mathbb{R}^d,\qquad k\ll d.$$

给定 $x$，编码得到 $z=\mu_\phi(x)$，再由 $\hat{x}=\mu_\theta(z)$ 重构。由于瓶颈 $k\ll d$ 无法逐像素照抄输入，网络必须优先保留有助于解码的结构。最基本的训练目标是

$$\mathcal{L}_{\mathrm{Recon}}(\phi,\theta) =\mathbb{E}_{x\sim p_{\mathrm{data}}} \left[\left\lVert\mu_\theta(\mu_\phi(x))-x\right\rVert_2^2\right].$$

该损失只检查“每个训练样本能否绕一圈回来”，并不关心编码之间的空白区域。编码器完全可以把相似图像送到相距很远的点，或把训练样本挤成许多彼此分离的细簇；解码器只在这些点附近受过训练，从簇间采样或插值便可能得到无意义结果。若潜在空间还要承担生成任务，除可重构之外，还需要相邻编码对应相近数据，并让一个简单分布能够覆盖常用区域。VAE 用概率编码和先验约束同时塑造这种结构。

VAE 损失

VAE 不再把一张图压到一个没有体积的点，而是编码成一个带均值和尺度的局部概率云。相近样本的概率云可以重叠，解码器也会在每个样本周围而非单点上接受训练，从而降低潜在空间对微小扰动的敏感性。常见的对角高斯参数化为

$$q_\phi(z\mid x) =\mathcal{N}\!\left(z;\mu_\phi(x), \operatorname{diag}(\sigma_\phi^2(x))\right),$$

$$p_\theta(x\mid z) =\mathcal{N}\!\left(x;\mu_\theta(z), \sigma_\theta^2(z)I_d\right).$$

编码器输出均值和逐维方差，表示“给定 $x$ 后哪些潜在解释仍然合理”；解码器给出潜在变量下的数据分布。方差很小会接近确定性编码并保存更多样本细节，方差较大则让邻域更平滑，却要求同一个区域能解码出更多变化。编码和解码可写为

$$z\sim q_\phi(\cdot\mid x),\qquad \tilde{x}\sim p_\theta(\cdot\mid z).$$

为了判断采到的 $z$ 是否仍保留了足够信息，重构项采用负条件对数似然：

$$\mathcal{L}_{\mathrm{rec}}(\phi,\theta) =-\mathbb{E}_{x\sim p_{\mathrm{data}},\,z\sim q_\phi(\cdot\mid x)} \left[\log p_\theta(x\mid z)\right].$$

若解码器是上式的各向同性高斯，则忽略与参数无关的常数后，

$$\mathcal{L}_{\mathrm{rec}} =\mathbb{E}_{x,z} \left[ \frac{\lVert x-\mu_\theta(z)\rVert_2^2} {2\sigma_\theta^2(z)} +\frac d2\log\sigma_\theta^2(z) \right]+\mathrm{const}.$$

第一项用预测方差衡量误差，第二项则惩罚随意增大方差。缺少后者时，解码器可以声称“任何像素都很不确定”来降低第一项，而不必改善均值；两项共同构成合法高斯似然。实践中常固定解码方差为常数 $\tilde{\sigma}^2$，把这种不确定性取舍预先确定，此时重构项与随机编码下的 MSE 成固定比例：

$$\mathcal{L}_{\mathrm{rec}} =\mathbb{E}_{x,z} \left[ \frac{1}{2\tilde{\sigma}^2} \lVert x-\mu_\theta(z)\rVert_2^2 \right]+\mathrm{const}.$$

仅靠随机编码仍可能把概率云散落在任意位置，无法回答生成时应从哪里取 $z$。于是预先选择易采样的 $p(z)=\mathcal{N}(0,I_k)$，并要求每个样本的编码分布不要偏离它太远：

$$\mathcal{L}_{\mathrm{prior}}(\phi) =\mathbb{E}_{x\sim p_{\mathrm{data}}} \left[ D_{\mathrm{KL}}\!\left(q_\phi(z\mid x)\,\Vert\,p(z)\right) \right].$$

对两个对角高斯

$$q=\mathcal{N}(\mu_q,\operatorname{diag}(\sigma_q^2)),\qquad p=\mathcal{N}(\mu_p,\operatorname{diag}(\sigma_p^2)),$$

KL 散度具有闭式解：

$$D_{\mathrm{KL}}(q\Vert p) =\frac12\sum_{j=1}^k \left[ \log\frac{\sigma_{p,j}^2}{\sigma_{q,j}^2} +\frac{\sigma_{q,j}^2+(\mu_{q,j}-\mu_{p,j})^2} {\sigma_{p,j}^2} -1 \right].$$

代入标准正态先验，得到训练中常用的形式：

$$D_{\mathrm{KL}}\!\left( q_\phi(z\mid x)\,\Vert\,\mathcal{N}(0,I_k) \right) =\frac12\sum_{j=1}^k \left( \mu_{\phi,j}^2(x) +\sigma_{\phi,j}^2(x) -\log\sigma_{\phi,j}^2(x) -1 \right).$$

其中 $\mu_{\phi,j}^2$ 把概率云中心拉回原点，$\sigma_{\phi,j}^2-\log\sigma_{\phi,j}^2-1$ 在方差过大或过小时都会增大，并在 $\sigma_{\phi,j}^2=1$ 时取最小值。因此，KL 的作用并非简单地把所有编码压成一个点，而是让它们在共享的标准尺度内占据连续区域。重构项希望 $z$ 保留足够多的个体信息，KL 项则希望编码容易从共同先验覆盖，两股力量形成 $\beta$-VAE 的目标：

$$\mathcal{L}_{\beta}(\phi,\theta) =\mathcal{L}_{\mathrm{rec}}(\phi,\theta) +\beta\mathcal{L}_{\mathrm{prior}}(\phi), \qquad \beta\ge0.$$

$\beta=0$ 时模型退回带随机性的自编码器，潜在分布可能难以采样；增大 $\beta$ 会提高不同样本编码的重叠程度并限制信息容量，通常以部分重构细节为代价。极端情况下，解码器可以只凭自身能力建模数据，编码器便把所有 $x$ 都映射到近似先验，出现 posterior collapse，潜变量几乎不再携带输入信息。需要注意，逐样本 KL 只鼓励聚合潜在分布接近先验，并不保证二者严格相等，这也为后续学习更复杂的潜在先验留下了空间。

ELBO

VAE 的损失并非任意拼接“重构”和“正则”两个经验项，它来自对数据似然的可计算近似。生成模型真正希望最大化

$$\log p_\theta(x) =\log\int p_\theta(x\mid z)p(z)\,\mathrm dz,$$

但积分汇总了所有可能的 $z$，高维神经解码器下通常无法精确计算。将编码器 $q_\phi(z\mid x)$ 作为重要性分布插入积分，并用 Jensen 不等式，可得

$$\begin{aligned} \log p_\theta(x) &=\log\mathbb E_{q_\phi(z\mid x)} \left[ \frac{p_\theta(x\mid z)p(z)}{q_\phi(z\mid x)} \right]\\ &\ge \mathbb E_{q_\phi(z\mid x)}[\log p_\theta(x\mid z)] -D_{\mathrm{KL}}\!\left(q_\phi(z\mid x)\Vert p(z)\right). \end{aligned}$$

右侧就是证据下界：第一项奖励能够解释 $x$ 的潜变量，第二项防止编码器只在先验几乎采不到的位置解释它。VAE 同时定义的两条联合分布为

$$q_\phi(x,z)=p_{\mathrm{data}}(x)q_\phi(z\mid x), \qquad p_\theta(x,z)=p(z)p_\theta(x\mid z).$$

在 $\beta=1$ 时，联合 KL 与 VAE 损失存在精确关系：

$$\begin{aligned} D_{\mathrm{KL}}\!\left(q_\phi(x,z)\,\Vert\,p_\theta(x,z)\right) &= \mathbb{E}_{x\sim p_{\mathrm{data}}} \left[\log p_{\mathrm{data}}(x)\right] \\ &\quad -\mathbb{E}_{x,z} \left[\log p_\theta(x\mid z)\right] +\mathbb{E}_{x\sim p_{\mathrm{data}}} \left[ D_{\mathrm{KL}}\!\left(q_\phi(z\mid x)\Vert p(z)\right) \right] \\ &= \mathbb{E}_{p_{\mathrm{data}}}[\log p_{\mathrm{data}}(x)] +\mathcal{L}_{\beta=1}(\phi,\theta). \end{aligned}$$

第一项只由数据分布决定，因此降低标准 VAE 损失，就是让“真实样本经过编码得到的 $(x,z)$”与“先从先验取 $z$ 再解码得到的 $(x,z)$”难以区分。这个等价关系不适用于一般的 $\beta\ne1$；后者主动改变重构与容量的平衡，不再是同一生成模型的标准似然下界。

对单个样本，定义

$$\operatorname{ELBO}(x;\phi,\theta) =\mathbb{E}_{q_\phi(z\mid x)}[\log p_\theta(x\mid z)] -D_{\mathrm{KL}}\!\left(q_\phi(z\mid x)\Vert p(z)\right).$$

它与真实对数似然之间满足恒等式

$$\log p_\theta(x) =\operatorname{ELBO}(x;\phi,\theta) +D_{\mathrm{KL}}\!\left( q_\phi(z\mid x)\,\Vert\,p_\theta(z\mid x) \right).$$

KL 非负，所以 ELBO 永远不会超过 $\log p_\theta(x)$；当编码器恰好等于真实后验时，最后一项为零，下界变紧。最大化 ELBO 一方面改善生成模型，另一方面也让可计算的编码器追赶生成模型的真实后验。当且仅当 $\beta=1$ 时，

$$\mathcal{L}_{\beta=1} =-\mathbb{E}_{x\sim p_{\mathrm{data}}} \left[\operatorname{ELBO}(x;\phi,\theta)\right].$$

这里的 $\mathcal L_{\beta=1}$ 指对数据分布取期望后的标准 VAE 目标。

“变分”指的是在一个可计算的分布族中寻找 $q_\phi(z\mid x)$，近似通常不可计算的真实后验 $p_\theta(z\mid x)$。即使训练充分，下界也可能因两个不同限制而无法贴紧。对给定的模型参数，设

$$q_x^\star =\arg\max_{q\in\mathcal{Q}} \operatorname{ELBO}(x;q,\theta),$$

其中 $\mathcal{Q}$ 是允许的近似后验族，例如对角高斯族。则选定编码器的推断误差可拆为

$$\begin{aligned} \log p_\theta(x) -\operatorname{ELBO}(x;q_\phi,\theta) &= \left[ \log p_\theta(x) -\operatorname{ELBO}(x;q_x^\star,\theta) \right] \\ &\quad+ \left[ \operatorname{ELBO}(x;q_x^\star,\theta) -\operatorname{ELBO}(x;q_\phi,\theta) \right]. \end{aligned}$$

前一项是近似后验族表达能力有限带来的近似缺口，例如对角高斯无法表示多个分离的后验模态；后一项是一个共享编码器必须立即为所有输入给出答案，无法像逐样本优化那样到达各自最优解而产生的摊销缺口。二者都描述“用 $q$ 逼近真实后验”的误差，而聚合潜在分布 $q_\phi(z)=\int p_{data}(x)q_\phi(z\mid x)\,\mathrm dx$ 与先验的失配描述“从哪里采样”的问题，含义并不相同。

重参数化

重构损失依赖从 $q_\phi(z\mid x)$ 取出的随机样本。如果把“从一个由 $\phi$ 决定的分布抽签”当作黑盒，解码器的误差无法沿这次离散抽签直接传回均值和方差。重参数化把抽签改写成固定噪声经过可导变换：

$$\epsilon\sim\mathcal{N}(0,I_k),\qquad z=\mu_\phi(x)+\sigma_\phi(x)\odot\epsilon.$$

这样，$\epsilon$ 是与参数无关的外部随机源；在一次前向—反向传播中把它视作常量，$z$ 就是关于 $\mu_\phi(x)$ 和 $\sigma_\phi(x)$ 的普通可导函数，重构误差能够沿解码器、$z$ 一直传回编码器。这种路径梯度通常比只根据采样结果估计分布参数的梯度方差更低。

一个 mini-batch 的训练过程由同一条可导计算图串起：

1. 编码器输出每个样本的 $\mu_i$ 与 $\log\sigma_i^2$，再计算 $\sigma_i=\exp(\tfrac12\log\sigma_i^2)$。
2. 采样 $\epsilon_i\sim\mathcal{N}(0,I_k)$，并构造 $z_i=\mu_i+\sigma_i\odot\epsilon_i$。
3. 解码器输出 $\hat{x}_i=\mu_\theta(z_i)$。
4. 计算高斯负对数似然形式的重构损失，以及到 $\mathcal{N}(0,I_k)$ 的 KL 损失。
5. 按 $\mathcal{L}_\beta=\mathcal{L}_{\mathrm{rec}}+\beta\mathcal{L}_{\mathrm{prior}}$ 反向传播。

编码器输出 $\log\sigma^2$ 而非直接输出方差，可以避免训练过程中产生负值，并把很大的尺度范围变成更易优化的实数范围。固定解码方差通常也比直接预测像素级方差稳定。若图像重构只使用逐像素 MSE，多个同样合理的高频细节会被均值化，结果容易发糊；感知损失或对抗损失可让解码器更重视人眼可见的结构，但也改变了纯高斯似然的解释，应视作额外的工程目标。KL 预热在训练早期让解码器先学会使用 $z$，随后再逐步加强先验约束，可降低模型一开始就选择忽略潜变量的风险。

潜在扩散

VAE 本身可以从 $p(z)=\mathcal{N}(0,I_k)$ 采样并解码，但强迫复杂图像的聚合潜在分布严格服从简单高斯，往往需要牺牲重构；放松 KL 又会使直接从高斯采到的 $z$ 落在编码器很少使用的区域。潜在扩散不要求一个模型同时解决“如何压缩像素”和“哪些潜变量最常出现”，而是把两件事分成两个阶段优化。它针对的是聚合潜在分布的建模和像素空间的计算成本，并非专门修复摊销缺口。

第一阶段训练自编码器或 VAE，使编码器把 $x\sim p_{\mathrm{data}}$ 映射到空间尺寸更小、通道更紧凑但仍保留感知与语义信息的潜变量，并使解码器能够可靠重构。此阶段相当于为生成模型选择工作坐标：压缩越强，第二阶段越便宜，但被编码器抹去的文字、细线和纹理以后无法靠扩散模型找回。编码器与解码器通常在第二阶段冻结，使潜在数据分布保持固定，否则生成模型会追逐不断移动的目标。

第二阶段将训练样本编码为

$$z\sim q_\phi(z\mid x),\qquad x\sim p_{\mathrm{data}},$$

并在这些潜变量上训练扩散或流模型 $p_\psi(z)$；有条件生成时，目标相应变为 $p_\psi(z\mid c)$。该模型直接学习编码器实际产生的复杂分布，相当于用数据驱动的潜在先验替代过于简单的标准高斯。扩散过程内部仍可从高斯噪声开始，但它负责把噪声逐步变成聚合潜变量，不要求第一阶段已经把二者完全对齐。潜变量的数值尺度常需固定归一化，使预设噪声调度与其方差处在相近量级。

推理时先从训练好的潜在扩散或流模型采样

$$z\sim p_\psi(z\mid c),$$

再使用解码器的均值输出

$$\hat{x}=\mu_\theta(z).$$

取均值而不再从 $p_\theta(x\mid z)$ 随机采样，可避免在已经完成的生成结果上再叠加一次独立解码噪声。训练第二阶段时，真实图像提供目标 $z$，可以直接随机加噪并监督单个时刻；推理时没有目标 $z$，必须从噪声运行完整的流或扩散采样，再解码一次。主要计算由数百万像素转移到更小的潜在张量，速度和显存显著改善，但整个系统也形成两道上限：潜在生成模型必须覆盖编码器产生的分布，第一阶段丢掉的细节则无法在第二阶段恢复。