Neural Discrete Representation Learning

摘要

学习无监督的有用表示仍然是机器学习中的关键挑战。在本文中，我们提出了一种简单而强大的生成模型，可以学习这种离散表示。我们的模型，被称为Vector Quantised-Variational AutoEncoder (VQ-VAE)，在两个方面不同于VAE：编码器网络输出离散code，而不是连续code；先验是学习到的而不是静态的。为了学习离散的潜在表示，我们将向量量化（VQ）的思想纳入其中。使用VQ方法允许模型来避免“后验崩塌”的问题，即当潜在表示与一个强大的自回归解码器配对时，通常会被忽略的问题，这通常在VAE框架中观察到。将这些VQ表式与自回归一起使用，该模型可以生成图像，视频和演讲的质量，同时能做到高质量的speaker转换和无监督的音素学习。

1.介绍

最近在图像，音频和视频的生成式建模上进步已经产生令人印象深刻的样例和应用。同时，相应地具有挑战性的任务，如小样本学习，领域自适应或强化学习严重依赖于从原始数据中学习到的表示，但以无监督方式训练的表示的有用性仍然远差与主流的方法。
　　最大似然和重构误差是用于在像素领域中训练无监督模型的两个常见目标，但它们只能获取到特定应用所关注的特征。我们的目标是实现一种模型，能够在潜在空间中保留数据的重要特征的同时，来优化最大似然。工作[7]中认为，最好的生成模型（通过对数似然测量）将是没有潜在表示，而具有强大的解码器（例如PixelCNN）。相反，在本文中，我们认为学习离散的有用潜在表示是非常重要的。
　　许多以前的工作主要关注于学习具有连续特征的表示，但我们关注于更加适合许多场景的离散表示。语言本质上是离散的，类似地语音通常表示为一个符号序列。图像通常可以通过语言简明扼要地描述。此外，离散表示对于复杂的推理，规划和预测学习是一种自然的契合（例如，“ if it rains, I will use an umbrella)”）。在深度学习中使用离散潜在表示的同时，已经开发了强大的自回归模型，用于在离散表示上建模分布。
　　在我们的工作中，我们引入了一种新的生成模型框架，给定一个观测值，通过（离散）潜在表示的后验分布的新参数化，将变分自编码器（VAE）框架与离散潜在表示相结合。我们的模型依赖于vector quantization（VQ），且易于训练，不会遭受大的偏差，从而避免“后验崩塌”问题。此外，它是第一个离散潜在表示的VAE模型，与相应的连续模型相比，能得到相似的性能，同时提供离散分布的灵活性。 我们将我们的模型称为VQ-VAE。
　　由于VQ-VAE可以有效地利用潜在空间，因此可以成功建模通常跨越数据空间中许多维度的重要特征（例如，跨越图像中的许多像素的目标，语音中的音素，文本片段中的消息等），而不是仅聚焦于某些局部特征。
　　最后，一旦VQ-VAE发现了一种良好的离散潜在结构，我们将在这些离散的随机变量code上训练一个强大的先验，从而产生了有趣的样本和有用的应用。例如，当在语音上训练时，我们发现了语言的潜在结构，这是没有任何监督或关于音素或单词的知识的。此外，我们可以用speaker标识带入我们的解码器，从而允许speaker迁移，即，在不改变内容的情况下将语音从一个speaker传输到另一个speaker。我们还表明了该框架在RL学习长期环境中有希望的结果。
　　因此，我们的贡献可以归纳为：

• 引入了简单的VQ-VAE模型，使用离散潜在表示，不会遭受“后验崩塌”且没有方差问题。
• 我们表明，在连续模型的对数似然方法上，同样可以执行离散的潜在模型（VQ-VAE）。
• 当与强大的先验联合使用时，我们模型在各种应用场景能生成高质量的样本，例如语音和视频。
• 我们通过原始演讲显示出学习语言的证据，没有任何监督，并显示无监督者转换的应用。

2.相关工作

在这项工作中，我们提出了用离散潜在变量训练变分自编码器的新方法。在深度学习中使用离散变量已经证明具有挑战性，当前大部分工作中都是连续潜变量建模占主导地位，即使某些应用方法是离散的。
　　目前存在许多训练离散VAE的可替代方案。NVIL估计器使用单个样本目标来优化变分的下限，并使用各种方差减少技术来加速训练。VIMCO优化了多样本目标，通过使用来自推理网络的多个样本进一步加速收敛。
　　最近，一些作者提出了基于所谓的Concrete或Gumbel-Softmax分布的新的连续重参数化方法，这是一种连续分布，并且在训练期间可以使用一个温度系数来进行退火，以收敛到离散的分布。在训练开始时，梯度的方差小但偏差高，并且朝向训练结束时，方差变得高但无偏差。
　　然而，上述方法都没法达到与使用高斯重参数化技巧的连续VAE相同的性能。此外，大多数这些技术通常在诸如MNIST等的相对较小的数据集上进行评估，并且潜在分布的维度小（例如，低于8）。在我们的工作中，我们使用了三个复杂的图像数据集（CIFAR10，Imagenet和DeepMind实验室）和原始语音数据集（VCTK）。
　　我们的工作还扩展了在VAE的解码器或先验中使用自回归分布的研究。一些研究使用LSTM解码器进行语言建模，另外一些则使用dilated convolutional解码器。PixelCNNs是卷积自回归模型，其也被用作VAE解码器中的分布。
　　最后，我们的方法还涉及用神经网络的图像压缩工作。Theis et. al.使用标量量化在算术编码之前压缩有损图像的激活。其他作者也提出一种具有向量量化的类似压缩模型的方法。
　　这些工作中提出了向量量化的连续松弛，随着时间的推移而退火，以获得一个硬聚类。在他们的实验中，他们首先训练一个AutoEncoder，然后将向量量化应用于编码器的激活，最后，使用具有小的学习率的soft-to-hard松弛，整个网络进行微调。在我们的实验中，由于解码器总是能够在训练期间颠覆连续松弛，因此我们无法从头使用soft-to-hard松弛方法，因此没有发生实际量化。

3.VQ-VAE

与我们的方法相关的工作是VAEs。VAE由以下部分组成：编码器网络，给定输入数据$x$，一个先验分布p(z)，由离散潜在随机变量参数化一个后验分布$q(z|x)$；一个具有输入数据分布$p(x|z)$的解码器网络。
　　通常，VAE中的后验和先验通常被假设为具有对角方差的正态分布，这允许使用高斯重参数化的技巧。VAE的相关扩展工作有：autoregressive prior and posterior models, normalising flows, and inverse autoregressive posteriors。
　　在这项工作中，我们引入了VQ-VAE，受到向量量化（VQ）的启发，我们使用具有新训练方式的离散潜在变量。后验和先验的分布是可分类的，并且用从这些分布中获取的样本来索引嵌入表。然后将这些嵌入用作解码器网络中的输入。

3.1 Discrete Latent variables

　　我们定义一个潜在嵌入空间$e\in R^{K\times D}$，其中$K$是离散潜在空间的大小（即，K-way 类），$D$是每个潜在嵌入向量$e_i$的维度。请注意，一共有$K$个嵌入向量$e_i\in R^D,i\in 1,2,...,K$。如图1所示，模型以$x$作为输入，该输入$x$通过编码器产生输出$z_e(x)$。然后，使用相同的嵌入空间$e$通过最近邻查找来计算离散潜变量$z$，如公式1中所示。解码器的输入是相应的嵌入向量$e_k$，如等式2中给出的。将一个前向计算管道作为常规自编码器，具有特定的非线性，将潜在表示映射到嵌入向量。该模型的完整参数集是编码器，解码器和嵌入空间$e$的参数。为简单起见，我们使用单个随机变量$z$表示本节中的离散潜在变量，但是对于语音，图像和视频我们实际上分别提取1D，2D和3D潜在特征空间。
　　后验分类分布$q(z|x)$概率定义为如下所示的one-hot表示：
$$q(z=k|x)=\{\begin{array}{ll}1& fork=argmi{n}_j||z_e(x)-e_j|{|}_2,\\ 0& otherwise\end{array},\text{\hspace{1em}(1)}$$
其中$z_e(x)$是编码器网络的输出。我们将此模型视为一个VAE，我们可以将$logp(x)$与ELBO绑定。我们提出的分布$q(z=k|x)$是确定性的，并且通过在$z$上定义一个简单的均匀先验，我们获得KL散度常数并等于$logK$。
　　表示$z_e(x)$通过离散瓶颈，然后映射到嵌入$e$中最近邻元素上，如等式1和2所示。
$$z_q(x)=e_k,wherek=argmi{n}_j||z_e(x)-e_j|{|}_2\text{\hspace{1em}(2)}$$

3.2 Learning

注意，对于等式2并没有定义真正的梯度，但是我们近似与直通估计器类似的梯度，并且只需将梯度从解码器输入$z_q(x)$复制到编码器输出$z_e(x)$。人们还可以通过量化操作使用子梯度，但是这个简单的估计器在本文的初始实验中工作得就很好。
　　在前向计算期间，最近邻的嵌入$z_q(x)$（等式2）传递给解码器，并且在反向传播期间，梯度${\nabla }_{z}L$直接复制到编码器。由于编码器的输出表示和解码器的输入共享相同的$D$维空间，因此梯度包含编码器如何改变其输出以降低重构损失的有用信息。
　　如图1（右）所示，梯度可以将编码器的输出推到另外一个离散嵌入，因为等式1中的分配是不同的。
　　等式3指定了一个整体的损失函数。它具有三个组件，用于训练VQ-VAE的不同部分：
　　第一项是优化解码器和编码器（通过上面说明的估计器）的重构损失（或称为数据项）。由于从$z_e(x)$到z_q(x)映射的直接梯度估计，嵌入$e_i$从重构损失$logp(x|z_q(x))$中没有接收梯度。因此，为了学习嵌入空间，我们使用最简单的字典学习算法之一，Vector Quantisation (VQ)。
　　VQ目标使用$l_2$误差将嵌入向量$e_i$沿着编码器输出$z_e(x)$移动，如等式的第二项中所示。因为该损失项仅用于更新字典，所以$z_e(X)$的移动平均函数同样可用于替代$l_2$来更新字典（在这项工作中没有用于实验）。有关详细信息，请参阅附录A.1。
　　最后，由于嵌入空间的容量是无量纲的，因此如果嵌入$e_i$不和编码器参数以相同速度来训练，则可以任意地增长。为了确保编码器提交到嵌入及其它的输出不会增长，我们增加了一个commitment损失，在等式3中作为第三个损失项。因此，总训练目标成为：
$$L=logp(x|z_q(x))+||sg[z_e(x)]-e|{|}_2^2+\beta ||z_e(x)-sg[e]|{|}_2^2\text{\hspace{1em}(3)}$$
其中$sg$代表stopgradient运算符，该运算符被定义为前向计算时间的标识，并且具有零导数，从而有效地限制其操作数量是非更新的常量。解码器仅优化第一个损失项，编码器优化第一个损失项和最后一个损失项，并且离散嵌入空间由第二个损失项进行优化。我们发现算法对$\beta$的选择非常鲁棒，因为训练结果没有因$\beta$值从0.1至2.0的变化而变化。我们在所有实验中使用$\beta =0.25$，但一般而言，这将取决于重构损失的规模。由于我们假设$z$服从均匀先验，因此通常在ELBO中出现的KL项是一个和编码器参数相关的常数，因此可以忽略。
　　在我们的实验中，我们定义了$N$个离散潜在变量（例如，对于ImageNet，我们使用32 x 32潜在值，对于CIFAR10则是8 x 8 x 10）。除了在N个k-means项的平均和commitment损失上，其余结果损失$L$是相同的。
　　完整模型的对数概率$logp(x)$可以如下估价：
$$logp(x)=log\sum _{k}p(x|z_k)p(z_k),$$
这是因为解码器$p(x|z)$使用从MAP-inference得到的$z=z_q(x)$训练，所以一旦它完全收敛，解码器不应将任何概率分配给满足$z\ne z_q(x)$的$p(x|z)$。因此，我们可以写为$logp(x)\approx logp(x|z_q(x))p(z_q(x))$。我们在第4节中经验性地评估了这种近似值。根据Jensen不等式，我们也可以写为$logp(x)\ge logp(x|z_q(x))p(z_q(x))$。

3.3 Prior

离散潜在值$p(z)$的先验分布是一个分类分布，并且可以根据特征图中的其他$z$进行自回归。虽然训练VQ-VAE，但是先验保持为均匀分布。在训练之后，我们拟合$z,p(z)$的自回归分布，因此我们可以通过ancestral采样生成$x$。我们在离散潜在变量上为，针对图像使用PixelCNN，针对原始音频使用WaveNet。先验和VQ-VAE共同训练，同样可以加强我们的结果，这留作未来的研究。