Deep neural networks with discrete latent variables offer the promise of better symbolic reasoning, and learning  abstractions that are more useful to new tasks. There has been a surge in interest in discrete latent variable models,  however, despite several recent improvements, the training of discrete latent variable models has remained  challenging and their performance has mostly failed to match their continuous counterparts. Recent work on vector quantized autoencoders (VQ-VAE) has made substantial progress in this direction, with its perplexity almost matching that of a VAE on datasets such as CIFAR-10. In this work, we investigate an alternate training technique for VQ-VAE, inspired by its connection to the Expectation Maximization (EM) algorithm. Training the discrete autoencoder with EM and combining it with sequence  level knowledge distillation alows us to develop a non-autoregressive machine translation model whose accuracy almost matches a strong greedy autoregressive baseline Transformer, while being 3.3 times faster at inference.
 Unsupervised learning of meaningful representations is a fundamental problem in machine learning since obtaining labeled data can often be very expensive. Continuous representations have largely been the workhorse of unsupervised deep learning models of images BID4 BID16 BID28 BID25 , audio BID27 , and video . However, it is often the case that datasets are more naturally modeled as a sequence of discrete symbols rather than continuous ones. For example, language and speech are inherently discrete in nature and images are often concisely described by language, see e.g., . Improved discrete latent variable models could also prove useful for learning novel data compression algorithms BID32 , while having far more interpretable representations of the data.We build on Vector Quantized Variational Autoencoder (VQ-VAE) , a recently proposed training technique for learning discrete latent variables. The method uses a learned code-book combined with nearest neighbor search to train the discrete latent variable model. The nearest neighbor search is performed between the encoder output and the embedding of the latent code using the 2 distance metric. VQ-VAE adopts the standard latent variable model generative process, first sampling latent codes from a prior, P (z), which are then consumed by the decoder to generate data from P (x | z). In van den , the authors use both uniform and autoregressive priors for P (z). The resulting discrete autoencoder obtains impressive results on unconditional image, speech, and video generation. In particular, on image generation, VQ-VAE was shown to perform almost on par with continuous VAEs on datasets such as CIFAR-10 (van den ). An extension of this method to conditional supervised generation, out-performs continuous autoencoders on WMT English-German translation task .The work of introduced the Latent Transformer, which set a new stateof-the-art in non-autoregressive Neural Machine Translation. However , additional training heuristics, namely, exponential moving averages (EMA) of cluster assignment counts, and product quantization BID24 were essential to achieve competitive results with VQ-VAE. In this work, we show that tuning for the code-book size can significantly outperform the results presented in . We also exploit VQ-VAE's connection with the expectation maximization (EM) algorithm BID3 , yielding additional improvements. With both improvements, we achieve a BLEU score of 22.4 on English to German translation, outperforming by 2.6 BLEU. Knowledge distillation BID7 BID12 ) provides significant gains with our best models and EM, achieving 26.7 BLEU, which almost matches the autoregressive transformer model with no beam search at 27.0 BLEU, while being 3.3× faster.Our contributions can be summarized as follows:1. We show that VQ-VAE from van den can outperform previous state-of-the-art without product quantization. 2. Inspired by the EM algorithm, we introduce a new training algorithm for training discrete variational autoencoders, that outperforms the previous best result with discrete latent autoencoders for neural machine translation. 3. Using EM training, and combining it sequence level knowledge distillation BID7 BID12 , allows us to develop a non-autoregressive machine translation model whose accuracy almost matches a strong greedy autoregressive baseline Transformer, while being 3.3 times faster at inference. 4. On the larger English-French dataset, we show that denoising discrete autoencoders gives us a significant improvement (1.0 BLEU) on top of our non-autoregressive baseline (see Section D). We investigate an alternate training technique for VQ-VAE inspired by its connection to the EM algorithm. Training the discrete autoencoder with EM and combining it with sequence level knowledge distillation, allows us to develop a non-autoregressive machine translation model whose accuracy almost matches a greedy autoregressive baseline, while being 3.3 times faster at inference. While sequence distillation is very important for training our best model, we find that the improvements from EM on harder tasks is quite significant. We hope that our results will inspire further research on using vector quantization for fast decoding of autoregressive sequence models.
Understand the VQ-VAE discrete autoencoder systematically using EM and use it to design non-autogressive translation model matching a strong autoregressive baseline. This paper introduces a new way of interpreting the VQ-VAE and proposes a new training algorithm based on the soft EM clustering. The paper presents an alternative view on the training procedure for the VQ-VAE using the soft EM algorithm