VQ-VAE离散潜在表示与矢量量化

VQ-VAE由van den Oord等人于2017年（DeepMind）提出，通过离散潜在代码和矢量量化（Vector Quantization）显著提升了VAE的生成质量和保真度。

核心动机： 标准VAE假设潜在z服从连续高斯分布，但这会导致后验塌缩和后验分布过于平滑，生成图像模糊。VQ-VAE使用离散的潜在表示，每个潜在编码是一个来自有限码本的索引，迫使模型学习更清晰的模式。

架构：

1. 编码器E：输入图像x → 连续特征图z_e(x)，形状H'×W'×D。
2. 矢量量化层（核心）：
   • 码本（Codebook）：K个D维嵌入向量{e_1, e_2, ..., e_K}（K=512或更大的离散词汇量）。
   • 对每个空间位置的特征向量z_e(x){i,j}，在码本中找到最近邻： z_q(x) = e_k，其中k = argmin_j ||z_e(x){i,j} - e_j||_2
   • 将连续特征量化为离散索引，组成潜在映射[H'×W']，每个位置存储码本索引。
3. 解码器D：从离散特征z_q(x)重建图像。

训练与损失函数： L = ||x - D(z_q)||² + ||sg[z_e] - e||² + β·||z_e - sg[e]||²

• 重建损失：第一项，确保解码器能重建原图。
• 码本学习：第二项（VQ损失），使码本向量接近编码器输出。其中sg是stop-gradient操作。
• 承诺损失：第三项（β=0.25），确保编码器输出'承诺'靠近码本向量。

先验学习（Prior Learning）： 训练完VQ-VAE后，在离散潜在空间上学习先验分布p(z)（如PixelCNN或Transformer）。生成时：

1. 从先验采样离散索引序列。
2. 通过解码器生成图像。

与VQ-GAN的关系： VQ-GAN（2021）将VQ-VAE的感知损失替换为GAN的对抗损失+L1损失，并引入PatchGan判别器，进一步提升了生成图像的感知质量。

应用：

• 高质量图像生成（如生成512×512人脸）。
• 图像压缩（如CompressAI）。
• VQ-VAE的离散潜在空间也为图像-文本模型的跨模态表示提供了基础。