VQ-VAE原理

基础模型概述

自编码器 (Autoencoder, AE)

原理： AE 旨在学习一个潜在表示（Latent Representation），使得输入 可以通过 被准确地重建出来。

$$x \xrightarrow{\text{Encoder}} z \xrightarrow{\text{Decoder}} \hat{x}$$

目标： 最小化重建误差（如均方误差 MSE）。

$$L_{\text{AE}} = ||x - \text{Decoder}(\text{Encoder}(x))||^2$$

缺点： 潜在空间 没有结构约束。无法直接从潜在空间采样生成有意义的新数据，因为 附近的区域可能对应无意义的重建 。

变分自编码器 (Variational Autoencoder, VAE)

原理： VAE 将编码器视为一个学习输入 潜在分布 的过程，通常假设这个分布是高斯分布，由均值 和标准差 参数化。它通过对潜在空间施加约束，使其服从一个简单的先验分布 （如标准正态分布 ），从而实现生成能力。

目标： 最大化证据下界（Evidence Lower Bound, ELBO）。

$$L_{\text{VAE}} = - \text{ELBO} = L_{\text{Reconstruction}} + L_{\text{KL}}$$ $$L_{\text{VAE}} = -\mathbb{E}_{q(z|x)}[\log p(x|z)] + D_{\text{KL}}(q(z|x) || p(z))$$

• 重建损失 ： 确保 尽可能接近 。
• KL 散度损失 ： 迫使潜在分布 接近先验分布 ，保证潜在空间的连续性和可采样性。

缺点： 标准 VAE 的潜在变量 z 是 连续的，而且其先验是一个 各向同性高斯分布

$$p(z) = \mathcal{N}(0, I)$$

这带来两个关键问题：

① 连续潜在空间要求解码器对输入变化“平滑”响应

VAE 的训练目标包含 KL 项：

$$\mathrm{KL}(q(z|x)\;\|\; p(z))$$

KL 项会强制每个样本的后验分布 都要向一个 光滑、统一、分布接近高斯 的空间挤压。

这意味着不同样本的 latent 表示必须在空间中彼此靠得比较近，而且整个映射必须是 连续且光滑 的，解码器不得不产生更“中间态/平均化”的输出。

于是：

• 类似但不完全相同的输入在 latent 空间中被拉得更近
• 解码器为了保持连续性，会对过度靠近的 latent 给出折中输出
• 这种折中输出往往表现为 模糊的重建（尤其图像边缘、纹理）

② VAE 训练目标偏向“平均化预测”

重建损失通常是像素 L2 或 BCE，加入 KL 正则后，整体更鼓励生成“平均像素值”，这同样促进模糊。

向量量化变分自编码器 (VQ-VAE)

向量量化变分自编码器（VQ-VAE）是一种生成模型，它结合了自编码器的重建能力与向量量化的离散潜在空间优势，解决了标准 VAE 潜在空间连续性带来的模糊性问题。

VQ-VAE 引入了一个码本（Codebook），也称为嵌入空间 ，其中 是码本大小，每个 是一个 维的嵌入向量（即Code Vector）。

整个流程如下：

1. 编码（Encoding）： 输入 经过编码器 得到连续的潜在向量 。 $$z_e = \text{Encoder}(x)$$
2. 量化（Quantization）： 这是 VQ-VAE 的核心步骤。对于 中的每一个向量，我们在码本 中找到一个与其欧氏距离最近的 Code Vector ，并将其替换为 。这个被选中的离散表示记为 。 $$z_q = \text{Quantize}(z_e) = e_k \quad \text{where } k = \arg\min_j ||z_e - e_j||^2$$
3. 解码（Decoding）： 量化后的向量 传入解码器，重建输出 。 $$\hat{x} = \text{Decoder}(z_q)$$

VQ-VAE 损失函数推导与训练

VQ-VAE 的损失函数由三个主要部分组成：重建损失、量化损失和承诺损失。

$$L_{\text{VQ-VAE}} = L_{\text{Reconstruction}} + L_{\text{VQ}} + L_{\text{Commitment}}$$

重建损失 (Reconstruction Loss)

这与 AE 或 VAE 的重建项相同，用于优化解码器 和码本 （通过 ）。

$$L_{\text{Reconstruction}} = ||x - \text{Decoder}(z_q)||^2$$

梯度会从解码器 传递回量化向量 ，但不会直接传递过量化操作 。

量化损失 (VQ Loss / Codebook Loss)

该项用于更新码本 。我们希望被选中的 Code Vector 能够尽可能接近它所代表的编码器输出 。

$$L_{\text{VQ}} = ||\text{sg}[z_e] - e_k||^2$$

其中 是 Stop Gradient（停止梯度） 操作，表示在反向传播时，该操作的输入（即 ）不会接收梯度。因此，此损失项仅用于更新 （码本）。

承诺损失 (Commitment Loss)

该项用于更新编码器 。它确保编码器的输出 必须足够接近它所映射到的 Code Vector 。这阻止了 的方差增长过大，确保 能够「承诺」（Commit）于码本中的一个向量。

$$L_{\text{Commitment}} = \beta ||z_e - \text{sg}[e_k]||^2$$

其中 停止了梯度流向 ，因此此损失项仅用于更新 （编码器）， 是一个超参数，用于平衡此项的重要性。

梯度传播机制

由于量化操作 是一个不可微分的操作，我们无法直接通过常规的反向传播来训练编码器 。

VQ-VAE 采用了 直通估计器（Straight-Through Estimator, STE） 来解决这个问题：

• 前向传播（Forward Pass）： 执行量化操作 。
• 反向传播（Backward Pass）：
  • 假设量化操作是一个恒等函数（即 ）。
  • 因此，解码器的梯度 被直接复制并传递给编码器的输出 。

$$\frac{\partial L}{\partial z_e} = \frac{\partial L}{\partial z_q} + \frac{\partial L_{\text{Commitment}}}{\partial z_e}$$

结合所有损失项，VQ-VAE 的完整优化目标是：

$$\min_{\text{Encoder}, \text{Decoder}, E} \quad L_{\text{Reconstruction}} + L_{\text{VQ}} + \beta L_{\text{Commitment}}$$

VAE 与 VQ-VAE 的关键区别

VAE 的潜在空间是连续的，编码器会输出均值和方差，用重参数技巧对潜在变量进行采样，并通过最大化 ELBO 来训练。它从高斯先验中直接采样潜在向量进行生成。由于潜在空间是连续的，生成结果有时会出现平均化或模糊的问题。

相比之下， VQ-VAE 采用离散的码本作为潜在空间，编码器首先产生一个连续向量，然后将其量化为码本中的最近向量。它的损失由重建损失、码本损失和 commitment 损失组成，不包含 KL 项。在生成阶段，会对离散码序列训练一个自回归先验，再通过解码器生成样本。由于离散表示的特性，VQ-VAE 更容易保留纹理与细节，也更适合用于压缩和层次化建模。

VQ-VAE 通过引入可学习的码本和不可微分的量化操作，成功地在自编码框架中使用了离散潜在空间，使其在处理结构化、高频细节的数据生成任务中，展现出比传统 VAE 更强大的能力。