Vqvae
VQ-VAE (Vector Quantized Variational Autoencoder) is the generative model that learns discrete latent representations by mapping encoder outputs to the nearest vector in a learned codebook — replacing the continuous Gaussian latent space of standard VAEs with a finite set of embedding vectors, enabling high-fidelity reconstruction, serving as the foundation for modern image/audio generation systems like DALL-E and SoundStream, and bridging continuous neural representations with discrete token-based generation.
Architecture
1. Encoder: Input x → continuous latent representation z_e(x).
2. Vector Quantization: Map z_e to nearest codebook vector: $z_q = e_k$ where $k = \arg\min_j z_e - e_j _2$.
3. Decoder: Reconstruct input from quantized latent: x̂ = Decoder(z_q).
4. Codebook: K learnable embedding vectors {e₁, e₂, ..., eₖ}, typically K=512-8192.
Training Loss
$L = x - \hat{x} _2^2 + \text{sg}[z_e] - e_k _2^2 + \beta z_e - \text{sg}[e_k] _2^2$
- Term 1: Reconstruction loss.
- Term 2: Codebook loss — move codebook vectors toward encoder outputs. (sg = stop gradient.)
- Term 3: Commitment loss — encourage encoder to commit to codebook vectors.
Straight-Through Estimator
- Problem: argmin (nearest neighbor lookup) is non-differentiable.
- Solution: Copy gradients from decoder input to encoder output, skipping the quantization.
- Forward: z_q = nearest codebook vector. Backward: gradients flow as if z_q = z_e.
VQ-VAE-2 (Hierarchical)
- Two-level codebook: Top level captures global structure, bottom level captures details.
- Top latent: Low resolution (32×32) → overall layout, color scheme.
- Bottom latent: High resolution (64×64) → fine details, textures.
- Two-stage generation: Train PixelCNN/Transformer on top → condition bottom on top.
Applications
| Application | System | How VQ-VAE Is Used |
|---|---|---|
| Image generation | DALL-E (v1) | VQ-VAE encodes images to discrete tokens → Transformer generates tokens |
| Audio compression | SoundStream, Encodec | VQ-VAE with residual quantization → neural audio codec |
| Video generation | VideoGPT | VQ-VAE for video frames → Transformer for temporal generation |
| Music generation | MusicGen, Jukebox | VQ-VAE tokenizes audio → language model generates music |
| Image tokenizer | LlamaGen, Parti | VQ tokenizer → autoregressive image generation |
Residual Vector Quantization (RVQ)
- Instead of single codebook: Apply VQ in multiple stages, each quantizing the residual error.
- Stage 1: Quantize z_e → residual r₁ = z_e - z_q₁.
- Stage 2: Quantize r₁ → residual r₂ = r₁ - z_q₂.
- Repeat for D stages → total representation: z_q₁ + z_q₂ + ... + z_qD.
- Used in neural audio codecs (SoundStream, Encodec) for variable-bitrate compression.
VQ-VAE is the foundational architecture that enabled the tokenization of continuous signals for discrete generation — by converting images, audio, and video into sequences of codebook indices, it allows powerful autoregressive transformers and language models to generate these modalities as naturally as generating text.