有限标量量化:简化向量量化的变分自编码器.

向量量化(Vector Quantize, VQ)是一种常用的图像Tokenizer技术,即把在连续实数空间的图像编码为一个整数的序列,与文本Tokenizer的输出形式一致,统一了不同模态数据的形式。

VQ-VAE模型建立一个字典存储一系列容量为$K$的编码表,在这个字典中找到和隐变量最接近的一个编码,用这个编码的index来代表这个隐变量。本文提出的有限标量量化(Finite Scalar Quantization, FSQ)则采用一种更简单的量化形式:

\[\hat{z} = \text{round} \left( \lfloor \frac{L}{2} \rfloor \tanh(z) \right)\]

其中$\tanh$函数把隐变量$z\in R$映射到$[-1,1]$,预先指定一个整数$L$,则上式通过四舍五入操作把输出限制在了$L$个整数之中,从而实现了离散化。对于$d$维向量$z$,被离散为$L^d$个整数之一。

由于四舍五入操作无法提供梯度,因此使用Straight-Through Estimator方法,梯度的直通估计是指前向传播的时使用目标变量(即使不可导),而反向传播时使用自己设计的梯度。即把四舍五入操作设计为:

\[x \leftarrow x + \text{sg}(\text{round}(x)-x)\]

其中$\text{sg}$表示stop gradient,即在反向传播时不计算其梯度,在pytorch中可以通过.detach()方法实现。

FSQ的编码数量为$K=L^d$。作者比较了VQ-VAEFSQ在具有相同编码数量的情况下的性能表现。结果表明当编码总数$K$比较小时,FSQ的效果通常不如VQ-VAE;编码表大小明显超过1000时,FSQ占优。