浅学一下VQ-VAE的内容，记录一下方便以后查阅。
Vector Quantized Variational Autoencoders （VQ-VAE）
- 基本理念
- [Why do it?](#Why do it?)
- Architecture
- [Vector Quantization Layer](#Vector Quantization Layer)
- [Embedding Space](#Embedding Space)
- 模型为何训练的起来
- [Loss Function](#Loss Function)
- [VAE 与 VQ-VAE 的潜在空间的区别](#VAE 与 VQ-VAE 的潜在空间的区别)
- 参考资料
基本理念
VAE 由三部分组成：
- 一个编码器网络，是一个后验概率 q θ ( z ∣ x ) q_θ(z|x) qθ(z∣x)，观测到 x x x推测出 z z z
- 一个先验分布 p ( z ) p(z) p(z)
- 一个解码器网络，是一个先验概率 p ϕ ( x ∣ z ) p_ϕ(x|z) pϕ(x∣z) ，给与 z z z推测出 x x x
通常情况下，我们假设先验值和后验值为正态分布，方差为对角线。然后使用编码器预测后验的均值和方差。
在VQ-VAE的工作中，作者使用的是离散潜变量（而不是连续正态分布）。后验分布和先验分布是分类的，从这些分布中抽取的样本会索引一个嵌入表。换句话说
- 编码器对分类分布进行建模，从中进行采样以获得整数值
- 这些整数值用于索引嵌入字典
- 然后将索引值传递给解码器
Why do it?
现实世界中的许多重要对象都是离散的。例如，在图像中，我们可能会有 "猫"、"汽车 "等类别，在这些类别之间进行插值可能没有意义。离散表示法也更容易建模，因为每个类别都有一个单一的值，而如果我们有一个连续的潜在空间，那么我们就需要对这个密度函数进行归一化处理，并学习不同变量之间的依赖关系，这可能会非常复杂。
离散化的好处是，不同的类别可以获得不同的信息空间，而不是所有类别的大杂烩，因为大杂烩很容易影响互信息。
此外，作者还声称，他们的模型不会出现 "后塌陷"（posterior collapse）问题，而这个问题一直困扰着 VAE，并妨碍了复杂解码器的使用。
Architecture
上图显示了架构中的各个顶层组件以及每个步骤的尺寸。假设在图像数据上运行我模型，下面是一些字母的解释：
- n n n：批量大小
- h h h：图像高度
- w w w：图像宽度
- c c c：输入图像的通道数
- d d d：隐藏状态下的通道数
现在可以按以下步骤解释工作原理：
- 编码器接收图像 x : ( n 、 h 、 w 、 c ) x:(n、h、w、c) x:(n、h、w、c)并输出 z e : ( n 、 h 、 w 、 d ) z_e:(n、h、w、d) ze:(n、h、w、d)
- 矢量量化（Vector Quantization）层接收 z e z_e ze，并根据距离从字典中选择嵌入，然后输出 z q z_q zq
- 解码器消耗 z q z_q zq 并输出 x ′ x' x′，试图重现输入 x x x
Vector Quantization Layer
如上图所示，VQ 层的工作可分为六个步骤：
- Reshape：将除最后一个维度外的所有维度合并为一个维度，这样就有了 n ∗ h ∗ w n*h*w n∗h∗w 个矢量，每个矢量的维度为 d d d
- Calculating distances：对于 n ∗ h ∗ w n*h*w n∗h∗w 向量中的每个向量，我们计算其与嵌入字典中 k k k 个向量中的每个向量的距离，从而得到一个形状为 ( n ∗ h ∗ w , k ) (n*h*w, k) (n∗h∗w,k) 的矩阵。
- Argmin：对于 n ∗ h ∗ w n*h*w n∗h∗w 个向量中的每个向量，我们从字典中找出 k k k 个向量中最接近的向量的索引。
- Index from dictionary：为 n ∗ h ∗ w n*h*w n∗h∗w 个向量中的每个向量从字典中找出最接近的向量索引
- Reshape：转换回形状 ( n 、 h 、 w 、 d ) (n、h、w、d) (n、h、w、d)
- Copying gradients：如果你一直看到现在，就会意识到不可能通过反向传播来训练这个架构，因为梯度不会流经 argmin。因此，尝试通过将梯度从 z q z_q zq 复制到 z e z_e ze 来进行近似。这样，我们实际上并没有最小化损失函数，但仍能传递一些信息回去进行训练。
Embedding Space
在空间上，论文定义 e ∈ R K × D e \in \mathbb{R}^{K \times D} e∈RK×D，意思是有 K K K 个长度为 D D D 的潜在空间组成。 K K K 代表有 K K K 种类； D D D 就是有 D D D 个维度。
上式为编码器的后分类分布概率。
对于编码器输入的 z z z，等式计算得到的潜在空间与嵌入空间之间每个潜在空间的 L 2 L2 L2 距离，生成映射 q ( z ∣ x ) q(z|x) q(z∣x)。将最近的嵌入空间作为解码器的输入。如下面的等式所示 z q ( x ) z_{q}(x) zq(x)：
在论文中的上图里，在 z z z 的每一格中找到对应最接近的潜在空间，并且作为解码器的输 入 入 入。
模型为何训练的起来
离散化主要的问题是在神经网络做反向传播时。因为它是离散的，这也意味着它不可微分。
因此在 VQ-VAE 中，将解码器返回的梯度直接回流到编码器输出端，不通过嵌入空间或 q ( z ∣ x ) q(z|x) q(z∣x)。
不过为什么可以直接将decoder的梯度回传到encoder？论文中提到，因为decoder的目的是要尽量接近嵌入空间的距离，而encoder也是接收嵌入空间的内容。一直以来都是学习由维度 D D D 组成的空间，如下方图片所示，解码器回传的 ∇ z L \nabla_z L ∇zL 藏有可观的信息，间接推理出编码器应该如何靠近那些向量。
Loss Function
总损失实际上由三部分组成：
-
Reconstruction loss：优化解码器和编码器：
r e c o n s t r u c t i o n L o s s = − l o g ( p ( x ∣ z q ) ) reconstructionLoss = -log( p(x|z_q) ) reconstructionLoss=−log(p(x∣zq)) -
Codebook loss：由于梯度会绕过嵌入，因此采用了字典学习算法，利用 L 2 L2 L2 误差将嵌入向量 e i eᵢ ei 移向编码器输出：
c o d e b o o k L o s s = ‖ s g [ z e ( x ) ] − e ‖ 2 codebookLoss = ‖ sg[z_e(x)]− e ‖^2 codebookLoss=‖sg[ze(x)]−e‖2其中 s g sg sg 代表停止梯度运算符，表示没有梯度流过该运算符所应用的对象。
-
Commitment loss：由于嵌入空间的体积是无量纲的，如果嵌入 e i eᵢ ei 的训练速度不及编码器参数的训练速度，它就会任意增长。
c o m m i t m e n t L o s s = β ‖ z e ( x ) − s g [ e ] ‖ 2 commitmentLoss = β‖ z_e(x)− sg[e] ‖^2 commitmentLoss=β‖ze(x)−sg[e]‖2其中 β 是一个超参数，用于控制与其他成分相比，我们希望在多大程度上权衡承诺损失
VAE 与 VQ-VAE 的潜在空间的区别
下图明了 VAE 和 VQ-VAE 的根本区别。潜空间中的点在 VQ-VAE 中是离散的，而在 VAE 中却是连续的。
参考资料
Understanding Vector Quantized Variational Autoencoders (VQ-VAE)
Tutorial - What is a variational autoencoder?
[筆記][NeurIPS2017] VQ-VAE: Neural Discrete Representation Learning
A Crash Course on VAEs, VQ-VAEs, and VAE-GANs
Difference between AutoEncoder (AE) and Variational AutoEncoder (VAE)