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 散度具有两个重要性质：
-
不对称性
：显然，K–L 散度对于
和
来说是不对称的。
-
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
paper：Neural Discrete Representation Learning
背景：
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. VQ-VAE 流程图
-
-
codebook 是一个K*D 的 table（紫色方块）：
-
-
从
到
的变化可以理解为聚类，如图中右子图所示，由于变化后的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，例如从原始图像的256
256变为8
8。这样
进一步增加下采样因子
减少分辨率，使得 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 时有两个优势：
-
使用单独的码本需要广泛的超参数搜索，以确定每一层的码本大小，而共享码本只需确定总码本大小 K。
-
共享码本使得所有的 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)：
