Skip to content

Neural Discrete Representation Learning (VQ-VAE)

Authors: Aaron van den Oord, Oriol Vinyals, Koray Kavukcuoglu (DeepMind)

Published: November 2017 (arXiv 1711.00937), NeurIPS 2017

Paper Link: https://arxiv.org/abs/1711.00937

Why This Paper Matters

VQ-VAE introduced vector quantization into the deep learning autoencoder framework, producing a model that compresses continuous data (images, audio, video) into discrete tokens that can then be modeled with the same powerful tools used for language. This bridge between continuous perception data and discrete sequence modeling was foundational - without it, there is no DALL-E 1, no VQ-GAN, no SoundStream, no Encodec, no MusicLM, and arguably no Sora.

The paper's core trick - replacing the latent vector of an autoencoder with the nearest entry in a learned codebook, and using a straight-through estimator to backpropagate through the non-differentiable lookup - has become a standard primitive in modern multimodal generative AI.

The Problem Before VQ-VAE

By 2017, autoencoders had two dominant flavors:

  • VAEs (Kingma & Welling, 2013) - continuous latents with a Gaussian prior. Mathematically principled but often produced blurry images and underutilized the latent space ("posterior collapse")
  • GANs - sharp images but no inference network, no likelihood, no useful latent representation

Meanwhile, the most powerful sequence models - autoregressive Transformers and RNNs like PixelCNN, WaveNet, and the new Transformer architecture - were exquisitely good at modeling discrete sequences (text, MIDI, pixel intensities). But applying them directly to images or audio at high resolution was impractical: PixelCNN at the pixel level was slow, and naive byte-level audio was hopelessly long.

The dream: get the strengths of discrete sequence modeling for continuous data. Compress images/audio/video into short sequences of discrete tokens, then train a powerful sequence model over those tokens.

The obstacle: how do you train a network to produce discrete tokens with gradient descent, when discrete operations aren't differentiable?

The Core Innovation: Vector Quantization in the Latent Space

VQ-VAE answers the question with three ingredients:

  1. A codebook - a learnable table of K embedding vectors, each of dimension D
  2. A quantization step - replace the encoder's continuous output with the nearest codebook vector
  3. A straight-through gradient estimator - pretend quantization is identity during the backward pass

Concretely:

class VQVAE:
    def __init__(self, K=512, D=64):
        self.encoder = ConvEncoder()
        self.decoder = ConvDecoder()
        self.codebook = Embedding(K, D)  # K codes, each D-dimensional

    def quantize(self, z_e):
        # z_e: continuous encoder output (B, H, W, D)
        # Find nearest codebook vector for each spatial position
        distances = pairwise_distance(z_e, self.codebook.weight)
        indices = argmin(distances, dim=-1)         # discrete!
        z_q = self.codebook(indices)                # (B, H, W, D)

        # Straight-through estimator: copy gradient from z_q to z_e
        z_q = z_e + (z_q - z_e).detach()
        return z_q, indices

    def forward(self, x):
        z_e = self.encoder(x)
        z_q, indices = self.quantize(z_e)
        x_recon = self.decoder(z_q)
        return x_recon, indices

The encoder maps a 256x256 image to (say) a 32x32 grid of vectors, each replaced by a codebook entry - yielding 32x32 = 1024 discrete token indices that represent the whole image.

The Loss Function

VQ-VAE's loss has three terms:

L = ||x - decoder(z_q)||^2              # reconstruction (trains encoder + decoder)
  + ||sg(z_e) - codebook_entries||^2    # codebook loss (trains codebook to match encodings)
  + beta * ||z_e - sg(codebook_entries)||^2 # commitment loss (encourages encoder to commit to codes)

sg(.) is the stop-gradient operator. beta is a hyperparameter (typically 0.25).

The three terms together ensure: - The autoencoder reconstructs well - The codebook entries move toward the encoder outputs - The encoder outputs stay close to existing codebook entries

Modeling the Prior

After training the autoencoder, you have a way to compress images into sequences of discrete tokens. Now train a powerful autoregressive model (PixelCNN in the original paper) over the token sequence, not the pixels. This prior model learns p(tokens), and at generation time you:

  1. Sample a token sequence from the prior
  2. Look up the corresponding codebook vectors
  3. Decode through the decoder to get the image

Critically, the prior model operates on a much shorter sequence (1024 tokens for a 256x256 image instead of 196,608 pixel values), making powerful sequence modeling tractable.

How VQ-VAE Avoids Posterior Collapse

A notorious problem with VAEs is posterior collapse: the decoder becomes powerful enough to ignore the latent and reconstruct from the prior alone, so the latent space carries no information.

VQ-VAE avoids this elegantly: - The latent is discrete and the encoder must commit to a specific code per spatial position - There's no KL term to a global prior pushing the encoder toward uninformative outputs - The decoder cannot ignore the latent because the code indices directly determine reconstruction

This is why VQ-VAE makes good use of its latent space where many VAEs fail to.

Key Results

Image Reconstruction

On ImageNet at 128x128: - Trained a VQ-VAE with K=512 codes - Achieved high-fidelity reconstruction - far sharper than continuous VAEs - Compressed each image to a 32x32 grid of token indices (a ~50x compression)

Generation

By training a PixelCNN prior on the discrete code grid, the authors generated novel images comparable to direct PixelCNN at the pixel level - but much faster, since the sequence is 50x shorter.

Audio (Most Striking Result)

VQ-VAE on raw audio (with a WaveNet decoder) demonstrated: - High-fidelity speech reconstruction - Speaker conversion - encode one speaker's audio to discrete codes, decode with a different speaker conditioning; the words are preserved but the voice changes - Strong evidence that the discrete codes captured phonemes without any phonetic supervision

The phoneme discovery was a major demonstration that the codebook learns semantically meaningful, modality-appropriate units.

Video

Trained on the action-conditioned video datasets; the discrete codes captured high-level structure while a small Transformer modeled the temporal evolution.

VQ-VAE-2 and Hierarchical Extensions

In 2019, the same authors released VQ-VAE-2, which scaled the approach using: - Hierarchical codes at multiple resolutions (top level for global structure, bottom for details) - Self-attention in the prior - Larger codebooks and bigger networks

VQ-VAE-2 produced ImageNet samples competitive with BigGAN - sharp, diverse, and at 256x256 and 1024x1024. This was the first non-adversarial method to reach this image quality.

Why This Paper Was Foundational

The discrete-token bridge VQ-VAE built between perception and language models had enormous downstream consequences:

DALL-E 1 (2021)

OpenAI's first text-to-image model used a VQ-VAE-like discrete image tokenizer plus a Transformer autoregressively predicting image tokens given text tokens. The whole approach is unthinkable without VQ-VAE's "images-as-token-sequences" trick.

VQ-GAN (2020)

Esser, Rombach, and Ommer added an adversarial loss and perceptual loss to VQ-VAE, dramatically improving reconstruction quality at high compression ratios. VQ-GAN became the standard image tokenizer for many downstream models.

Latent Diffusion / Stable Diffusion (2022)

Stable Diffusion runs diffusion in the latent space of an autoencoder. The autoencoder is essentially a VQ-VAE or VAE variant inspired by this lineage.

Neural audio codecs

SoundStream (Google) and Encodec (Meta) are essentially VQ-VAEs for audio, using residual vector quantization for higher fidelity. They power MusicLM, AudioLM, and most modern audio generation systems.

Speech and music models

AudioLM, MusicLM, and Whisper-style discrete-token audio approaches all build on the VQ-VAE foundation: compress audio to discrete tokens, then train language models over those tokens.

Video tokenizers

Sora's video tokenization, MAGVIT, and other modern video generators use VQ-VAE-derived spatiotemporal codebooks.

Conceptual Contributions

  1. Discrete latents are useful - they bridge perception and language model architectures
  2. Vector quantization is trainable with the straight-through estimator
  3. The commitment loss is essential for stable codebook learning
  4. Modality-appropriate units emerge - codes naturally align with phonemes in audio, patches in images
  5. Two-stage generation works well - autoencoder + prior is a flexible, powerful paradigm

Limitations

  • Codebook collapse - many codes go unused; later work (residual VQ, FSQ, look-up free quantization) addresses this
  • Quantization is lossy - high compression ratios cause visible degradation
  • Two-stage training is more involved than end-to-end approaches
  • Reconstruction-vs-generation trade-off - better compression hurts reconstruction
  • The codebook size is fixed at design time and affects capacity in non-obvious ways
  • Sampling can be slow with large autoregressive priors

Connections to Other Papers

  • VAE (#57): The direct ancestor. VQ-VAE replaces continuous Gaussian latents with discrete codebook entries, avoiding posterior collapse
  • Attention Is All You Need (#1): Once VQ-VAE compresses images/audio into token sequences, Transformers become the natural prior model
  • GANs (#2): Competing approach to image generation. VQ-VAE-2 caught up to BigGAN on image quality without adversarial training
  • Stable Diffusion (#7): Uses a VAE/VQ-VAE-style autoencoder to compress images so diffusion can run in a compact latent space
  • DDPM (#6) and Diffusion Models (#6): Define the continuous-token branch of generative modeling; VQ-VAE defines the discrete-token branch. Modern systems often combine them
  • DALL-E 1 and DALL-E 3 (#48): Built on VQ-VAE-style tokenization for the image side
  • GPT-3 (#4) and GPT-2 (#64): VQ-VAE makes it possible to apply GPT-style autoregressive modeling to images and audio by converting them to discrete tokens
  • Sora / DiT (#44): Modern video models use VQ-VAE-derived video tokenizers
  • MAE (#88): An alternative self-supervised approach - MAE reconstructs raw pixels without quantization; VQ-VAE compresses into discrete codes. Both are foundational autoencoder paradigms

Key Takeaways

  1. Discrete latents via vector quantization are trainable using the straight-through estimator and unlock the power of sequence-model priors for continuous data
  2. The codebook + commitment loss is the key recipe - encoder commits to nearby codes, codes move toward encoder outputs
  3. Two-stage generation (autoencoder, then prior over tokens) is dramatically more efficient than direct pixel-level autoregression
  4. Semantic units emerge from VQ training - phonemes in audio, recurring patches in images
  5. Foundational for multimodal AI - DALL-E, VQ-GAN, Stable Diffusion, neural audio codecs, MusicLM, AudioLM, Sora's tokenizer, and most modern multimodal systems trace back to the ideas introduced here