一文详解 codebook 技术史（从 VAE 到 VQ/RQ-VAE 到 FSQ）

VAE：

VAE (variational autoencoder, 变分自编码器) 是一种强大的生成模型, Encoder 把数据编码到隐空间 , 其学习条件概率, Decoder把数据从隐空间中重建回来 ，其学习另一个条件概率 。VAE 额外有一个限制条件是让 满足 Gaussian分布。这样做的好处就是训练结束后可以扔掉 Encoder, 直接从这个先验分布 上随便采样 , 然后通过 Decoder 就能生成一个 。

VAE 最主要的是这个 ELBO ：

ELBO，即evidence low bound，evidence指的就是 ，而 ELBO 表示 evidence 的最小期望。我们要让这个 lower bound 尽可能变大，得到的模型就会更可能产生我们期望看到的 。

为解释 ELBO 是怎么来的，我们一步一步来看。

K-L散度：

我们首先讲解 KL 散度，为衡量模型生成的分布与原始分布的相似度，常用的便是 K–L（ Kullback–Leibler ）散度。定义如下，对于两个具有概率密度函数 和 的分布：

K–L 散度具有两个重要性质：

1. 不对称性：显然，K–L 散度对于 和 来说是不对称的。

2. Gibbs 不等式：它总是【非负】的，并且当且仅当 和 在每一处都相同时才为 0。

为了理解这一点，我们可以将 KL 散度分解为两部分：

第二项带有负号，其对应的是 的信息熵；第一项也带有负号，代表 和 之间的交叉熵。第一项始终不大于每个给定符号下的第二项，这便是 Gibbs 不等式；而 Gibbs 不等式的证明可以使用 Jensen 不等式：

若 是凸函数，则有：

设 由于 所以其为凸函数，以及

那么：

由

VAE 理论框架（联合概率建模角度）：

VAE 框架可以从多个角度建立，例如概率分布视角、贝叶斯视角 以及 联合概率视角，这里我选用联合概率这一简单的方法来阐述：

假设原始数据样本为 ，分布为 ，我们希望借助隐变量 （标准正态分布）来建模 ，因此我们设立 来逼近 :

是标准正态分布， 是我们的生成式模型；此外还需明确的是 是 的原始分布， 是encoder生成的 ，训练时要让其逼近正态分布。

我们直接采用联合建模的角度，原来我们的目的是让 来逼近 ，我们转变下思路变为让 与 越相近越好，注意除了中也有参数：

KL 散度便是我们的终极目标，我们将从这个 KL 散度推导出最终的 ELBO：

这里被我们拆开为两项，第一项：

无论 是什么，它一定是确定的，故第一项是常数

第二项：

因此我们很快便得到了最终的 ELBO，注意多了个负号。

ELBO：

ELBO 有两项，分别为： 以及 – ，这两部分可以理解为【重构误差项】 以及【KL散度项】：

重构误差项：这部分度量了模型生成数据的质量，即解码器 使用从编码器 采样的 来重构输入 的准确性，这是负对数似然，表明给定潜在变量 后，重构原来的 的概率有多大。目标是最大化这部分期望值，即希望模型能生成与输入 尽可能接近的数据。

KL散度项：- 是后验分布 和先验分布 之间的负K–L 散度，以此衡量编码器的输出分布与标准正态分布的差异。目标是最小化KL散度，确保潜在变量 z 尽可能接近正态分布。

至此我们推导出了VAE的损失函数，了解了ELBO的原理。

VQ-VAE：

paper：[Neural Discrete Representation Learning]

https://arxiv.org/abs/1711.00937

背景：

VAE中的隐变量 z 的每一维都是一个连续值， 而VQ-VAE 中 的每一维都是离散的整数，这些整数便可 index 到已训练好的 codebook（码本，本质上就是一批 embedding）。这样做符合自然界模态的特点，例如语言本质上就是由很多字符组成，每个字符都可以是用数字索引到字符库里的某个字符，NLP中可以理解为token_id索引到vocab里的某个token，所以VQ-VAE可以理解为 【图像tokenization】 的过程，事实上这种思想可以借鉴引用到很多领域，例如广告推荐里将广告用一串索引表示。

文章还指出，VAE 存在后验坍塌（Posterior Collapse） 的问题，这一般是由散度消失（KL-Vanishinig）导致的，因此该问题也称为KL-vanishing。简单来说就是解码器太强，模型的 潜在空间（latent space）无效化，即编码器 退化为与先验 相同的分布，ELBO里的KL散度项为0，而忽略了输入数据的信息。

方法：

将隐变量 离散化的关键操作是VQ, 即 vector quatization。

1. 图像 输入至 encoder 中得到
2. codebook 是一个K*D 的 table（紫色方块）：
3. 将中每一维都映射为 codebook 中K个embedding之一
4. 全部替换后图中绿色的变为紫色的，然后进行重构

从到的变化可以理解为聚类，如图中右子图所示，由于变化后的embedding位于codebook内，当然就可以只用整数来表示。

训练：

ELBO 损失项：

我们先看原有的 ELBO ，这里p和q互换以与图示对应，q代表encoder，p代表decoder；

这里后验分布 里都是one-hot向量，如下所示：

而非之前VAE里的正态分布，由此 预估的每一维都是codebook里每个embedding的概率；我们假设采样的先验分布 是均匀分布，则每一维对于某个embedding选取概率有 ，则有：

第一项表示one-hot中为1对应的那一维对KL散度的贡献，第二项代表其他维的贡献。

因此 ELBO 中第二项可以忽略，只有重构损失项。

那我们再看第一项损失，可以简单写为：

然而 包含了argmin，这个操作是没有梯度的，无法更新 encoder；VQ-VAE 使用了一个很精巧也很直接的方法，称为 Straight-Through Estimator，称为“直通估计(https://papers.cool/arxiv/1308.3432)”。其思想是在前向传播的时候可以任意变量（可以不可导），而反向传播的时候，直接 跳过 这个不可导的操作。对应图1中红色箭头，表明跳过 的操作。

根据这个思想，我们设计的目标函数是

sg 代表阻止梯度回传

codebook 损失项：

为使得 与 尽量接近，设置损失： ；

这里我们理解下：是编码器得到的，是离得最近的embedding，两者都有可训练的参数；因此在实际训练时，codebook相对自由宽松，没什么限制条件，而编码器生成的要保证重建效果，我们更希望主要靠近，并且因为 的梯度等于以及梯度之和，故可拆解为：

第一项可以理解为不变，主要靠近，第二项相反，由此我们可以给第二项设置一个相对较小的权重，来达到更希望主要靠近的效果。

整体损失项：

文中指出，实验发现 设置[0,1]均具有鲁棒性，故使用 ，还可以使用滑动平均的方式更新，下面阐述。

滑动平均方法：

具体来说使用 指数移动平均（EMA）来更新 codebook ：

设 为编码器输出中最接近词典项 的一组 个元素，那么可以将损失写为：

理论上可以求得 的最优值，可以通过封闭形式的解求得，即该集合中所有元素的平均值：

这种更新方法通常用于 K-Means 等算法。然而，当处理小批量（minibatches）时，无法直接使用上述更新方式。因此，我们可以采用指数移动平均，作为该更新的在线版本：

其中， 的取值范围在 0 到 1 之间，论文发现 0.99 是一个不错的选择。

应用：

按照之前 VAE 的逻辑，使用时去掉encoder，在正态分布里采样即可生成图片；那么VQ-VAE呢？其假设先验分布为均匀分布，然而并没有直接在均匀分布里采样，而是使用 PixelCNN 来学习编码的分布（这里非常奇怪，在issue一节讨论），即学习 。

简单介绍下，PixelCNN 是一种采用自回归方式逐像素从左上角生成的图像生成模型，其中使用了mask conv操作，可以类比 GPT，使用 mask self-attention 操作。

所以最后我们通过 PixelCNN 来随机生成 ，然后再用VQ-VAE的 Decoder 来生成最后的图片。

Issue：

VQ-VAE 到底是不是 VAE ？

VAE 的核心是encoder学习一个先验分布，最后只需要从这个先验分布里采样就可以用来生成，然而VQ-VAE事实上并不行，其假设先验分布为均匀分布，但并不能从均匀分布里采样解码得到真实图像，这就说明这就不过只是一个AE 类模型。

那么问题出在哪了？回顾 VQ-VAE 的设计，发现并没有类似 VAE 里的 KL散度loss 来迫使先验分布逼近均匀分布。你可能会问假设分布是均匀分布，KL散度是一个常数呀，上面不是还推导了？那么我们再回顾一下：

KL散度是常数，那么这一项就不会优化，也就不存在要让 更逼近 的说法，也就是 不会被更新，其生成的分布根本不可控。

那么继续深究，这一项为何会是常数？原因就在于 始终是一个one-hot分布，无论怎么优化都是如此，而one-hot分布和均匀分布的 KL散度 始终是 logK，因此 ELBO里 的这一项毫无意义。

其实本质上VQ-VAE 做的是【图像 tokenization】的工作，生成模型部分交给自回归模型 PixelCNN 去负责了。

此外：苏神在 博客 评论里还指出 VQ-VAE里边从均匀分布采样离散的code直接传入decoder，生成结果也不至于差得完全不可看，还是勉强能看的，比纯AE要好点，但要保证质量，还是得 pixelcnn。

VQ-VAE 的核心贡献？

核心贡献不在于其提出了一种新的 VAE 架构，而在于提供了一个序列压缩技术。正如上所说，其本质是一个利用codebook 做图像 tokenization 的工作，然而这种 codebook 的思想不仅可以应用于图像，音频、视频甚至短视频、广告都是可以的，所以我们才看到VQ-VAE的思想应用于各个领域，这才是VQ-VAE的魅力所在。

VQ-VAE-2：

论文：https://arxiv.org/pdf/1906.00446

主要变化就是把 VQ-VAE 的 encoder 和 decoder 都进行了分层, bottom层对local feature进行建模，top层采取全局自注意力机制。

RQ-VAE：

paper：https://arxiv.org/pdf/2203.01941

背景：

VQ-VAE 的序列长度较长，需要大量的codebook，这势必会导致_codebook collapse（码本摊缩）_问题，使得VQ-VAE的训练很不稳定；而 RQ-VAE 则采取一种 _residual quantization（残差量化）_的新方法，通过D轮迭代，将feature map表示为D个堆叠的离散编码，可以进一步减小feature map（可以理解为经过encoder后的表示）的spatial resolution，例如从原始图像的256256变为88。这样 进一步增加下采样因子 减少分辨率，使得 AR 模型能够减少计算成本、提高图像生成速度，并更好地学习codebook中各向量之间的长依赖关系。

方法：

RQ v.s. VQ：

VQ：

假设codebook表示为 ，对于向量 ，其映射为近邻向量的操作表示为：

给定图片输入为 ，提取的 feature map 为：，通过映射后得到的code map为： ，其中 是feature map中(h,w)位置上的向量。

假设 codebook 大小为 K，那么整个feature map为 个 bit，根据_rate-distortion theory（率失真理论）_，H和W每缩小一半，K都要增加到 ，因此说VQ-VAE需要大量的codebook。

RQ：

在RQ里，定义新的映射为近邻向量的操作：

可以看到并非之前单一的数字，而是一个元组，那么每一位的k如何选择？首先初始化残差 mathbf{r}_0=mathbf{z} ，然后按照如下方法计算：

可以这么理解，我要模拟 ，但是我模拟的 肯定和 有差距，我用 表示出来这两者的差，然后我继续模拟 ，但是我模拟的 肯定又和有差距，我用表示出来…… 因此每个 逐步相加，理论上和要模拟的 越来越逼近。

可以看出VQ将空间分为K个簇，而RQ将空间分为 个簇，来实现更精确的量化。

共享codebook机制：

虽然我们可以为每一层深度 d 分别构建一个码本，但在每个量化深度上使用的是单个共享码本。共享码本在构建 RQ 近似向量 z 时有两个优势：

1. 使用单独的码本需要广泛的超参数搜索，以确定每一层的码本大小，而共享码本只需确定总码本大小 K。

2. 共享码本使得所有的 embedding 在每一层量化时都可用。因此，每一层都可以使用相同的 embedding，以最大化其效用。

RQ-Transformer：

可以看出编码得到的 feature map 输入给 Transformer 来作为自回归任务的输入，整个 RQ-Transformer 分为Spatial Transformer和 Depth Transformer 两部分。

输入处理：

RQ-VAE 提取的代码映射 会按照 栅格扫描顺序（raster-scan order）重新排列为二维数组 ，其中 。每一行 包含 D 个代码：

自回归建模总公式为：

建模动机：

直接将 mathbf{S} 展开为长度 TD 的序列并输入传统 Transformer 的方法存在不足，无法利用导 RQ-VAE 降低后的长度 T的优势。此外，这种直接展开会增加计算成本。由此设计为 Spatial Transformer和 Depth Transformer 两部分。

空间 Transformer（Spatial Transformer）：

首先空间 Transformer的输入为每个位置上的 feature（各个残差项之和），并加上位置编码（PE），如下：

整个 Spatial Transformer 表示为：

深度 Transformer (Depth Transformer)：

深度 Transformer 的任务是在给定位置 t 自回归地预测 D 个残差项code，即

在深度 d 和位置 t 时，Transformer 的输入 被定义为之前深度的嵌入之和：

每个深度的预测基于之前所有深度的估计，使得每一层的估计更加精细。

是深度 d 的位置嵌入，且在所有位置 t 上共享。

整个 Depth Transformer 表示为：

训练：

RQ-VAE 的训练损失函数 包含两部分：

重构损失（Reconstruction Loss）：

这个损失度量的是输入 和重构结果 之间的欧氏距离，用于确保重构后的样本尽可能接近原始输入。这里同样会采用 Straight-Through Estimator。

承诺损失（Commitment Loss）：

（sg[·] 是 stop-gradient 操作符，用于在反向传播时阻止梯度的传递），该损失的作用是最小化每个维度 d 上的量化误差，从而鼓励编码器的输出 更接近量化后的值 。

论文内提及codebook会采用聚类特征的指数滑动平均来更新，从而提升模型的训练效果和稳定性。

RQ-VAE 同时还采用了对抗训练（Adversarial Training ）以提高重构图像的感知质量。采用了基于 patch 的对抗损失和感知损失。

负对数似然损失 (Negative Log-Likelihood, NLL) ：

用于训练 RQ-Transformer：

Trick：

曝光偏差 (Exposure Bias)：

曝光偏差是自回归（AR）模型中的常见问题。在训练和推断阶段，由于预测错误的累积，模型性能会下降。尤其是在 RQ-Transformer 中，随着深度 D 的增加，量化特征向量的估计变得更加困难，误差也会累积。

论文采用了软标签 (Soft Labeling) 和 随机采样 (Stochastic Sampling)策略：

软标签（Soft Labeling）：

基于 RQ-VAE 中代码嵌入之间的几何关系，定义了一个温度参数 控制的类别分布：

当 时，分布 会收缩为一个 one-hot 分布：

软标签的作用：

利用嵌入之间的几何距离，为目标代码的监督引入了软标签分布; 在位置 和深度 上，假设特征向量为 ，并令残差向量为 。负对数似然（NLL）损失使用了该软分布作为监督。

区别于 one-hot 标签，该监督机制使用了软化后的分布 。

随机采样（Stochastic Sampling）：

在原始的 RQ-VAE 中，代码选择是确定性的。然而，这里通过从软分布 中进行采样来选择代码 。当 时，随机采样等价于原始确定性代码选择。

优势：随机采样为特征映射提供了不同的代码组合，从而缓解了训练和推断中的不一致性。

FSQ：

paper：Finite Scalar Quantization: VQ-VAE Made Simple

https://arxiv.org/abs/2309.15505

方法：

论文提出使用 FSQ（Finite Scalar Quantization） 来替代 VQ-VAE中的“VQ”，其离散化思路非常简单，就是“四舍五入”。如上图所示，假设最后要把x映射为d维（图中d=3），我们把z的每一维用L个value表示（图中L=3），然后将z的每一维的L个value四舍五入（图中则变化为正方体的边线所在顶点处），由此便离散化了。

还有个区别图式中便是VQ里量化后的 会用一个单独的数字代替，表示codebook里的索引；而FSQ里会用L个数字组成的元组（例如(-1,0,1)）来替代，也表示索引，整体codebook数量为L^d，图里为9。

方案对比如下：

具体来说给定一个 d -维表示 ，我们的目标是将 量化为有限的码字集。为此，我们首先应用一个边界函数 ，然后将结果四舍五入为整数。我们选择 使得 取得 个唯一值之一（例如， ），上图的右子图可视化了这个转化，由于tanh取值范围为(-1，1)，由此z的范围是 () ，故四舍五入后便是L个取值，图中L=5，则有-2,-1,0,1,2这5个取值。

由此，我们得到 ，其中 便是码本，且 。

为了在整个四舍五入操作中传播梯度，使用了前述 STE（直通估计） 技巧 ，通过以下方式轻松实现“停止梯度（sg）”操作：

实验：

从图中可以看到，编码表大小2^10是一个分界点，在2^10左右时，FSQ与VQ的效果接近；超过2^10时，FSQ占优，反之小于2^10时，VQ占优。文中建议 ，并且d是个位数，相比之下VQ-VAE中d是三位数。

引用：

Elijha：VQ-VAE解读(https://zhuanlan.zhihu.com/p/91434658)

Variational Autoencoders(https://amaires.github.io/VAE/)

变分自编码器（二）：从贝叶斯观点出发 – 科学空间|Scientific Spaces(https://spaces.ac.cn/archives/5343)

VQ-VAE的简明介绍：量子化自编码器 – 科学空间|Scientific Spaces(https://spaces.ac.cn/archives/6760)

简单得令人尴尬的FSQ：“四舍五入”超越了VQ-VAE – 科学空间|Scientific Spaces(https://www.spaces.ac.cn/archives/9826)

公众号后台回复“数据集”获取100+深度学习各方向资源整理