VAE 和 VQVAE

python pytorch

本篇博客重点讲解一下 VAE 和 VQVAE 的原理，VAE 是 Variational Autoencoder 的缩写，中文名为变分自编码器，是一种生成模型。VQVAE 则是 Vector Quantized Variational Autoencoder 的缩写，中文名为向量量化变分自编码器，是对 VAE 的一种改进，这种改进也是从连续空间到离散空间的一种尝试，下面分别介绍这两种模型

1. 什么是生成模型？

生成模型的目标是学习数据点 $X$ 的真实分布 $P_{\text{gt}}(X)$（这里的数据点 $X$ 可以理解为训练集所在的空间分布），并构建一个模型 $M$，通过从随机噪声中采样变量（通常是高斯噪声）来生成新的数据，使得生成数据的分布 $P_M(X)$ 尽可能接近 $P_{\text{gt}}(X)$。换句话说，我们希望模型能够生成像 $X$ 数据集中的样本一样的样本，从而完成对于数据真实分布的建模。

生成模型的核心思想：通过定义一个潜在空间（通常用潜在变量 $z$ 表示），通过模型将潜在变量映射到 $X$ 的空间，从而获得可生成数据的分布。生成模型的最终目标是最大化数据点 $X$ 的概率，使得模型分布 $P(X)$ 接近真实数据分布 $P_{\text{gt}}(X)$。这里我们假设它假设 数据 $X$ 并不是直接生成，而是先从一个潜在空间中采样一个隐藏的变量 $z$，是一个潜在的低维结构，代表数据 $X$ 的压缩信息或者隐藏特征。比如，在图像生成中，$z$ 可以表示对象的类型、颜色、背景等信息。可以将 $P(X)$ 分解为一个公式：

$$
P(X) = \int P(X|z; \theta) P(z) dz
$$

• $\theta$：模型的参数，定义了解码器（或者生成器）。
• $z$：随机噪声，通常假设遵循某种简单的分布（例如标准正态分布）。
• $P(z)$：先验分布，用来定义 $z$ 的采样边界。
• $P(X|z; \theta)$：条件概率分布，表示在给定潜在变量 $z$ 及模型的参数 $\theta$ 时，生成数据点 $X$ 的概率。

这一公式的意义在于，它描述了数据点 $X$ 的生成过程：通过从潜在空间中采样 $z$，生成模型 $P(X|z; \theta)$ 中会输出分布 $X$。最终，我们希望调整模型参数 $\theta$，使得 $P(X)$ 最大化。

为什么引入潜在变量 $z$？

从数学上看，直接建模 $P(X)$ 是非常困难的，尤其当 $X$ 是高维数据（如图像、音频、文本等）时。考虑如下例子：对于一张 $256 \times 256$ 的图像，每个像素可以取值 $0-255$。那这张图像可能的总组合数高达：$256^{256 \times 256}$，直接推断 $P(X)$ 是几乎不可能的。因此，引入潜在变量 $z$ 是为了简化问题：

1. 降低复杂度：$z$ 是在一个潜在空间中定义的变量，通常是低维的（比如 10 维或 100 维）。这比直接操作高维数据 $X$ 要简单得多。
2. 让生成过程可控：$z$ 是对数据 $X$ 的压缩或者抽象表达，它捕获了数据的关键特征。生成模型通过学习 $P(X|z)$ 的方式，可以灵活控制生成过程的输出。

2. 为什么需要变分推断？

在生成模型中，我们的目标是最大化 $P(X)$。然而，直接优化 $P(X)$ 是非常困难的，原因如下

1. 潜在空间太广泛：我们通过 $P(X) = \int P(X|z) P(z) , dz$ 定义了生成数据 $X$ 的边际概率，但这一公式需要对所有可能的潜在变量 $z$ 求积分。因为 $P(z)$ 通常是一个非常宽泛的分布（如标准高斯分布），它覆盖了整个潜在空间，计算代价非常高。
2. 大多数 $z$ 无效：在 $P(z)$ 的采样中，大部分潜在变量 $z$ 生成数据的概率 $P(X|z)$ 通常非常小或接近 0，这意味着这些 $z$ 对生成数据没有帮助。这让采样过程不但慢，而且效率极低。

为了解决这些问题，我们需要一个更有效的方式来缩小采样范围。因此，引入了新的分布 $Q(z|X)$ 来替代 $P(z)$ 进行采样。$Q(z|X)$ 是一个近似后验分布，它描述了给定数据 $X$，可能的潜在变量 $z$ 的分布，使采样只集中在和 $X$ 相关的 $z$ 范围内，而不是盲目搜索整个潜在空间。相比于直接从 $P(z)$ 中无脑采样，$Q(z|X)$ 更加智能。例如，$Q(z|X)$ 专注于哪些 $z$ 会生成与数据 $X$ 更接近的样本，从而显著提高了生成模型的效率。

变分推断的核心公式，通过 $Q(z|X)$，我们可以优化以下新的期望目标：
$$
\mathbb{E}_{z \sim Q} P(X|z) = \int P(X|z) Q(z) dz
$$

这一公式的核心思想是：我们仍然想优化生成概率 $P(X)$，但将重点从宽泛的 $P(z)$ 分布转移到了与数据 $X$ 更相关的 $Q(z|X)$ 分布上（即：用 $Q(z|X)$ 替代 $P(z)$ 进行积分，可以理解为这样就更相关了嘛）

这里再强调一下 $Q(z|X)$ 的重要性

1. 缩小了潜在空间的范围：$Q(z|X)$ 专注于和 $X$ 相关的潜在变量 $z$，避免了无效采样，从而显著减少了计算需求。
2. 生成过程更加高效：$Q(z|X)$ 使得生成模型可以快速找到有效的 $z$，从而生成更符合 $X$ 的数据。
3. 优化模型的方法：通过不断优化 $Q(z|X)$，它可以越来越接近真实后验分布 $P(z|X)$，从而间接优化了模型的总体目标 $P(X)$。

实际上，变分自编码器（VAE） 的核心正是借助变分推断：它通过优化一个近似分布 $Q(z|X)$ 代替难解的后验分布 $P(z|X)$，让模型可以优化

3. 从 KL 散度到 VAE 的目标函数

从上面的推理中，我们可以用 $Q(z|X)$ 来替代 $P(z|X)$，作者将它设置为一个概率分布，并使用KL 散度来进行衡量

$$
D_{\mathrm{KL}}[Q(z|X) | P(z|X)] = \mathbb{E}_{z \sim Q} \big[ \log Q(z|X) - \log P(z|X) \big]
$$

这一目标的核心是让 $Q(z|X)$ 尽可能接近 $P(z|X)$。然后根据 贝叶斯公式 将 $P(z|X) = \frac{P(z)P(X|z)}{P(X)}$ 带入，再进一步化简如下：

$$
\begin{aligned}
D_{\mathrm{KL}}[Q(z|X) | P(z|X)] &= \mathbb{E}_{z \sim Q} \big[ \log Q(z|X) - \log \frac{P(z) P(X|z)}{P(X)} \big] \\
&= \mathbb{E}_{z \sim Q} \big[ \log Q(z|X) - \log P(z) - \log P(X|z) + \log P(X) \big] \\
&= \mathbb{E}_{z \sim Q} \big[ \log Q(z|X) - \log P(z) - \log P(X|z) \big] + \log P(X).
\end{aligned}
$$

由于 $\log P(X)$ 与 $z$ 无关，它可视为一个常数。因此通过重新整理上式，我们得到：

$$
\log P(X) - D_{\mathrm{KL}}[Q(z|X) | P(z|X)] = \mathbb{E}_{z \sim Q} \big[ \log P(X|z) \big] - D_{\mathrm{KL}}[Q(z|X) | P(z)].
$$

这就是 VAE 最核心的公式。

• 左边：$\log P(X)$ 是我们想最大化的项，描述了模型生成真实数据 $X$ 的概率；通过最小化 $D_{\mathrm{KL}}[Q(z|X) | P(z|X)]$，希望 $Q(z|X)$ 逼近 $P(z|X)$。

• 右边：这部分公式可以通过梯度下降实现优化。

  • 第一项 $\mathbb{E}_{z \sim Q} \big[ \log P(X|z) \big]$：这是生成数据的重构概率，可以通过解码器直接实现，类似于自编码器（AE）。
  • 第二项 $D_{\mathrm{KL}}[Q(z|X) | P(z)]$：这是一个正则化项，约束潜在空间中的 $Q(z|X)$ 使其尽可能贴近真实的先验 $P(z)$。

以上推导清楚展现了 VAE 的本质：它通过变分推断将复杂的采样问题转化为更可控的优化问题，同时结合解码和正则化，使生成模型既能生成高质量数据，又能调整潜在空间的结构。

4. 从 VAE 到 VQVAE：离散潜在空间的探索

在上一部分，我们探讨了 VAE 的工作原理和公式推导，并了解到 VAE 使用了潜在变量 $z$ 的连续分布（通常是高斯分布），以此来高效生成数据。然而，对于某些生成任务（如离散自然语言建模、图像压缩、视频生成等），这种连续分布可能并不总是最优选择。相较于 VAE，VQVAE 提出了一种有趣的（非概率性的）离散潜在空间表示方法，对于图像、音频等高维数据的离散特征建模非常有用。在 Diffusion、Transformer 等模型大热之前，VQVAE 已经被证明在生成任务中非常有效。以下是经过精心设计的博客内容：

5. VQVAE 是什么？

VQVAE，全称为 Vector Quantized Variational Autoencoder（向量量化变分自编码器），是由 Aaron van den Oord 等人在 2017 年提出的。（论文《Neural Discrete Representation Learning》）

其核心思想是用离散空间（Discrete Space）替代了 VAE 中的连续潜在空间。具体来说，VQVAE 强制将潜在变量嵌入到一个 离散的向量空间 中，并通过向量量化操作将连续表示 $z$ 映射到离散的编码向量索引。它的离散化特性使其在一些任务上相比于 VAE 具有以下两点优势：

1. 离散化特性适合处理本质上离散的数据（如文本、编码、符号等）。
2. 简化了统计建模任务的复杂性，比如在生成阶段仅需要对有限数量的离散向量进行索引和采样，而不是连续分布的采样。

VQVAE 的这种离散化方式，让它更适合作为下游模型（例如 GAN、Transformer 模型）的离散层表示。

7. VQVAE 的基本结构

VQVAE 的高层结构与标准自编码器（Autoencoder）非常相似，仍然包括一个 编码器 和一个 解码器，但在潜在表示的处理上引入了新的离散化步骤。以下是 VQVAE 的基本结构流程：

编码器 (Encoder)

输入数据 $X$ 经过一个编码器网络 $E$，被映射到潜在空间中产生一个连续的潜在表示 $z_e(X)$：
$$
z_e(X) = E(X),
$$
其中 $z_e(X)$ 是编码器对输入数据 $X$ 的连续表示。

向量量化 (Vector Quantization)

与 VAE 不同，VQVAE 不直接使用 $z_e(X)$ 作为潜在变量。相反，它需要对 $z_e(X)$ 进行量化，使其映射到一个 离散的向量空间。

这一步通过一个预定义的向量集合 ${e_k}_{k=1}^K$ 实现。向量集合中包含 $K$ 个固定维度的向量（即 “代码簇”）。将 $z_e(X)$ 的每个编码向量 $z$ 用集合中最接近的向量 $e_k$ 替换：
$$
z_q(X) = e_k \quad \text{where } k = \arg\min_j | z_e(X) - e_j |_2
$$
也就是说，量化的过程是寻找 $z_e(X)$ 最近邻的离散向量 $e_k$。（注意，最近邻搜索是基于欧几里得距离实现的）

此时，连续潜在变量被离散化为离散的向量索引（通过 “代码簇” 表示）。

解码器 (Decoder)

量化后的 $z_q(X)$ 被送入解码器 $D$，重建出输入数据 $\hat{X}$：
$$
\hat{X} = D(z_q(X)).
$$

3. VQVAE 的损失函数

在训练过程中，VQVAE 的损失函数由两部分组成：

1. 重构误差
   重构误差是解码器输出 $\hat{X}$ 与输入 $X$ 的差异，它通过常见的 $L_2$ 范数来衡量：
   $$
   \mathcal{L}_{\text{recon}} = | X - \hat{X} |^2.
   $$

2. 向量量化误差
   向量量化过程中，编码器输出 $z_e(X)$ 和被量化后的向量 $z_q(X)$ 应该彼此接近。为了减少编码器和离散表示之间的不连续性，VQVAE 作者引入了一个量化误差项：
   $$
   \mathcal{L}_{\text{vq}} = | \text{sg}[z_e(X)] - z_q(X) |^2,
   $$
   其中 $\text{sg}$ 是 “停止梯度” 操作，表示这一项不会对 $z_q(X)$ 传播梯度（因为 $z_q$ 是离散值）。

此外，为了让编码器 $z_e(X)$ 学会更好地输出接近量化空间的特征，还需要添加一个 编码约束项：
$$
\mathcal{L}_{\text{commit}} = \beta | z_e(X) - \text{sg}[z_q(X)] |^2,
$$
其中 $\beta$ 是一个超参数，用于调节编码约束的权重。

最终，VQVAE 的总损失函数为：
$$
\mathcal{L}{\text{VQVAE}} = \mathcal{L}{\text{recon}} + \mathcal{L}{\text{vq}} + \mathcal{L}{\text{commit}}.
$$

4. VQVAE 的优势与应用

为什么引入离散潜在空间？

相比于连续潜在变量，离散潜在变量具有以下几点优势：

1. 离散空间更适合某些离散序列任务，例如语言建模、离散编码符号等。
2. 离散潜在空间的表示更加紧凑，有助于降低数据生成时的复杂度。
3. 后续的生成模型（例如基于 Transformer 的建模）更适合直接处理这些离散潜在表示。

VQVAE 的应用

VQVAE 在以下应用领域都取得了巨大的成功：

1. 图像生成：VQVAE 可以将图像压缩为离散的潜在空间，然后使用 PixelCNN 或 Transformers 模型对离散编码序列进行建模，从而实现高质量的图像生成。
2. 音频信号处理：VQVAE 曾被应用于语音生成任务（如 WaveNet），有效地提高了音频数据的生成质量。
3. 视频生成：VQVAE 可将视频的每帧压缩为离散表示，再通过序列建模生成视频。
4. 模型压缩：通过离散化潜在空间，VQVAE 可以被用于高效的神经网络权重压缩。

5. 总结：从 VAE 到 VQVAE 的转变

相比于 VAE，VQVAE 的本质转变在于离散化潜在变量，这为多种离散任务的生成建模提供了极大的便利。VAE 假设潜在变量服从连续分布，而 VQVAE 将其映射到一个固定的向量库，从而提高了建模效率，同时保留了生成模型的灵活性。

VQVAE 为后续的生成模型（如 VQVAE-GAN、VQVAE-Transformer）开辟了新的方向。在现代生成任务中，VQVAE 凭借其离散特性和强大的重构能力，仍然是不可忽视的重要工具。

本文由 Yonghui Wang 创作，采用 知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: Jul 17, 2025 01:01 pm