Neural Discrete Representation Learning (VQ-VAE) 简介

VQ-VAE

VAE（Variational AutoEncoder ）是一种生成模型，VQ-VAE（Vector Quantised Variational AutoEncoder ）是VAE的变种。
VAE的隐含变量是连续的，符合高斯分布，而VQ-VAE的隐含变量是离散的。离散的隐含变量对于自然语言，推理都比较有帮助。著名的DALL-E就使用了类似VQ-VAE的离散隐含变量来从文本生成图像。

VQ-VAE通过vector quantisation (VQ) 将隐含变量离散化。
假设$\mathbf{e}\in {\mathbb{R}}^{K\times D}$是codebook。其中$K$是codebook中embeddings的个数，$D$是embedding的维度。${\mathbf{e}}_i$是其中一个embedding。
encoder的输出$E(\mathbf{x})={\mathbf{z}}_e$将通过最近邻查找的方式找到自己属于的embedding向量${\mathbf{e}}_k$:
$${\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$$并且要求embedding向量${\mathbf{e}}_k$通过decoder$D(.)$的输出将尽可能与$\mathbf{x}$相似。
于是有VQ-VAE的优化目标：
$$L_{VQVAE}=\underset{\text{reconstruction loss}}{\underbrace{\Vert \mathbf{x}-D({\mathbf{e}}_k){\Vert }_2^2}}+\underset{\text{VQ loss}}{\underbrace{\Vert \text{sg}[E(\mathbf{x})]-{\mathbf{e}}_k{\Vert }_2^2}}+\underset{\text{commitment loss}}{\underbrace{\beta \Vert E(\mathbf{x})-\text{sg}[{\mathbf{e}}_k]{\Vert }_2^2}}$$其中sq表示stop_gradient。

普通VAE优化的目标是最小化负ELBO（即最大化ELBO）：$$L_{VAE}=-{\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}|\mathbf{x})}\log p_{\theta }(\mathbf{x}|\mathbf{z})+D_{\text{KL}}(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p_{\theta }(\mathbf{z}))$$由于VQ-VAE假设先验$\mathbf{z}$是均匀分布，$p_{\theta }(\mathbf{z})=\frac{1}{K}$，$q_{\varphi }(\mathbf{z}|\mathbf{x})$是中只有一维为1，其余为0。
$$D_{\text{KL}}(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p_{\theta }(\mathbf{z}))=\sum _{\mathbf{z}}{q}_{\varphi }(\mathbf{z}|\mathbf{x})\ln \frac{q_{\varphi }(\mathbf{z}|\mathbf{x})}{p_{\theta }(\mathbf{z})}={\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\ln \frac{q_{\varphi }(\mathbf{z}|\mathbf{x})}{p_{\theta }(\mathbf{z})}={\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\ln K{q}_{\varphi }(\mathbf{z}|\mathbf{x})$$ELBO中的KL散度项是常数，因此KL散度项在训练时可以忽略。

为了使用batch的方式更新codebook，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)}$$其中$\gamma \in (0,1)$，$n_i^{(t)}$是第t个batch中属于${\mathbf{e}}_i$的数量（最近邻选择的是${\mathbf{e}}_i$），$N_i^{(t)}$是第t次更新时属于${\mathbf{e}}_i$的累计计数，${\mathbf{m}}_i^{(t)}$是第t次更新时属于${\mathbf{e}}_i$的累计向量。

图像生成

一幅图像并不是经过编码器变为一个$D$维的embedding，而是用卷积编码为$m\times m$个$D$维的embedding。这样可以保留图像的空间结构，并且用小的码本表示大量多变的图像。每个embedding对应于codebook中的一个位置，所以图像可以表示为$m\times m$的整数矩阵，这就实现了图像的离散编码。
训练过程中，隐含表示$\mathbf{z}$的服从均匀分布。在训练后，作者使用了自回归的方式拟合$m\times m$个$\mathbf{z}$。这样这些embedding之间是存在关系的，而不是独立的。具体地，作者对图像生成任务使用PixelCNN来生成隐含表示。
有点LDMs（Latent Diffusion Models）的感觉，隐含表示用PixelCNN生成，编码器和解码器用VQ-VAE。而LDMs则是隐含表示用diffusion model生成，编码器解码器用VAE。

参考

lilianweng.github.io VQ-VAE
https://www.spaces.ac.cn/archives/6760