2309.15505

paper

Motivation

由于codebook的查表操作是不可微的, 并且大规模codebook的利用率不高, 因此有codebook collapse的问题.

因此, 提出了FSQ作为VQ-VAE的替代, 舍弃了显式codebook查找和复杂的辅助loss函数, 通过简单的定界函数(bounding function)和取整来实现quantize

Methods

首先, 将原始的domain编码到feature domain: $z=\mathcal{E}(x)$

然后, 使用ste进行量化. ste的过程为:

1. 使用便捷映射函数. 此处使用tanh作为激活函数: $z_{bound,i}=\lfloor \frac{L_i}{2}\rfloor \cdot \tanh (z_i)$, 其中,
   • $z_i$是原始feature $z$的第$i$维度的分量
   • $L_i$是给定的超参数, 表示$i$-dim需要量化到多少个离散的值
   • $z_{bound,i}$是bounding之后的分量
2. 然后, 将$z_{bound,i}$取整, 得到离散的整数索引${\hat{z}}_i=\text{round}(z_{bound,i})$. 由于round函数大多数位置导数为0(不可导), 使用这种方法计算round: ${\hat{z}}_i=roun{d}_{s}te(z_{bound,i})=z_{bound,i}+\text{stop\_gradient}(\text{round}(z_{bound,i})-z_{bound,i})$

接着根据quantize之后的隐变量预测输入: $\hat{x}=\mathcal{D}(\hat{z})$

由于使用了round_ste, 能够将梯度回传, 因此loss函数可以直接写成${\nabla }_{\varphi }\mathcal{L}\approx {\nabla }_{\varphi }\Vert x-{\mathcal{D}}_{\theta }(\hat{z}){\Vert }^2$

Proof

FSQ表明了, 语义信息并不保存在codebook中, 而是保存在encoder和decoder里面:

将两张不同的图片的feature拼接, FSQ和VQ都能够将内容融合(边缘变化平滑, 纹理、光照连续变化). 这证明了图片的语义并不存在codebook中, 因此可以直接抛弃codebook