Towards a better understanding of Vector Quantized Autoencoders

1 Introduction . 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 ( Goodfellow et al. , 2014 ; van den Oord et al. , 2016 ; Kingma et al. , 2016 ; Salimans et al. , 2017 ; Parmar et al. , 2018 ) , audio ( Van Den Oord et al. , 2016 ; Reed et al. , 2017 ) , and video ( Kalchbrenner et al. , 2016 ) . 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. , Vinyals et al . ( 2015 ) . Improved discrete latent variable models could also prove useful for learning novel data compression algorithms ( Theis et al. , 2017 ) , while having far more interpretable representations of the data . We build on Vector Quantized Variational Autoencoder ( VQ-VAE ) ( van den Oord et al. , 2017 ) , 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 Oord et al . ( 2017 ) , 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 Oord et al. , 2017 ) . An extension of this method to conditional supervised generation , out-performs continuous autoencoders on WMT English-German translation task ( Kaiser et al. , 2018 ) . The work of Kaiser et al . ( 2018 ) 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 ( Norouzi & Fleet , 2013 ) 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 Kaiser et al . ( 2018 ) . We also exploit VQ-VAE ’ s connection with the expectation maximization ( EM ) algorithm ( Dempster et al. , 1977 ) , yielding additional improvements . With both improvements , we achieve a BLEU score of 22.4 on English to German translation , outperforming Kaiser et al . ( 2018 ) by 2.6 BLEU . Knowledge distillation ( Hinton et al. , 2015 ; Kim & Rush , 2016 ) 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 Oord et al . ( 2017 ) 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 ( Hinton et al. , 2015 ; Kim & Rush , 2016 ) , 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 ) . 2 VQ-VAE and the Hard EM Algorithm . The connection between K-means , and hard EM , or the Viterbi EM algorithm is well known ( Bottou & Bengio , 1995 ) , where the former can be seen a special case of hard-EM style algorithm with a mixture-of-Gaussians model with identity covariance and uniform prior over cluster probabilities . In the following sections we briefly explain the VQ-VAE discrete autoencoder for completeness and it ’ s connection to classical EM . 2.1 VQ-VAE discretization algorithm . VQ-VAE models the joint distribution PΘ ( x , z ) where Θ are the model parameters , x is the data point and z is the sequence of discrete latent variables or codes . Each position in the encoded sequence has its own set of latent codes . Given a data point , the discrete latent code in each position is selected independently using the encoder output . For simplicity , we describe the procedure for selecting the discrete latent code ( zi ) in one position given the data point ( xi ) . The encoder output ze ( xi ) ∈ RD is passed through a discretization bottleneck using a nearest-neighbor lookup on embedding vectors e ∈ RK×D . Here K is the number of latent codes ( in a particular position of the discrete latent sequence ) in the model . More specifically , the discrete latent variable assignment is given by , zi = arg min j∈ [ K ] ‖ze ( xi ) − ej‖2 ( 1 ) The selected latent variable ’ s embedding is passed as input to the decoder , zq ( xi ) = ezi The model is trained to minimize : L = lr + β ‖ze ( xi ) − sg ( zq ( xi ) ) ‖2 , ( 2 ) where lr is the reconstruction loss of the decoder given zq ( x ) ( e.g. , the cross entropy loss ) , and , sg ( . ) is the stop gradient operator defined as follows : sg ( x ) = { x forward pass 0 backward pass To train the embedding vectors e ∈ RK×D , van den Oord et al . ( 2017 ) proposed using a gradient based loss function ‖sg ( ze ( xi ) ) − zq ( xi ) ‖2 , ( 3 ) and also suggested an alternate technique of training the embeddings : by maintaining an exponential moving average ( EMA ) of all the encoder hidden states that get assigned to it . It was observed in Kaiser et al . ( 2018 ) that the EMA update for training the code-book embedding , results in more stable training than using gradient-based methods . We analyze this in more detail in Section 5.1.1 . Specifically , an exponential moving average is maintained over the following two quantities : 1 ) the embeddings ej for every j ∈ [ 1 , . . . , K ] and , 2 ) the count cj measuring the number of encoder hidden states that have ej as it ’ s nearest neighbor . The counts are updated in a mini-batch of targets as : cj ← λcj + ( 1− λ ) ∑ i 1 [ zq ( xi ) = ej ] , ( 4 ) with the embedding ej being subsequently updated as : ej ← λej + ( 1− λ ) ∑ i 1 [ zq ( xi ) = ej ] ze ( xi ) cj , ( 5 ) where 1 [ . ] is the indicator function and λ is a decay parameter which we set to 0.999 in our experiments . This amounts to doing stochastic gradient in the space of both code-book embeddings and cluster assignments . These techniques have also been successfully used in minibatch K-means ( Sculley , 2010 ) and online EM ( Liang & Klein , 2009 ; Sato & Ishii , 2000 ) . The generative process for our latent variable NMT model , P ( y , z | x ) , begins by autoregressively sampling a sequence of discrete latent codes from a model conditioned on the input x , P ( z | x ) = |z|∏ i=1 P ( zi | z1 , ... , ( i−1 ) , x ) , ( 6 ) which we refer to as the Latent Predictor model ( Kaiser et al. , 2018 ) . The decoder then consumes this sequence of discrete latent variables to generate the target y all at once , where P ( y | z , x ) = |y|∏ j=1 P ( yj | z , x ) . ( 7 ) The autoregressive learned prior prior is fit on the discrete latent variables produced by the autoencoder . Our goal is to learn a sequence of latents , that is much shorter than the targets , |z| |y| , thereby speeding up decoding significantly with no loss in accuracy . The architecture of the encoder , the decoder , and the latent predictor model are described in further detail in Section 5 . 2.2 Hard EM and the K-means algorithm In this section we briefly recall the hard Expectation maximization ( EM ) algorithm ( Dempster et al. , 1977 ) . Given a set of data points ( x1 , . . . , xN ) , the hard EM algorithm approximately solves the following optimization problem : Θ∗ = arg max Θ max z1 , ... , zN PΘ ( x1 , . . . , xN , z1 , . . . , zN ) , ( 8 ) Hard EM performs coordinate descent over the following two coordinates : the model parameters Θ , and the hidden variables z1 , . . . , zN . In other words , hard EM consists of repeating the following two steps until convergence : 1 . E step : ( z1 , . . . , zN ) ← arg maxz1 , ... , zN PΘ ( x1 , . . . , xN , z1 , . . . , zN ) , 2 . M step : Θ← arg maxΘ PΘ ( x1 , . . . , xN , z1 , . . . , zN ) A special case of the hard EM algorithm is K-means clustering ( MacQueen et al. , 1967 ; Bottou & Bengio , 1995 ) where the likelihood is modelled by a Gaussian with identity covariance matrix . Here , the means of the K Gaussians are the parameters to be estimated , Θ = 〈µ1 , . . . , µK〉 , µk ∈ RD . With a uniform prior over the hidden variables ( PΘ ( zi ) = 1K ) , the marginal is given by PΘ ( xi | zi ) = N ( µzi , I ) ( xi ) . In this case , equation ( 8 ) is equivalent to : ( µ1 , . . . , µK ) ∗ = arg max µ1 , ... , µK min z1 , ... , zN N∑ i=1 ‖µzi − xi‖22 ( 9 ) Note that optimizing equation ( 9 ) is NP-hard , however one can find a local optima by applying coordinate descent until convergence : 1 . E step : Cluster assignment is given by , zi ← arg min j∈ [ K ] ∥∥µj − xi∥∥22 , ( 10 ) 2 . M step : The means of the clusters are updated as , cj ← N∑ i=1 1 [ zi = j ] ; µ j ← 1 cj N∑ i=1 1 [ zi = j ] xi . ( 11 ) We can now easily see the connections between the training updates of VQ-VAE and K-means clustering . The encoder output ze ( x ) ∈ RD corresponds to the data point while the discrete latent variables corresponds to clusters . Given this , Equation 1 is equivalent to the E-step ( Equation 10 ) and the EMA updates in Equation 4 and Equation 5 converge to the M-step ( Equation 11 ) in the limit . The M-step in K-means overwrites the old values while the EMA updates interpolate between the old values and the M step update .

Summary: The paper presents an alternative view on the training procedure for the VQ-VAE. The authors have noticed that there is a close connection between the original training algorithm and the well-known EM algorithm. Then, they proposed to use the soft EM algorithm. In the experiments the authors showed that the soft EM allows to obtain significantly better results than the standard learning procedure on both image and text datasets.