2203.01941

paper

Motivation

为了增强VAE的生成质量, 通常的方法是增加空间分辨率(即, feature的分辨率). 但是如果使用自回归生成, 增加feature的size会导致建模长度指数级增加, 计算成本过高. 因此RQ-VAE提出: 不再增加feature size, 而是增加quantize的深度.

Method

RVQ

1. 首先将原始的图片经过Encoder, 得到feature: $z=\mathcal{E}(I)$. 其中:
   • $I$是原始的图片
   • $\mathcal{E}$是encoder, 在这里使用了ResNet的架构, 使用CNN+Residual Blocks
   • $z$是隐空间向量, 或者说叫做feature vector
2. 初始残差: $r_0=z$
3. 对于每一层(深度, 用$d$表示), 首先计算quantize: $k_d={\mathrm{argmin}}_{j\in \{1,\cdots ,K\}}\Vert r_{d-1}-e(j){\Vert }^2$, 其中:
   1. $e(j)$表示codebook中第$j$个向量. 实际上, 用$e_j$表示更加合理
   2. $K$是codebook的大小
   3. $r_d$表示残差, 是深度为$d$的时候的残差, 计算方法为$r_d=r_{d-1}-e(k_d)$
   4. $k_d$表示, 在深度为$d$的时候, 应该使用codebook中序号为$k_d$的基向量
4. 最终, 将所有选中的基向量进行加和, 得到最终的重建结果: $\hat{z}={\sum }_{d=1}^{D}e(k_d)$

这个方法类似:

假设有$\{100,50,20,10,5,1,0.5,0.1\}$这几种面值的钱币, 需要凑出来173.6, 应该如何做:

1. 第一步quantize选择$e(k_1)=100$, 残差为$r_1=73.6$
2. 第二步quantize选择$e(k_2)=50$, 残差为$r_2=23.6$
3. …

RQ-Transformer

至于推理, 提出了RQ-Transformer架构. 使用两个Transformer, 一个生成当前patch的隐变量$h_t$, 一个自回归生成当前隐变量对应的基向量序号序列$s_{t,d}{|}_{d=1,\cdots ,D}$. 然后, 根据序号求和得到最终的feature ${\hat{z}}_t={\sum }_{d=1}^{D}e(s_{t,d})$.

优点:

1. 从粗糙到精细的建模方式
2. 显著减少feature的size, 降低compute cost
3. codebook复用率高, 无需很大(但仍然效率低, 利用率不高)

缺点:

1. RQ-Transformer建模会有误差累积
2. 不能仅依赖reconstruction loss进行训练, 会导致重建保真度降低, 需要引入类似GAN的对抗损失进行训练
3. 训练时间成本高