[Paper Review] Neural Discrete Representation Learning

Aäron van den Oord, Oriol Vinyals+1 more|arXiv (Cornell University)|Nov 2, 2017

Speech Recognition and Synthesis1,919 citations

TL;DR

Introduces VQ-VAE, a variational autoencoder with discrete latent variables learned via vector quantisation, avoiding posterior collapse and enabling high-quality generation with autoregressive priors.

ABSTRACT

Learning useful representations without supervision remains a key challenge in machine learning. In this paper, we propose a simple yet powerful generative model that learns such discrete representations. Our model, the Vector Quantised-Variational AutoEncoder (VQ-VAE), differs from VAEs in two key ways: the encoder network outputs discrete, rather than continuous, codes; and the prior is learnt rather than static. In order to learn a discrete latent representation, we incorporate ideas from vector quantisation (VQ). Using the VQ method allows the model to circumvent issues of "posterior collapse" -- where the latents are ignored when they are paired with a powerful autoregressive decoder -- typically observed in the VAE framework. Pairing these representations with an autoregressive prior, the model can generate high quality images, videos, and speech as well as doing high quality speaker conversion and unsupervised learning of phonemes, providing further evidence of the utility of the learnt representations.

Motivation & Objective

• Motivate learning useful representations without supervision across images, audio, and video.
• Develop a discrete-latent VAE that overcomes posterior collapse seen with powerful decoders.
• Show that discrete latents can match continuous VAEs in likelihood while enabling strong priors for generation.
• Demonstrate applications including image/video generation, speech understanding, and unsupervised speaker conversion.

Proposed method

• Define a latent embedding space e in R^{K x D} with K discrete codes.
• Encoder outputs z_e(x); z is obtained by nearest-neighbor lookup against embedding space e (z_q(x)=e_k).
• Train with a three-term loss: reconstruction log p(x|z_q(x)), a VQ loss updating e toward encoder outputs, and a commitment loss to keep encoder outputs close to embeddings (using stop-gradient).
• Use straight-through estimator to propagate gradients through the discrete quantisation step.
• Assume a uniform prior over z, making the KL term constant; later fit an autoregressive prior over z (PixelCNN for images, WaveNet for audio) for generation.
• Evaluate log p(x) via log p(x|z_q(x)) p(z_q(x)) as an approximation and compare with continuous VAEs.

Experimental results

Research questions

• RQ1 Can a discrete latent VAE (VQ-VAE) achieve competitive log-likelihood with continuous VAEs on standard datasets?
• RQ2 Does discretising latent variables help avoid posterior collapse when using strong decoders while preserving reconstruction quality?
• RQ3 Can a learned autoregressive prior over discrete latents enable coherent, high-quality generation in images, audio, and video?
• RQ4 Do discrete latent representations capture meaningful high-level structure (e.g., phonemes in speech) in an unsupervised manner?

Key findings

• VQ-VAE achieves competitive likelihoods with continuous VAEs on CIFAR-10 (4.67 bits/dim for VQ-VAE vs 4.51 bits/dim for continuous VAE, 5.14 for VIMCO).
• Discrete latents enable high-quality reconstructions on ImageNet (128x128x3) using a 32x32x1 latent space (K=512) with a PixelCNN prior.
• For audio, the model learns a latent space invariant to low-level waveform details, supporting unsupervised phoneme-like structure and enabling speaker conversion with a separate speaker embedding.
• In video modeling, the latent space supports long-term sequence generation by sampling z from a learned prior and decoding to frames, maintaining local geometry without pixel-level generation.
• The model avoids posterior collapse and uses a simple, robust training regime with a straightforward dictionary update (VQ) and a commitment term.