Neural Discrete Representation Learning (VQ-VAE) 简介

VQ-VAE

VAE是一种生成模型。
Vector QuantisedVariational AutoEncoder (VQ-VAE)是VAE的变种，其隐含变量是离散的。离散的隐含变量对于自然语言，推理都比较有帮助。著名的DALL-E就使用了类似VQ-VAE的离散隐含变量来从文本生成图像，DALL-E对VQ-VAE进行了一些改进。

VQ-VAE通过vector quantisation (VQ) 将隐含变量离散化。
假设$\mathbf{e}\in {\mathbb{R}}^{K\times D}$是codebook。其中$K$是codebook中embeddings的个数，$D$是codebook的维度。${\mathbf{e}}_i$是其中一个embedding。
encoder的输出$E(\mathbf{x})={\mathbf{z}}_e$将通过最近邻查找的方式找到自己属于的embedding向量。这个embedding向量通过decoder$D(.)$的输出将尽可能与$\mathbf{x}$相似。
$${\mathbf{z}}_q(\mathbf{x})=\text{Quantize}(E(\mathbf{x}))={\mathbf{e}}_k\text{ where }k=\arg \underset{i}{\min }\Vert E(\mathbf{x})-{\mathbf{e}}_i{\Vert }_2$$

VQ-VAE优化的是下面的目标：
$$L=\underset{\text{reconstruction loss}}{\Vert \mathbf{x}-D({\mathbf{e}}_k){\Vert }_2^2}+\underset{\text{VQ loss}}{\Vert \text{sg}[E(\mathbf{x})]-{\mathbf{e}}_k{\Vert }_2^2}+\underset{\text{commitment loss}}{\beta \Vert E(\mathbf{x})-\text{sg}[{\mathbf{e}}_k]{\Vert }_2^2}$$其中sq表示stop_gradient。

codebook中的embedding向量使用EMA (exponential moving average)学习。
$$N_i^{(t)}=\gamma N_i^{(t-1)}+(1-\gamma )n_i^{(t)}{\mathbf{m}}_i^{(t)}=\gamma {\mathbf{m}}_i^{(t-1)}+(1-\gamma )\sum _{j=1}^{n_i^{(t)}}{\mathbf{z}}_{i,j}^{(t)}{\mathbf{e}}_i^{(t)}={\mathbf{m}}_i^{(t)}/N_i^{(t)}$$

参考

lilianweng.github.io VQ-VAE