VQ-VAE

相比于普通的VAE，VQ-VAE能利用codebook机制把图像编码成离散向量，为图像生成类任务提供了一种新的思路。VQ-VAE的这种建模方法启发了无数的后续工作，包括声名远扬的Stable Diffusion。

神经网络离散表征学习

在机器学习领域，无监督学习有用表征仍是关键挑战。当前生成模型在图像、音频和视频等方面取得进展，但无监督学习的通用表征实用性仍有待提升。许多研究聚焦于连续特征表征学习，而本文关注离散表征，因其更契合多种数据模态（如语言、语音）且适用于复杂推理、规划和预测学习。

离散潜变量
- 定义潜嵌入空间$e \in R^{K×D}$，$K$为离散潜空间大小，$D$为每个潜嵌入向量维度。模型输入$x$经编码器输出$z_{e}(x)$，离散潜变量$z$通过在共享嵌入空间$e$中最近邻查找计算得出（公式1），解码器输入为对应的嵌入向量$e_{k}$（公式2）。后验分布$q(z|x)$为确定性的one - hot分布，通过定义简单均匀先验，KL散度为常数$\log K$。
学习过程
- 由于方程2无真实梯度，采用类似直通估计器方法，将梯度从解码器输入$z_{q}(x)$复制到编码器输出$z_{e}(x)$。总损失函数$L$（公式3）由三部分组成：重建损失（优化解码器和编码器）、向量量化（VQ）目标（学习嵌入空间，使嵌入向量向编码器输出靠近）和承诺损失（确保编码器输出稳定）。解码器仅优化重建损失，编码器优化重建损失和承诺损失，嵌入向量由VQ目标优化。实验中发现该算法对$\beta$（承诺损失权重）稳健，通常取值0.25。模型假设$z$的先验为均匀分布，训练时可忽略KL项。
先验分布
- 离散潜变量的先验分布$p(z)$为分类分布，训练VQ - VAE时先验保持恒定均匀，训练后拟合自回归分布（如对图像使用PixelCNN，音频使用WaveNet）以便通过祖先采样生成$x$。联合训练先验和VQ - VAE有待未来研究。

与连续变量模型对比
- 在CIFAR10数据集上，将VQ - VAE与普通VAE（连续变量）及VIMCO（具有独立高斯或分类先验）对比。VQ - VAE在压缩表示下实现类似VAE的良好重建效果，且是首个在性能上挑战连续VAE的离散潜变量模型。
与其他离散VAE模型对比
- 现有训练离散VAE的方法如NVIL和VIMCO，存在性能差距或在复杂数据集上评估受限等问题。本文VQ - VAE模型通过新的训练方式解决了这些问题，在复杂图像和语音数据集上表现出色。

以$128×128×3$图像为例，通过VQ - VAE压缩到$32×32×1$离散空间（$K = 512$），用PixelCNN建模先验。重建图像虽稍模糊但保留大部分信息，训练PixelCNN先验后生成的样本视觉效果良好。在DeepMind Lab环境的$84×84×3$帧实验中，重建效果类似。
进一步训练第二阶段VQ - VAE（含PixelCNN解码器），尽管因压缩无法完美重建图像，但表明模型可有效利用潜空间，避免了“后验崩溃”问题。

在VCTK语音数据集上，训练VQ - VAE（编码器含6个步长卷积层，解码器基于WaveNet）。通过编码和解码实验表明，模型可提取仅保留长期相关信息的潜空间，学习到与语音内容相关的高层抽象空间，不受低层级特征影响。
在更大语音数据集上训练模型，通过训练先验生成的无条件样本包含清晰单词和部分句子，表明VQ - VAE能以无监督方式建模基本音素级语言模型。
进行说话人转换实验，成功实现将一个说话人的语音内容转换为另一个说话人的声音，证明编码表示可分解出说话人特定信息。
将离散潜变量与真实音素序列对比，分类准确率达49.3%（随机潜空间仅7.2%），表明离散潜变量是与音素紧密相关的高层语音描述符。

在DeepMind Lab环境中，训练VQ - VAE生成基于给定动作序列的视频。模型可在潜空间生成视频序列，再通过确定性解码器映射到像素空间，生成的视频序列视觉质量良好，保持局部几何形状，且能成功根据给定动作生成帧序列，无视觉质量下降。

未来研究方向可包括联合训练先验和VQ - VAE以进一步提升模型性能，探索更适合的损失函数（如基于感知的损失函数）以提高图像重建质量，以及进一步研究离散潜变量在其他领域的应用等。