之前探索了变分自编码器VAE中Encoder和Decoder的理论基础
https://blog.csdn.net/liliang199/article/details/156280212
这里进一步探索向量量化的变分自编码器VQ-VAE的理论基础。
所用示例参考和修改自网络资料。
1. 核心思想与动机
VQ-VAE指Vector Quantized Variational AutoEncoder。
VQ-VAE在VAE上引入离散潜在空间，通过向量量化将连续编码映射到离散码本中。
相比VAE其关键差异在于：
1）学习离散表示，更接近自然信号（语言、语音、图像）的本质
2）避免VAE中的后验坍缩问题
3）为后续自回归建模（如PixelCNN）提供高质量的离散潜在表示
这里从多个角度分析和探索VQ-VAE相比VAE的这些差异改进。
2. 模型架构与数学定义
2.1 编码器Encoder
编码器产生连续向量：
其中是编码器神经网络，输出维度为
。
2.2 向量量化层
向量量化层主要包括码本、向量量化两部分。
码本（Codebook）：
其中是码本大小，离散潜在变量的取值数，d为离散潜在变量的维度。
量化操作：
给定编码器输出，通过最近邻搜索找到码本中最接近的向量：
这产生了离散潜在表示。
2.3 解码器Decoder
解码器从量化后的向量重建数据：
3. 损失函数与训练目标
3.1 总体损失函数
VQ-VAE的损失函数包含三部分：
其中：
表示停止梯度操作，
是超参数通常设为0.25。
3.2 各部分说明
1）重构损失（Reconstruction Loss）
-
使解码器能从量化向量
中重建输入
-
对于图像：通常用L2损失或感知损失
-
对于离散数据：用交叉熵损失
2）码本损失（Codebook Loss）
-
更新码本向量
-
sg操作：码本更新时不考虑编码器的梯度
3）承诺损失（Commitment Loss）
-
鼓励编码器输出承诺于特定的码本向量
-
防止编码器输出在码本向量之间频繁跳跃
3.3 梯度流处理
量化操作不可导，VQ-VAE使用直通估计器(Straight-Through Estimator)处理梯度。
前向传播：使用量化后的向量
反向传播：将解码器的梯度直接复制到编码器输出
数学上表示为：
3.4 两阶段训练流程
VQ-VAE训练分两个阶段。
第一阶段：训练VQ-VAE
-
训练编码器、码本、解码器
-
学习高质量的数据离散表示
-
得到压缩的潜在空间：
第二阶段：训练先验模型
在训练好的VQ-VAE之上，学习离散潜在变量的先验分布：
常用自回归模型有PixelCNN、Transformer，可捕获离散潜在变量之间的长程依赖关系。
4 相比VAE的优化
4.1 VAE vs VQ-VAE
这里动从多个角度对比VAE和VQ-VAE。
1）潜在空间
VAE采用连续连续空间表示，VQ-VAE采用离散表示。
2）后验分布
VAE采用连续的高斯分布，VQ-VAE采用分类分布(ont-hot)。
3）KL项
VAE采用显式KL散度，VQ-VAE则通过量化隐式约束。
4）正则化
VAE采用KL散度正则化，VQ-VAE则采用码本与承诺损失。
5）从高斯先验采样
VAE从从高斯先验采样，VQ-VAE则从码本索引采样。
4.2 VQVAE理论优势
相比VAE，VQ-VAE有一些列优势。
1）避免了后验坍缩
VAE中当解码器过于强大时，可能忽略潜在变量，导致KL散度项趋向于0。
VQ-VAE通过强制离散表示避免了这一问题。
2）解耦表示学习与生成建模
VQ-VAE专注于学习有效的离散表示，先验模型专注于学习序列依赖关系。
这种分离使得每部分可以独立优化
3）支持多尺度架构
VQ-VAE可以扩展到多尺度，形成层次化离散表示。
每一层捕获不同抽象级别的信息。
5. VQ-VAE核心洞察
5.1 目标函数推导
VQ-VAE可视为最小化以下目标
其中第二项保证编码器输出与量化向量的接近。
VQ-VAE与EM算法也有关联，VQ-VAE的训练类似于期望最大化（EM）算法。
E步相当于将数据点分配到最近的码本向量，具体为最近邻搜索。
M步相当于更新码本向量和网络参数以最小化重构误差。
5.2 实际应用扩展
VQ-VAE有一系列进一步改进，比如VQ-VAE-2、VQ-GAN。
VQ-VAE-2使用多尺度层次化潜在空间，底层捕获局部细节，高层捕获全局结构。
VQ-GAN则结合对抗损失生成更高质量的图像。
VQ-VAE在音频领域也应用广泛。
比如将波形编码为离散token序列，使用Transformer学习token序列的分布。
5.3 VQ-VAE总结
VQ-VAE的核心贡献在于：
1）离散潜在表示，更符合许多模态（语言、语音）的本质
2）解耦训练，表示学习与序列建模分离
3）可扩展性，支持层次化、多尺度架构
4）实用性，为自回归生成模型提供高质量的离散输入
VQ-VAE离散表示在生成建模中优势明显，被后续一系列生成式框架采用，
如DALL-E、AudioLM、Codeformer。
reference
VAE中Encoder和Decoder的理论基础的探索
https://blog.csdn.net/liliang199/article/details/156280212
PixcelCNN-Pytorch
https://github.com/singh-hrituraj/PixelCNN-Pytorch
What Is DALL-E?
https://www.coursera.org/articles/what-is-dall-e
audiolm-pytorch
https://github.com/lucidrains/audiolm-pytorch
VQ-VAE-2
https://github.com/vvvm23/vqvae-2
VQGAN
https://github.com/Westlake-AI/VQGAN
CodeFormer