12 第十二章离散潜空间

近年来随着Transformer被广泛运用，离散潜在空间也变得越来越受欢迎。其中VQ-VAE和RVQ是离散潜空间模型的代表。

12.1 VQ-VAE

考虑卷积类型的编码器解码器架构，无论是一维序列输入或二维图像输入，潜变量都可以表示为特征加多通道的形式：

VQ-VAE的策略是，对于潜变量，维护一个码本，对于编码器计算出的潜变量特征图，以通道方向为纵方向，把每个通道方向的向量替换成码本里最接近的向量。

因此码本的高度\(D\)就是通道个数，码本的宽度\(K\)常见的取值范围通常在512到8192之间，并且习惯设为2的幂次方：

我们将码本表示为包含 \(K\) 个向量的集合 \(E = \{e_1, e_2, \dots, e_K\}\)，其中每个码字 \(e_i \in \mathbb{R}^D\)。对于编码器在某个位置输出的连续特征向量 \(z_e \in \mathbb{R}^D\)，量化过程其实就是计算 \(z_e\) 与码本中每一个向量 \(e_i\) 的欧氏距离，并挑选出距离最近的那一个的索引 \(k\)： \[ k = \arg\min_{i} \|z_e - e_i\|_2 \tag{12.1}\] 找到索引 \(k\) 之后，把对应的码字 \(e_k\) 拿出来替代原先的 \(z_e\)，得到量化后的潜向量 \(z_q\)： \[ z_q = e_k \tag{12.2}\]

特征图完成这样的就近替换后，形成一个量化了的特征图，接着喂给解码器恢复出原始的图像或音频信号。

12.2 损失函数

在传统自编码器中，误差可以通过梯度反向传播直接一路传回编码器。然而在 VQ-VAE 中，量化操作涉及式 12.1 的 \(\arg\min\) 寻找最近邻，这是一个不可求导的离散操作，将导致解码器的梯度无法直接传回编码器。

为了解决这个问题，VQ-VAE 引入了直通估计器（Straight-Through Estimator, STE）：在前向传播时，送入解码器的确实是离散量化后的码字 \(z_q\)，但在反向传播时，直接把解码器对 \(z_q\) 的梯度原封不动地复制给编码器的输出 \(z_e\)，即： \[ \nabla_{z_e} \mathcal{L} \approx \nabla_{z_q} \mathcal{L} \tag{12.3}\]

然而，仅靠 STE 并不能解决所有问题，STE只负责把解码器的梯度传给编码器，为了训练网络，需要让输出接近输入，也就是优化重构损失。

同时双向地优化编码器的输出和码本，让编码器的输出尽量接近码本，也让码本尽量接近编码器的输出。

因此，VQ-VAE 将训练目标精妙地拆解为三部分： \[ \mathcal{L} = \underbrace{\mathcal{L}_{\text{recon}}(x, \hat{x})}_{\text{重构损失}} + \underbrace{\beta \|z_e - \text{sg}[e_k]\|_2^2}_{\text{让潜变量接近码本 (承诺损失)}} + \underbrace{\|\text{sg}[z_e] - e_k\|_2^2}_{\text{让码本接近潜变量 (码本损失)}} \tag{12.4}\]

其中 \(\text{sg}[\cdot]\) 表示梯度截断算子，其包裹的变量在前向传播时正常参与计算，但在反向传播时不会产生梯度（当成是一个常量）。

下面分别对这三部分进行详细说明：

重构损失

重构损失负责保证模型的保真度。它衡量解码器通过量化潜变量 \(z_q\) 还原出的样本 \(\hat{x}\) 与原始输入样本 \(x\) 之间的差异。
承诺损失

该项写为 \(\beta \|z_e - \text{sg}[e_k]\|_2^2\)，其中 \(\beta\) 是一个控制权重的超参数（原论文中通常取 \(\beta = 0.25\)）。它仅优化编码器。通过截断码字 \(e_k\) 的梯度，把选中的码字视为固定的“靶标”，强制要求编码器的输出 \(z_e\) 向离它最近的码字靠拢，对当前的码字选择做出“承诺”，防止 \(z_e\) 肆意发散。
码本损失

该项写为 \(\|\text{sg}[z_e] - e_k\|_2^2\)，它的作用对象是码本向量 \(e_k\)，让码本能够随着学习到数据的特征分布（换个角度理解，优化码本就像选出一些最具代表性的VAE潜变量，量化替换后输出效果差不多的同时实现数据压缩）。在实际工程实现中，这一项通常被指数移动平均（EMA）替代。