Source-linked AI summary
Neural Discrete Representation Learning
Aaron van den Oord, Oriol Vinyals, Koray Kavukcuoglu
TL;DR
Unsupervised learning of useful generic representations remains challenging. VQ-VAE learns discrete latent representations with vector quantisation and, paired with an autoregressive prior, generates meaningful images, videos, and speech while achieving likelihoods nearly matching continuous-latent counterparts.
Problem
Useful generic representations trained unsupervised from raw data remain insufficiently established for challenging learning tasks.
Method
VQ-VAE combines variational autoencoders with vector quantisation to learn discrete latent representations and a learned prior.
Results
VQ-VAEs model long-term dependencies across images, videos, and audio, generate meaningful speech, support speaker conversion, and achieve likelihoods nearly matching continuous counterparts.
Takeaways & Limitations
The learned discrete latent space captures important data features without supervision and supports coherent generative modeling across modalities.
Takeaways & Limitations
The encoder-decoder mapping can make each discrete latent depend on preceding latents, potentially requiring more advanced phoneme mapping for higher accuracy.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper introduces VQ-VAE, a generative model that learns discrete latent representations through vector quantization while avoiding posterior collapse. Combined with a learned autoregressive prior, it supports coherent, high-quality generation and unsupervised discovery of useful structure across modalities.
- Motivation: Unsupervised representations remain insufficiently useful despite advances in generative modeling for images, audio, and video.Such representations are important for few-shot learning, domain adaptation, and reinforcement learning.
- Motivation: The paper targets latent spaces that preserve important data features while optimizing maximum likelihood, rather than relying solely on pixel-domain reconstruction.The introduction notes that powerful decoders can produce strong likelihoods while ignoring latent variables.
- Model: VQ-VAE combines the VAE framework with discrete latent representations using a vector-quantized posterior parameterization.The model is described as simple to train and free from large-variance issues.
- Model: VQ-VAE avoids posterior collapse, in which powerful decoders ignore latent variables, while matching the performance of continuous latent VAE counterparts.Discrete representations are also presented as a natural fit for language and speech, and potentially for images and complex reasoning.
- Results: A learned prior over VQ-VAE’s discrete variables yields coherent, high-quality samples and useful applications across speech and video generation.On speech, the model discovers latent language structure without supervision or prior knowledge of phonemes or words, and supports speaker conversion when given speaker identity.
2 Related Work
The paper introduces a method for training VAEs with discrete latent variables, addressing challenges that have limited their use relative to continuous representations. It positions this approach against alternative discrete-VAE estimators, continuous relaxations, autoregressive VAE models, and neural image-compression methods.
- Discrete latent variables: The work presents a new method for training variational autoencoders with discrete latent variables, despite continuous latent variables dominating current deep-learning research.The authors note that using discrete variables in deep learning has proven challenging, even for inherently discrete modalities.
- Alternative discrete-VAE methods: NVIL optimizes a single-sample variational lower-bound objective with variance reduction, while VIMCO uses multiple inference-network samples to accelerate convergence.Both methods are described as alternatives for training discrete VAEs.
- Continuous relaxations: Concrete and Gumbel-softmax methods use temperature-annealed continuous distributions that approach discrete distributions, trading low biased gradient variance early for high unbiased variance late.The relaxation is continuous during training and converges to a discrete distribution in the limit.
- Limitations of prior methods: Existing discrete-latent methods do not close the performance gap with continuous-latent VAEs using Gaussian reparameterization and are usually tested on small datasets with latent dimensionality below 8.The paper instead evaluates on CIFAR10, ImageNet, DeepMind Lab, and a raw speech dataset.
- Autoregressive models and compression: Related work also includes autoregressive decoders or priors in VAEs, including LSTM, dilated-convolution, and PixelCNN-based models, plus scalar- and vector-quantization approaches for neural image compression.The paper further discusses a soft-to-hard vector-quantization relaxation that is annealed over time, but reports difficulty training it from scratch.
3 VQ-VAE
VQ-VAE learns discrete latent representations by combining categorical variables, nearest-neighbor vector quantization, and embedding-table lookup. Its training objective separately updates the decoder, encoder, and embeddings, after which an autoregressive prior models the discrete latents for generation.
- Model definition: VQ-VAE uses categorical posterior and prior distributions whose sampled indices select embeddings from a shared embedding table for decoder input.This replaces continuous latent representations with discrete codes trained using vector-quantization ideas.
- Quantization: The encoder output ze(x) is discretized by assigning it to the nearest vector among K embedding vectors in the shared space e ∈ R^K×D.The embedding space contains K vectors, each with dimensionality D.
- Optimization: A straight-through estimator copies decoder gradients from zq(x) to ze(x), allowing encoder outputs to change their discrete assignment in later forward passes.The nearest embedding is used during forward computation, while gradients pass unaltered to the encoder during backpropagation.
- Optimization: The three-part loss combines reconstruction, embedding, and commitment terms, with the decoder optimizing reconstruction, the encoder optimizing reconstruction and commitment, and embeddings optimizing the middle term.The commitment term prevents encoder outputs from growing without binding to an embedding; performance was robust for β values from 0.1 to 2.0.
- Autoregressive prior: The prior is initially uniform and fixed during VQ-VAE training, then replaced by an autoregressive distribution over discrete latents for ancestral generation.The paper uses PixelCNN for image latents and WaveNet for raw audio.
4 Experiments
Experiments show that VQ-VAE learns compact discrete representations that retain useful structure across images, audio, and video. The model challenges continuous-VAE performance, avoids posterior collapse, and supports reconstruction, generation, speaker conversion, and phoneme discovery.
- CIFAR10: On CIFAR10, VQ-VAE achieves 4.67 bits/dim, compared with 4.51 for a VAE and 5.14 for VIMCO.All reported likelihoods are lower bounds, and the continuous-VAE result is comparable to a reported 4.54 bits/dim for a Deep convolutional VAE.
- CIFAR10: VQ-VAE is the first discrete-latent model described as challenging continuous-VAEs while combining VAE-like reconstructions with compressed symbolic representations.The paper presents this combination as a key characteristic of the trained VQ-VAEs.
- Image modeling: ImageNet images are compressed from 128×128×3 pixels to a 32×32×1 discrete latent space with K=512, yielding reconstructions only slightly blurrier than the originals.The latent-space compression reduces the representation by approximately 42.6 in bits, while a PixelCNN prior models global image structure.
- Image modeling: On DeepMind Lab frames, VQ-VAE produces reconstructions nearly identical to the originals and samples from a PixelCNN prior over a 21×21×1 latent space.A second-stage VQ-VAE using three K=512 latent variables meaningfully uses its latents instead of suffering posterior collapse.
- Audio modeling: For speech, VQ-VAE compresses the VCTK waveform representation 64 times and preserves long-term relevant information rather than exact sample-level detail.The decoder is conditioned on both the discrete latents and a speaker embedding.
- Audio modeling: Speech experiments show speaker conversion with preserved content, while unsupervised latent-to-phoneme mapping reaches 49.3% accuracy on 41-way classification.The phoneme experiment uses a 128-dimensional discrete space running at 25 Hz.
- Video modeling: In DeepMind Lab video generation, VQ-VAE generates future frames conditioned on action sequences entirely in latent space, without generating the images themselves.The demonstrated actions are repeated moves forward or right.
5 Conclusion
The paper introduces VQ-VAE, combining variational autoencoders with vector quantisation to learn discrete latent representations. Its compressed discrete latent space supports long-term dependency modeling across images, video, and speech, while encoder/decoder context may complicate phoneme mapping.
- 5 Conclusion: VQ-VAE combines VAEs with vector quantisation to obtain a discrete latent representation.The authors present it as a new family of models.
- 5 Conclusion: The compressed discrete latent space models very long-term dependencies in 128 × 128 colour images, action-conditional video, and speech generation.Even an unconditional audio model generated surprisingly meaningful chunks of speech.
- 5 Conclusion: Context-dependent encoder/decoder mappings could improve compression but require more advanced phoneme mapping for higher accuracy.Each discrete latent may depend on previous latents, including bi-grams or tri-grams.
A Appendix · A.1 VQ-VAE dictionary updates with Exponential Moving Averages
This appendix section describes an exponential-moving-average alternative for updating VQ-VAE dictionary items. It derives the average-based update, explains its minibatch adaptation, and reports γ = 0.99 as effective in practice.
- A.1 VQ-VAE dictionary updates with Exponential Moving Averages: VQ-VAE dictionary items can be updated with exponential moving averages instead of the loss-term update from Equation 3.
- A.1 VQ-VAE dictionary updates with Exponential Moving Averages: For dictionary item e_i, the optimal update is the average of encoder outputs assigned to that item.
- A.1 VQ-VAE dictionary updates with Exponential Moving Averages: This average-based dictionary update is typically used in algorithms such as K-Means.
- A.1 VQ-VAE dictionary updates with Exponential Moving Averages: The direct average update cannot be used with minibatches.
- A.1 VQ-VAE dictionary updates with Exponential Moving Averages: Exponential moving averages provide an online version of the average-based update for minibatch training.
- A.1 VQ-VAE dictionary updates with Exponential Moving Averages: The decay parameter γ lies between 0 and 1, and γ = 0.99 worked well in practice.