Original: Towards a better understanding of Vector Quantized Autoencoders
Revised: Towards a better understanding of Vector Quantized Autoencoders

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: Abstract
Revised: Abstract

Original: Deep neural networks with discrete latent variables oﬀer the promise of better symbolic reasoning, and learning abstractions that are more useful to new tasks.
Revised: Deep neural networks with discrete latent variables oﬀer the promise of better symbolic reasoning, and learning abstractions that are more useful to new tasks.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In this work, we investigate an alternate training technique for VQ-VAE, inspired by its connection to the Expectation Maximization (EM) algorithm.
Revised: In this work, we investigate an alternate training technique for VQ-VAE, inspired by its connection to the Expectation Maximization (EM) algorithm.

Original: Training the discrete bottleneck with EM helps us achieve better image reconstruction results on SVHN together with  knowledge distillation, 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.
Revised: Training the discrete autoencoder with EM and combining it with sequence level knowledge distillation alows us to develop a nonautoregressive machine translation model whose accuracy almost matches a strong greedy autoregressive baseline Transformer, while being 3 . 3 times faster at inference.

Original:  Introduction
Revised:  Introduction

Original: Unsupervised learning of meaningful representations is a fundamental problem in machine learning since obtaining labeled data can often be very expensive.
Revised: Unsupervised learning of meaningful representations is a fundamental problem in machine learning since obtaining labeled data can often be very expensive.

Original: Continuous representations have largely been the workhorse of unsupervised deep learning models of images (5; 33; 14; 27; 23), audio (32; 25) and video (11).
Revised: 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).

Original: However, it is often the case that datasets are more naturally modeled as a sequence of discrete symbols rather than continuous ones.
Revised: However, it is often the case that datasets are more naturally modeled as a sequence of discrete symbols rather than continuous ones.

Original: For example, language and speech are inherently discrete in nature and images are often concisely described by language, see e.g., (36).
Revised: For example, language and speech are inherently discrete in nature and images are often concisely described by language, see e.g., Vinyals et al.

Original: Improved discrete latent variable models could also prove useful for learning novel data compression algorithms (31) while having far more interpretable representations of the data.
Revised: 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.

Original: We build on Vector Quantized Variational Autoencoder (VQ-VAE) (34), a recently proposed training technique for learning discrete latent variables.
Revised: 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.

Original: The method uses a learned code-book combined with nearest neighbor search to train the discrete latent variable model.
Revised: The method uses a learned code-book combined with nearest neighbor search to train the discrete latent variable model.

Original: The nearest neighbor search is performed between the encoder output and the embedding of the latent code using the (cid:96) 2 distance metric.
Revised: The nearest neighbor search is performed between the encoder output and the embedding of the latent code using the (cid:96) 2 distance metric.

Original: The generative process begins by sampling a sequence of discrete latent codes from an autoregressive model ﬁtted on the encoder latents, acting as a learned prior.
Revised: The generative process begins by sampling a sequence of discrete latent codes from an autoregressive model ﬁtted on the encoder latents, acting as a learned prior.

Original: The discrete latent sequence is then consumed by the decoder to generate data.
Revised: The discrete latent sequence is then consumed by the decoder to generate data.

Original: The resulting discrete autoencoder obtains impressive results on unconditional image, speech, and video generation.
Revised: The resulting discrete autoencoder obtains impressive results on unconditional image, speech, and video generation.

Original: In particular, on image reconstruction, our performance is almost on par with continuous VAEs on datasets such as CIFAR-10 (34).
Revised: 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).

Original: An extension of this method to conditional supervised generation, out-performs continuous autoencoders on WMT English-German translation task (10).
Revised: An extension of this method to conditional supervised generation, out-performs continuous autoencoders on WMT English-German translation task (Kaiser et al., 2018).

Original: The work of (10) introduced the Latent Transformer, which achieved impressive results using discrete autoencoders for fast neural machine translation.
Revised: The work of Kaiser et al. (2018) introduced the Latent Transformer, which set a new stateof-the-art in non-autoregressive Neural Machine Translation.

Original: However, additional training heuristics, namely, exponential moving averages (EMA) of cluster assignment counts, and product quantization (22) were essential to achieve competitive results with VQ-VAE.
Revised: 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.

Original: In this work, we show that tuning for the code-book size can signiﬁcantly outperform the results presented in (10).
Revised: In this work, we show that tuning for the code-book size can signiﬁcantly outperform the results presented in Kaiser et al.

Original: We also exploit VQ-VAE’s connection with the expectation maximization 
Revised: We also exploit VQ-VAE’s connection with the expectation maximization (EM) algorithm

Original: (EM) algorithm (4), yielding additional improvements.
Revised: (Dempster et al., 1977), yielding additional improvements.

Original: With both improvements, we achieve a BLEU score of 22 . 4 on English to German translation, outperforming (10) by 2 . 6 BLEU.
Revised: 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.

Original: Knowledge distillation (7; 12) provides signiﬁcant gains with our best models and EM, achieving 26 .
Revised: Knowledge distillation (Hinton et al., 2015; Kim & Rush, 2016) provides signiﬁcant gains with our best models and EM, achieving 26 .

Original: 7 BLEU, which almost matches the autoregressive transformer model with no beam search at 27 . 

Revised: 7 BLEU, which almost matches the autoregressive transformer model with no beam search at 27 .

Original: BLEU, while being 3 . 3 × faster.
Revised: BLEU, while being 3 . 3 × faster.

Original: Our contributions can be summarized as follows:
Revised: Our contributions can be summarized as follows:

Original: We show that VQ-VAE from (34) can outperform previous state-of-the-art without product quantization.
Revised: We show that VQ-VAE from van den Oord et al. (2017) can outperform previous state-of-the-art without product quantization.

Original: 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.
Revised: 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.

Original: Using EM training,we achieve better reconstructions on the SVHN, and with the additional use of knowledge distillation, 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.
Revised: 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.

Original: On the larger English-French dataset, we show that denoising discrete autoencoders gives us a signiﬁcant improvement (1.0 BLEU) on top of our non-autoregressive baseline (see Section D).
Revised: On the larger English-French dataset, we show that denoising discrete autoencoders gives us a signiﬁcant improvement (1.0 BLEU) on top of our non-autoregressive baseline (see Section D).

Original:  VQ-VAE and the Hard EM Algorithm
Revised:  VQ-VAE and the Hard EM Algorithm

Original: The connection between K -means, and hard EM, or the Viterbi EM algorithm is well known (2), 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.
Revised: 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.

Original: In the following sections we brieﬂy explain the VQ-VAE discretization algorithm for completeness and it’s connection to classical EM.
Revised: In the following sections we brieﬂy explain the VQ-VAE discrete autoencoder for completeness and it’s connection to classical EM.

Original: 2. VQ-VAE discretization algorithm
Revised: 2. VQ-VAE discretization algorithm

Original: 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.
Revised: 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.

Original: Each position in the encoded sequence has its own set of latent codes.
Revised: Each position in the encoded sequence has its own set of latent codes.

Original: Given a data point, the discrete latent code in each position is selected independently using the encoder output.
Revised: Given a data point, the discrete latent code in each position is selected independently using the encoder output.

Original: For simplicity, we describe the procedure for selecting the discrete latent code ( z i ) in one position given the data point ( x i ).
Revised: For simplicity, we describe the procedure for selecting the discrete latent code ( z i ) in one position given the data point ( x i ).

Original: The encoder output z e ( x i
Revised: The encoder output z e ( x i

Original: ) ∈
Revised: ) ∈

Original: R D is passed through a discretization bottleneck using a nearest-neighbor lookup on embedding vectors e ∈ R K × D .
Revised: R D is passed through a discretization bottleneck using a nearest-neighbor lookup on embedding vectors e ∈ R K × D .

Original: Here K is the number of latent codes (in a particular position of the discrete latent sequence) in the model.
Revised: Here K is the number of latent codes (in a particular position of the discrete latent sequence) in the model.

Original: More speciﬁcally, the discrete latent variable assignment is given by,
Revised: More speciﬁcally, the discrete latent variable assignment is given by,

Original: [Equation]
Revised: [Equation]

Original: The selected latent variable’s embedding is passed as input to the decoder,
Revised: The selected latent variable’s embedding is passed as input to the decoder,

Original: [Equation]
Revised: [Equation]

Original: The model is trained to minimize:
Revised: The model is trained to minimize:

Original: [Equation]
Revised: [Equation]

Original:  where l r is the reconstruction loss of the decoder given z q ( x ) (e.g., the cross entropy loss), and, sg ( . ) is the stop gradient operator deﬁned as follows:
Revised: (2) where l r is the reconstruction loss of the decoder given z q ( x ) (e.g., the cross entropy loss), and, sg ( . ) is the stop gradient operator deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: To train the embedding vectors e ∈ R K × D , (34) proposed using a gradient based loss function
Revised: To train the embedding vectors e ∈ R K × D , van den Oord et al. (2017) proposed using a gradient based loss function

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: It was observed in (10) that the EMA update for training the code-book embedding, results in more stable training than using gradient-based methods.
Revised: 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.

Original: We analyze this in more detail in Section 5.1.1.
Revised: We analyze this in more detail in Section 5.1.1.

Original: Speciﬁcally, an exponential moving average is maintained over the following two quantities: 1) the embeddings e j for every j ∈ [1 , . . .
Revised: Speciﬁcally, an exponential moving average is maintained over the following two quantities: 1) the embeddings e j for every j ∈ [1 , . . .

Original: , K ] and, 2) the count c
Revised: , K ] and, 2) the count c

Original: j measuring the number of encoder hidden states that have e j as it’s nearest neighbor.
Revised: j measuring the number of encoder hidden states that have e j as it’s nearest neighbor.

Original: The counts are updated in a mini-batch of targets as:
Revised: The counts are updated in a mini-batch of targets as:

Original: [Equation]
Revised: [Equation]

Original: with the embedding e j being subsequently updated as:
Revised: with the embedding e j being subsequently updated as:

Original: [Equation]
Revised: [Equation]

Original: where 1 [ . ] is the indicator function and λ is a decay parameter which we set to 0 .
Revised: where 1 [ . ] is the indicator function and λ is a decay parameter which we set to 0 .

Original: 999 in our experiments.
Revised: 999 in our experiments.

Original: This amounts to doing stochastic gradient in the space of both code-book embeddings and cluster assignments.
Revised: This amounts to doing stochastic gradient in the space of both code-book embeddings and cluster assignments.

Original: These techniques have also been successfully used in minibatch K -means (29) and online EM (17; 28).
Revised: These techniques have also been successfully used in minibatch K -means (Sculley, 2010) and online EM (Liang & Klein, 2009; Sato & Ishii, 2000).

Original: The generative process  begins by sampling a sequence of discrete latent codes from an autoregressive model, which we refer to as the Latent Predictor model.
Revised: The generative process for machine translation begins by sampling a sequence of discrete latent codes from an autoregressive model, which we refer to as the Latent Predictor model.

Original: The decoder then consumes this sequence of discrete latent variables to generate the data.
Revised: The decoder then consumes this sequence of discrete latent variables to generate the data.

Original: The autoregressive model which acts as a learned prior is ﬁtted on the discrete latent variables produced by the encoder.
Revised: The autoregressive model which acts as a learned prior is ﬁt on the discrete latent variables produced by the encoder.

Original: The architecture of the encoder, the decoder, and the latent predictor model are described in further detail in the experiments section.
Revised: The architecture of the encoder, the decoder, and the latent predictor model are described in further detail in the experiments section.

Original: 2. Hard EM and the K -means algorithm
Revised: 2. Hard EM and the K -means algorithm

Original: In this section we brieﬂy recall the hard Expectation maximization (EM) algorithm(4).
Revised: In this section we brieﬂy recall the hard Expectation maximization (EM) algorithm

Revised: (Dempster et al., 1977).

Original: Given a set of data points ( x 1 , . . .
Revised: Given a set of data points ( x 1 , . . .

Original: , x N ) , the hard EM algorithm approximately solves the following optimization problem:
Revised: , x N ) , the hard EM algorithm approximately solves the following optimization problem:

Original: [Equation]
Revised: [Equation]

Original: Hard EM performs coordinate descent over the following two coordinates: the model parameters Θ , and the hidden variables z 1 , . .
Revised: Hard EM performs coordinate descent over the following two coordinates: the model parameters Θ , and the hidden variables z 1 , . .

Original: . , z N .
Revised: . , z N .

Original: In other words, hard EM consists of repeating the following two steps until convergence:
Revised: In other words, hard EM consists of repeating the following two steps until convergence:

Original: [Equation]
Revised: [Equation]

Original: A special case of the hard EM algorithm is K -means clustering (19; 2) where the likelihood is modelled by a Gaussian with identity covariance matrix.
Revised: 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.

Original: Here, the means of the K Gaussians are the parameters to be estimated,
Revised: Here, the means of the K Gaussians are the parameters to be estimated,

Original: [Equation]
Revised: [Equation]

Original: With a uniform prior over the hidden variables ( P Θ ( z i ) = 1 K ) [Equation] the marginal is given by
Revised: With a uniform prior over the hidden variables ( P Θ ( z i ) = 1 K ), the marginal is given by

Original: P Θ ( x i
Revised: P Θ ( x i

Original: | z i [Equation] =N
Revised: | z i ) =

Original: [Equation]

Original:  µ z i , I )( x i ) .
Revised: N ( µ z i , I )( x i ) .

Original: In this case [Equation] equation (6) is equivalent to:
Revised: In this case, equation (6) is equivalent to:

Original: [Equation]
Revised: [Equation]

Original: Note that optimizing equation (7) is NP-hard [Equation] however one can ﬁnd a local optima by applying coordinate descent until convergence:
Revised: Note that optimizing equation (7) is NP-hard, however one can ﬁnd a local optima by applying coordinate descent until convergence:

Original: E step: Cluster assignment is given by,
Revised: E step: Cluster assignment is given by,

Original: [Equation]
Revised: [Equation]

Original: M step: The means of the clusters are updated as,
Revised: M step: The means of the clusters are updated as,

Original: [Equation]
Revised: [Equation]

Original: We can now easily see the connections between the training updates of VQ-VAE and K -means clustering.
Revised: We can now easily see the connections between the training updates of VQ-VAE and K -means clustering.

Original: The encoder output z e ( x ) ∈ R D corresponds to the data point while the discrete latent variables corresponds to clusters.
Revised: The encoder output z e ( x ) ∈ R D corresponds to the data point while the discrete latent variables corresponds to clusters.

Original: Given this, Equation 1 is equivalent to the E-step (Equation 8) and the EMA updates in Equation 4 and Equation 5 converge to the M-step (Equation 9) in the limit.
Revised: Given this, Equation 1 is equivalent to the E-step (Equation 8) and the EMA updates in Equation 4 and Equation 5 converge to the M-step (Equation 9) in the limit.

Original: The M-step in K -means overwrites the old values while the EMA updates interpolate between the old values and the M step update.
Revised: The M-step in K -means overwrites the old values while the EMA updates interpolate between the old values and the M step update.

Original:  VQ-VAE training with EM
Revised:  VQ-VAE training with EM

Original: In this section, we investigate a new training strategy for VQ-VAE using thesoft EM algorithm.
Revised: In this section, we investigate a new training strategy for VQ-VAE using the EM algorithm.

Original: 3. Soft EM
Revised: 3. Expectation Maximization

Original: First, we brieﬂy describe thesoft EM algorithm.
Revised: First, we brieﬂy describe the EM algorithm.

Original: While the hard EM procedure selects one cluster or latent variable assignment for a data point, here the data point is assigned to a mixture of clusters.
Revised: While the hard EM procedure selects one cluster or latent variable assignment for a data point, here the data point is assigned to a mixture of clusters.

Original: Now, the optimization objective is given by,
Revised: Now, the optimization objective is given by,

Original: [Equation]
Revised: [Equation]

Original: Coordinate descent algorithm is again used to approximately solve the above optimization algorithm.
Revised: Coordinate descent algorithm is again used to approximately solve the above optimization algorithm.

Original: The E and M step are given by:
Revised: The E and M step are given by:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: 3. Vector Quantized Autoencoders trained with EM
Revised: 3. Vector Quantized Autoencoders trained with EM

Original: Now, we describe vector quantized autoencoders training using thesoft EM algorithm.
Revised: Now, we describe vector quantized autoencoders training using the EM algorithm.

Original: As discussed in the previous section, the encoder output z e ( x ) ∈ R D corresponds to the data point while the discrete latent variables corresponds to clusters.
Revised: As discussed in the previous section, the encoder output z e ( x ) ∈ R D corresponds to the data point while the discrete latent variables corresponds to clusters.

Original: The E step instead of hard assignment now produces a probability distribution over the set of discrete latent variables (Equation 10).
Revised: The E step instead of hard assignment now produces a probability distribution over the set of discrete latent variables (Equation 10).

Original: Following VQ-VAE, we continue to assume a uniform prior over clusters, since we observe that training the cluster priors seemed to cause the cluster assignments to collapse to only a few clusters.
Revised: Following VQ-VAE, we continue to assume a uniform prior over clusters, since we observe that training the cluster priors seemed to cause the cluster assignments to collapse to only a few clusters.

Original: The probability distribution is modeled as a Gaussian with identity covariance matrix,
Revised: The probability distribution is modeled as a Gaussian with identity covariance matrix,

Original: [Equation]
Revised: [Equation]

Original: Since computing the  expectation  in the M step (Equation 11)is computationally infeasible in our case, weinstead perform Monte-Carlo Expectation Maximization (37) by drawing m samples z 1 i , · · · , z mi ∼ Multinomial (cid:16) − (cid:107) e 1 − z e ( x i ) (cid:107) 22 , . . 
Revised: As an alternative to computing the full expectation term in the M step (Equation 11) we perform Monte-Carlo Expectation Maximization (Wei & Tanner, 1990) by drawing m samples z 1 i , · · · , z mi ∼ Multinomial (cid:16) − (cid:107) e 1 − z e ( x i ) (cid:107) 22 , . . .

Original: , − (cid:107) e K − z e ( x i ) (cid:107) 22 (cid:17) , where
Revised: , − (cid:107) e K − z e ( x i ) (cid:107) 22 (cid:17) , where

Original: Multinomial( l 1 , . .
Revised: Multinomial( l 1 , . .

Original: , l K ) refers to the K -way multinomial distribution with logits l 1 , . . .
Revised: , l K ) refers to the K -way multinomial distribution with logits l 1 , . . .

Original: , l K .
Revised: , l K .

Revised: This results in a less diﬀuse target for the autoregressive prior.

Original: Thus, the E step can be ﬁnally written as:
Revised: Thus, the E step can be ﬁnally written as:

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The model parameters Θ are then updated to maximize this Monte-Carlo estimate in the M step given byM step:
Revised: The model parameters Θ are then updated to maximize this Monte-Carlo estimate in the M step given by

Original: [Equation]
Revised: [Equation]

Original: Instead of exactly following the above M step update, we use the EMA version of this update similar to the one described in Section 2.1.
Revised: Instead of exactly following the above M step update, we use the EMA version of this update similar to the one described in Section 2.1.

Original: When sending the embedding of the discrete latent to the decoder, instead of sending the posterior mode, argmax z P ( z | x ) , similar to hard EM and K -means, we send the average of the embeddings of the sampled latents:
Revised: When sending the embedding of the discrete latent to the decoder, instead of sending the posterior mode, argmax z P ( z | x ) , similar to hard EM and K -means, we send the average of the embeddings of the sampled latents:

Original: [Equation]
Revised: [Equation]

Original: Since m latent code embeddings are sent to the decoder in the forward pass, all of them are updated in the backward pass for a single training example.
Revised: Since m latent code embeddings are sent to the decoder in the forward pass, all of them are updated in the backward pass for a single training example.

Original: In hard EM training, only one of them is updated during training.
Revised: In hard EM training, only one of them is updated during training.

Original: Sending averaged embeddings also results in more stable training using thesoft EM algorithm compared to hard EM as shown in Section 5.
Revised: Sending averaged embeddings also results in more stable training using the EM algorithm compared to VQ-VAE as shown in Section 5.

Original: To train the latent predictor model (Section 2.1) in this case, we use an approach similar to label smoothing (24): the latent predictor model is trained to minimize the cross entropy loss with the labels being the average of the one-hot labels of z 1 i , . . .
Revised: To train the latent predictor model (Section 2.1) in this case, we use an approach similar to label smoothing (Pereyra et al., 2017): the latent predictor model is trained to minimize the cross entropy loss with the labels being the average of the one-hot labels of z 1 i , . . .

Original: , z mi .
Revised: , z mi .

Original:  Other Related Work
Revised:  Other Related Work

Original: Variational autoencoders were ﬁrst introduced by (14; 26) for training continuous representations; unfortunately, training them for discrete latent variable models has proved challenging.
Revised: Variational autoencoders were ﬁrst introduced by Kingma & Welling (2014) for training continuous representations; unfortunately, training them for discrete latent variable models has proved challenging.

Original: One promising approach has been to use various gradient estimators for discrete latent variable models, starting with the REINFORCE estimator of (38), an unbiased, high-variance gradient estimator.
Revised: One promising approach has been to use various gradient estimators for discrete latent variable models, starting with the REINFORCE estimator of Williams (1992), an unbiased, high-variance gradient estimator.

Revised: Subsequent work on improving the variance of the REINFORCE estimator are REBAR (Tucker et al., 2017) and RELAX (Grathwohl et al., 2017).

Original: An alternate approach towards gradient estimators is to use continuous relaxations of categorical distributions, for e.g., the Gumbel-Softmax reparametrization trick (8; 20).
Revised: An alternate approach towards gradient estimators is to use continuous relaxations of categorical distributions, for e.g., the Gumbel-Softmax reparametrization trick (Jang et al., 2016; Maddison et al., 2016).

Original: These methods provide biased but low variance gradients for training.
Revised: These methods provide biased but low variance gradients for training.

Original: Machine translation using deep neural networks have been shown to achieve impressive results (30; 1; 3; 35).
Revised: Machine translation using deep neural networks have been shown to achieve impressive results (Sutskever et al., 2014;

Revised: Bahdanau et al., 2014; Cho et al., 2014; Vaswani et al., 2017).

Original: The state-of-the-art models in Neural Machine Translation are all autoregressive, which means that during decoding, the model consumes all previously generated tokens to predict the next one.
Revised: The state-of-the-art models in Neural Machine Translation are all auto-regressive, which means that during decoding, the model consumes all previously generated tokens to predict the next one.

Original: Very recently, there have been multiple eﬀorts to speed-up machine translation decoding. 
Revised: Recently, there have been multiple eﬀorts to speed-up machine translation decoding. Gu et al.

Original: (6) attempts to address this issue by using the Transformer model (35) together with the REINFORCE algorithm (38), to model the fertilities of words.
Revised: (2017) attempts to address this issue by using the Transformer model (Vaswani et al., 2017) together with the REINFORCE algorithm (Williams, 1992), to model the fertilities of words.

Original: The main drawback of the approach of (6) is the need for extensive ﬁne-tuning to make policy gradients work, as well as the non-generic nature of the solution. 
Revised: The main drawback of the approach of Gu et al. (2017) is the need for extensive ﬁne-tuning to make policy gradients work, as well as the non-generic nature of the solution. Lee et al.

Original: (16) propose a non-autoregressive model using iterative reﬁnement.
Revised: (2018) propose a non-autoregressive model using iterative reﬁnement.

Original: Here, instead of decoding the target sentence in one-shot, the output is successively reﬁned to produce the ﬁnal output.
Revised: Here, instead of decoding the target sentence in one-shot, the output is successively reﬁned to produce the ﬁnal output.

Original: While the output is produced in parallel at each step, the reﬁnement steps happen sequentially.
Revised: While the output is produced in parallel at each step, the reﬁnement steps happen sequentially.

Original:  Experiments
Revised:  Experiments

Original: We evaluate our proposed methods for image reconstruction on the CIFAR-10 and SVHN dataset, dataset and supervised conditional language generation on the WMT English-toGerman translation task.
Revised: In this section we report our experiments with VQ-VAE and EM on the English-German translation task, with the aim of improving the decoding speed of autoregressive translation models.

Original: Our models and generative process follow the architecture proposed in (34) for unconditional image generation, and (10) for neural machine translation.
Revised: Our model and generative process follows the architecture proposed in Kaiser et al. and is depicted in Figure 1.

Original: For all our experiments, we use the Adam (13) optimizer and decay the learning rate exponentially after initial warm-up steps.
Revised: For all our experiments, we use the Adam (Kingma & Ba, 2014) optimizer and decay the learning rate exponentially after initial warm-up steps.

Original: Unless otherwise stated, the dimension of the hidden states of the encoder and the decoder is 512 , see Table 4 for a comparison of models with lower dimension.
Revised: Unless otherwise stated, the dimension of the hidden states of the encoder and the decoder is 512 , see Table 4 for a comparison of models with lower dimension.

Revised: For all conﬁgurations we select the optimal hyperparameters by using WMT’13 English-German as the validation set and reporting the BLEU score on the WMT’14 English-German test set.

Original: 5. Machine Translation
Revised: 5. Machine Translation

Original: In Neural Machine Translation with latent variables, we model P ( y, z | x ) , where y and x are the target and source sentence respectively.
Revised: In Neural Machine Translation with latent variables, we model P ( y, z | x ) , where y and x are the target and source sentence respectively.

Original: Our model architecture, depicted in Figure 2, is similar to the one in (10).
Revised: Our model architecture, depicted in Figure 1, is similar to the one in Kaiser et al.

Original: The encoder function is a series of strided convolutional layers with residual convolutional layers in between and takes target sentence y as input.
Revised: The encoder function is a series of strided convolutional layers with residual convolutional layers in between and takes target sentence y as input.

Original: The source sentence x is converted to a sequence of hidden states through multiple causal self-attention layers.
Revised: The source sentence x is converted to a sequence of hidden states through multiple causal self-attention layers.

Revised: In Kaiser et al.

Original: In (10), the encoder of the autoencoder attends additionally to this sequence of continuous representation of the source sentence.
Revised: (2018), the encoder of the autoencoder attends additionally to this sequence of continuous representation of the source sentence.

Original: We use VQ-VAE as the discretization algorithm.
Revised: We use VQ-VAE as the discretization algorithm.

Original: The decoders, applied after the bottleneck layer uses transposed convolution layers whose continuous output is fed to a transformer decoder with causal attention, which generates the output.
Revised: The decoders, applied after the bottleneck layer uses transposed convolution layers whose continuous output is fed to a transformer decoder with causal attention, which generates the output.

Original: The results are summarized in Table 1.
Revised: The results are summarized in Table 1.

Original: Our implementation of VQ-VAE achieves a signiﬁcantly better BLEU score and faster decoding speed compared to (10).
Revised: Our implementation of VQ-VAE achieves a signiﬁcantly better BLEU score and faster decoding speed compared to Kaiser et al.

Original: We found that tuning the code-book size (number of clusters) for using 2 12 discrete latents achieves the best accuracy which is 16 times smaller as compared to the code-book size in (10).
Revised: We found that tuning the code-book size (number of clusters) for using 2 12 discrete latents achieves the best accuracy which is 16 times smaller as compared to the code-book size in Kaiser et al.

Original: Additionally, we see a large improvement in the performance of the model by using sequence-level distillation (12), as has been observed previously in non-autoregressive models (6; 16).
Revised: Additionally, we see a large improvement in the performance of the model by using sequence-level distillation (Hinton et al., 2015; Kim & Rush, 2016), as has been observed previously in non-autoregressive models (Gu et al., 2017; Lee et al., 2018).

Original: Our teacher model is a base Transformer (35) that achieves a BLEU score of 28 . 1 and 27 . 

Revised: Our teacher model is a base Transformer (Vaswani et al., 2017) that achieves a BLEU score of 28 . 1 and 27 .

Original: on the WMT’14 test set using beam search decoding and greedy decoding respectively.
Revised: on the WMT’14 test set using beam search decoding and greedy decoding respectively.

Original: For distillation purposes, we use the  beam search decoded Transformer.
Revised: The distilled data is decoded from the base Transformer using a beam size of

Original: Our VQ-VAE model trained with soft EM and distillation, achieves a BLEU score of 26 . 7 , without noisy parallel decoding (6).
Revised: Our VQ-VAE model trained with soft EM and distillation, achieves a BLEU score of 26 . 7 , without noisy parallel decoding (Gu et al., 2017).

Original: This perforamce is 1 . 4 bleu points lower than an autoregressive model decoded with a beam size of 4 , while being 4 .
Revised: This perforamce is 1 . 4 bleu points lower than an autoregressive model decoded with a beam size of 4 , while being 4 .

Original: 1 × faster.
Revised: 1 × faster.

Original: Importantly, we nearly match the same autoregressive model with beam size 1 , with a 3 . 3 × speedup.
Revised: Importantly, we nearly match the same autoregressive model with beam size 1 (greedy decoding), with a 3 . 3 × speedup.

Original: The length of the sequence of discrete latent variables is shorter than that of target sentence y .
Revised: The length of the sequence of discrete latent variables is shorter than that of target sentence y .

Original: Speciﬁcally, at each compression step of the encoder we reduce its length by half.
Revised: Speciﬁcally, at each compression step of the encoder we reduce its length by half.

Original: We denote by n c , the compression factor for the latents, i.e. the number of steps for which we do this compression.
Revised: We denote by n c , the compression factor for the latents, i.e. the number of steps for which we do this compression.

Original: In almost all our experiments, we use n c = 3 reducing the length by 8.
Revised: In almost all our experiments, we use n c = 3 reducing the length by 8.

Original: We can decrease the decoding time further by increasing the number of compression steps.
Revised: We can decrease the decoding time further by increasing the number of compression steps.

Original: As shown in Table 1, by setting n c to 4, the decoding time drops to 58 milliseconds achieving 25.4 BLEU while a NAT model with similar decoding speed achieves only 18.
Revised: As shown in Table 1, by setting n c to 4, the decoding time drops to 58 milliseconds achieving 25.4 BLEU while a NAT model with similar decoding speed achieves only 18.

Original: BLEU.
Revised: BLEU.

Original: Note that, all NAT models also train with sequence level knowledge distillation from an autoregressive teacher.
Revised: Note that, all NAT models also train with sequence level knowledge distillation from an autoregressive teacher.

Original: 5.1. Analysis
Revised: 5.1. Analysis

Revised: Size of Discrete Latent Variable code-book: Table 2 in Appendix shows the BLEU score for diﬀerent code-book sizes for models trained using VQ-VAE without distillation.

Revised: While Kaiser et al.

Original:  Attention to Source Sentence Encoder: While the encoder of the discrete autoencoder in (10) attends to the output of the encoder of the source sentence, we ﬁnd that to be unnecessary, with both models achieving the same BLEU score with 2 12 latents.
Revised: (2018) use 2 16 as their code-book size, we ﬁnd that 2 12 gives the best performance. Attention to Source Sentence Encoder: While the encoder of the discrete autoencoder in Kaiser et al. (2018) attends to the output of the encoder of the source sentence, we ﬁnd that to be unnecessary, with both models achieving the same BLEU score with 2 12 latents.

Original: Note that the decoder of the discrete autoencoder in both (10) and our work does not attend to the source sentence.
Revised: Note that the decoder of the discrete autoencoder in both Kaiser et al. and our work does not attend to the source sentence.

Original: VQ-VAE vs Other Discretization Techniques:
Revised: VQ-VAE vs Other Discretization Techniques:

Original: We compare the Gumbel-Softmax of (8; 20) and the improved semantic hashing discretization technique proposed in (10) to VQ-VAE.
Revised: We compare the Gumbel-Softmax of (Jang et al., 2016; Maddison et al., 2016) and the improved semantic hashing discretization technique proposed in Kaiser et al. (2018) to VQ-VAE.

Original: When trained with sequence level knowledge distillation, the model using Gumbel-Softmax reached 23 .
Revised: When trained with sequence level knowledge distillation, the model using Gumbel-Softmax reached 23 .

Original: BLEU, the model using improved semantic hashing reached 24 .
Revised: BLEU, the model using improved semantic hashing reached 24 .

Original: 1 BLEU, and the model using VQ-VAE reached 26 .
Revised: 1 BLEU, and the model using VQ-VAE reached 26 .

Original: 4 BLEU on WMT’14 English-German.
Revised: 4 BLEU on WMT’14 English-German.

Original: Size of Discrete Latent Variable code-book: Table 2 in Appendix shows the BLEU score fordiﬀerent code-book sizes for models trained using hard EM without distillation.
Revised: Model Size: The eﬀect of model size on BLEU score for models trained with EM and distillation is shown in Table 4 in Appendix.

Original: While (10) use 2 16 as their code-book size, we ﬁnd that 2 12 gives the best performance.

Original: Robustness of  EM to Hyperparameters: Whilethe soft EM training gives a small performance improvement, we ﬁnd that it also leads to more robust training- both for translation (Figure 3) as well as for images (Figure 5). [Table] for with and without EM) with a hidden dimension of size 512 . Decoding is performed on a single CPU machine with an
Revised: Number of samples in Monte-Carlo EM update: While training with EM, we perform a Monte-Carlo update with a small number of samples (Section 3.2).

Original: NVIDIA GeForce GTX 1080 with a batch size of

Original: [Table]
Revised: Table 3 in

Original: Our experiments on image reconstruction on SVHN (21) in section A also highlight the robustness of EM training. The training approach from (34) exhibits high variance on reconstruction quality, while soft EM is much more stable, resulting in good reconstructions in almost all training runs. Model Size: The eﬀect of model size onBLEU score for models trained with soft EM and distillation is shown in Table 4 in Appendix. Gradient based update vs EMA update of code-book: The original VQ-VAE paper (34) proposed a gradient based update rule for learning the code-book where the code-book entries are trained by minimizing (cid:107) sg
Revised: Appendix shows the impact of number of samples on the ﬁnal BLEU score.

Original: ( z e ( x )) − z q ( x ) (cid:107) 2 . However, it was found in (10) that the EMA update worked better than this gradient based loss. Note that if the gradient based loss was minimized using SGD then the update rule for  the embeddings is [Equation]
Revised: Emergence of EOS/PAD latent: We observe that all the latent sentences for a speciﬁc experiment with VQ-VAE/EM end with a ﬁxed latent indicating the end of the sequence.

Original: for a learning rate η .

Original: This is quite similar to the EMA update rule ofEquation 5, with the only diﬀerence being that the latter also maintains an EMA over the counts c j . When using SGD with momentum or Adam, the update rule becomes quite diﬀerent however, since we now take the moving average of the gradient term itself, before subtracting it from current value of the embedding e j . This is similar to  the issue ofusing weight decay with Adam, where using the (cid:96) 2 penalty in the loss function results in worse performance (see e.g., (18)).
Revised: Since we always ﬁx the length of the latent sentence to be 2 n c times smaller than the true sentence, the model learns to pad the remainder of the latent sequence with this special code (see Table 5 for examples).

Original: [Table] update: While training with soft EM, we perform a Monte-Carlo update with a small number of samples (Section 3.2). Table 3 in Appendix shows the impact of number of samples on the ﬁnal BLEU score.
Revised: Note that one can speed up decoding even further by stopping the Latent Predictor (LP) model as soon as it outputs this special code.

Original:  Conclusion
Revised:  Conclusion

Original: We investigate an alternate training technique for VQ-VAE inspired by its connection to the EM algorithm.
Revised: We investigate an alternate training technique for VQ-VAE inspired by its connection to the EM algorithm.

Original: Training the discrete bottleneck with EMhelps us achieve better image reconstructions on SVHN, and together with  knowledge distillation, allows us to develop a non-autoregressive machine translation model whose accuracy almost matches the greedy autoregressive baseline, while being 3.3 times faster at inference. 
Revised: 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. [Table] distillation from Hinton et al. ; Kim & Rush (2016).

Revised: We used a code-book of size 2 12 for VQ-VAE (for with and without EM) with a hidden dimension of size 512 .

Revised: Decoding is performed on a single CPU machine with an NVIDIA GeForce GTXwith a batch size of

Original:  A Image Reconstruction
Revised: [Table] [Table] A Image Reconstruction

Revised: We use the notation x to denote the input image, with the output of the encoder z e ( x ) ∈ R D being used to perform nearest neighbor search to select the (sequence of) discrete latent variable.

Revised: The selected discrete latent is used to train the latent predictor model, while the embedding z q ( x ) of the selected discrete latent is passed as input to the decoder.

Revised: In this section we report additional experiments we performed using VQ-VAE and EM for the task of image reconstruction.

Original: We train a discrete autoencoder with VQ-VAE (34) andsoft EM on the CIFAR-10 data set, modeling the joint probability P ( x, z ) , where x is the image and z are the discrete latent codes.
Revised: We train a discrete autoencoder with VQ-VAE (van den Oord et al., 2017) and EM on the CIFAR-10 data set, modeling the joint probability P ( x, z ) , where x is the image and z are the discrete latent codes.

Original: We use a ﬁeld of 8 × 8 × 10 latents with a code-book of size 2 8 each containingdimensions.
Revised: We use a ﬁeld of 8 × 8 × 10 latents with a code-book of size 2 8 each containing 512 dimensions.

Original: We maintain the same encoder and decoder as used in Machine Translation.
Revised: We maintain the same encoder and decoder as used in Machine Translation.

Original: For the encoder, we use 4 convolutional layers, with kernel size 5 × 5 and strides 2 × 2 , followed by 2 residual layers, and a single dense layer.
Revised: For the encoder, we use 4 convolutional layers, with kernel size 5 × 5 and strides 2 × 2 , followed by 2 residual layers, and a single dense layer.

Original: For the decoder, we use a single dense layer,residual layers, and 4 deconvolutional layers.
Revised: For the decoder, we use a single dense layer, 2 residual layers, and 4 deconvolutional layers.

Original: Figure 4 shows that our reconstructions are on par with hard EM training.
Revised: Figure 3 shows that our reconstructions are on par with hard EM training.

Original: We also train discrete autoencoders on the SVHN dataset (21), with bothVQ-VAE (34) and
Revised: We also train discrete autoencoders on the SVHN dataset (Netzer et al., 2011), with both

Original:  EM.
Revised: VQ-VAE (van den Oord et al., 2017) and EM.

Original: The autoencoder is similar to our CIFAR-10 model, where each n x = 32 × 32 × 3 image is encoded into 640 discrete latents from a shared codebook of size 256 .
Revised: The autoencoder is similar to our CIFARmodel, where each n x = 32 × 32 × 3 image is encoded into 640 discrete latents from a shared codebook of size 256 .

Original: By contrasting the reconstructions from several training runs for VQ-VAE (left) and EM (right), we ﬁnd that training with EM is more reliable and the reconstructions are of high quality (Figure 5)
Revised: By contrasting the reconstructions from several training runs for VQ-VAE (left) and EM (right), we ﬁnd that training with EM is more reliable and the reconstructions are of high quality (Figure 4)

Original: B Ablation Tables
Revised: B Ablation Tables

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: C Additional Analysis
Revised: C Additional Analysis

Revised: Gradient based update vs EMA update of code-book: The original VQ-VAE paper (van den Oord et al., 2017) proposed a gradient based update rule for learning the code-book where the code-book entries are trained by minimizing (cid:107) sg

Revised: ( z e ( x )) − z q ( x ) (cid:107) 2 .

Revised: However, it was found in Kaiser et al.

Revised: (2018) that the EMA update worked better than this gradient based loss.

Original: Denoising autoencoder: We also use word dropout with a dropout rate of 0 . 3 and word permutation with a shuﬄe rate of 0 .5 as in (15).
Revised: Note that if the gradient based loss was minimized using SGD then the update rule for the embeddings is [Equation] for a learning rate η . This is quite similar to the EMA update rule of Equation 5, with the only diﬀerence being that the latter also maintains an EMA over the counts c j .

Revised: When using

Original: On the WMT English-German we did not notice any improvement from  usingthese regularization techniques, but on the larger WMT
Revised: SGD with momentum or Adam, the update rule becomes quite diﬀerent however, since we now take the moving average of the gradient term itself, before subtracting it from current value of the embedding e j . This is similar to the issue of using weight decay with Adam, where using the (cid:96) 2 penalty in the loss function results in worse performance (Loshchilov & Hutter, 2017). Robustness of EM to Hyperparameters

Original: English-French dataset, we observe that using a denoising autoencoder signiﬁcantly improves performance with a gain of 1.0 BLEU on VQ-VAE and 0.9 BLEU over EM (see Table 6).
Revised: : While EM training gives a small performance improvement, we ﬁnd that it also leads to more robust training for machine translation. Our experiments on image reconstruction on SVHN (Netzer et al., 2011) in section A also highlight the robustness of EM training.

Original: Emergence of EOS/PAD latent: We observe that all the latent sentences for a speciﬁc experiment with VQ-VAE/EM end with a ﬁxed latent indicating the end of the sequence.
Revised: The training approach from van den Oord et al. (2017) exhibits high variance on reconstruction quality, while EM is much more stable, resulting in good reconstructions in almost all training runs.

Original: Since we always ﬁx the length of the latent sentence to be 2 n c times smaller than the true sentence, the model learns to pad the remainder of  the latent sequence with this special code (see Table 5 for examples). Note that one can speed up decoding even further by stopping the Latent Predictor (LP) model as soon as it outputs this special code.
Revised: Denoising autoencoder: We also use word dropout with a dropout rate of 0 . 3 and word permutation with a shuﬄe rate of 0 . 5 as in Lample et al. On the WMT EnglishGerman we did not notice any improvement from using these regularization techniques, but on the larger WMT English-French dataset, we observe that using a denoising autoencoder signiﬁcantly improves performance with a gain of 1.0 BLEU on VQ-VAE and 0.9 BLEU over EM (see Table 6).

Original: Additional analysis on latents: In order to compute correlations between the discrete latents and n -grams in the original text, we computed Point-wise Mutual Information (PMI) and tf-idf scores where the latents are treated as documents.
Revised: Additional analysis on latents: In order to compute correlations between the discrete latents and n -grams in the original text, we computed Point-wise Mutual Information (PMI) and tf-idf scores where the latents are treated as documents.

Original: However, we were unable to see any semantic patterns that stood out in this analysis.
Revised: However, we were unable to see any semantic patterns that stood out in this analysis.

Original: [Table]
Revised: [Table]

Original: D Preliminary Results on English French
Revised: D Preliminary Results on English French

Original: In this section we report preliminary results on the WMT English-French dataset without using knowledge distillation from an autoregressive teacher (7).
Revised: In this section we report preliminary results on the WMT English-French dataset without using knowledge distillation from an autoregressive teacher (Hinton et al., 2015; Kim & Rush, 2016).

Original: We use a Transformer base model from (35).
Revised: We use a Transformer base model from Vaswani et al.

Original: Our best non-autoregressive base model trained on non-distilled targets gets 30 . 

Revised: Our best non-autoregressive base model trained on non-distilled targets gets 30 .

Original: BLEU compared to the autoregressive base model with the same choice of hyperparameters, which gets 33 .3 BLEU (see Table 6).
Revised: BLEU compared to the autoregressive base model with the same choice of hyperparameters, which gets 33 . BLEU (see Table 6).

Original: As in the case of English-German, we anticipate that using knowledge distillation (7) will likely close this gap.
Revised: As in the case of English-German, we anticipate that using knowledge distillation Hinton et al. (2015) will likely close this gap.

Original: [Table]
Revised: [Table]

Revised: with shuﬄe rate 0 .

Revised: 5 as described in Section C.

Revised: The hidden dimension of the codebook is 512 .

Revised: Decoding is performed on a single CPU machine with an NVIDIA GeForce GTX 1080 with a batch size of 1 .