VQ-VAE Vector Quantization - Generative Models

VQ-VAE Vector Quantization

What you'll build

The discrete bottleneck of a VQ-VAE (van den Oord et al., 2017): snap every encoder vector to its nearest codebook entry, and compute the two auxiliary losses that train the codebook and the encoder.

Background

A VQ-VAE replaces the Gaussian latent of a VAE with a discrete one: a grid of indices into a learned codebook of K vectors. Discrete codes are what image and audio tokenizers (VQGAN, SoundStream, the image tokens of many multimodal LLMs) produce, so a Transformer can then model them like text.

Algorithm: vq_quantize

1. Flatten the encoder output z_e from (B, H, W, D) to N = B·H·W vectors of size D.
2. Compute all squared distances to the K codes with the expansion ||z − e_k||² = ||z||² − 2·z·e_k + ||e_k||², giving an (N, K) matrix.
3. indices = argmin over K (ties go to the lowest index, as np.argmin does).
4. z_q = codebook[indices], reshaped back to (B, H, W, D); indices to (B, H, W).

Losses: vq_loss

sg[·] is stop-gradient (.detach() in PyTorch):

codebook_loss   = mean( (sg[z_e] − z_q)² )          # pulls codes toward encoder outputs
commitment_loss = β · mean( (z_e − sg[z_q])² )      # keeps the encoder near its codes

Means are over all elements. In a forward pass both squared errors have the same value (sg only changes gradients), so commitment_loss = β · codebook_loss. The full VQ-VAE loss is reconstruction + codebook_loss + commitment_loss. The straight-through estimator z_q = z_e + sg[z_q − z_e] copies the decoder's gradient to the encoder; it only matters in an autograd framework (see the PyTorch track and E2E: VQ-VAE).

Inputs

• z_e: float array (B, H, W, D)
• codebook: float array (K, D)
• beta: float, commitment weight (default 0.25)

Expected return format

• vq_quantize returns (z_q, indices): z_q float (B, H, W, D), indices int (B, H, W)
• vq_loss returns (codebook_loss, commitment_loss): Python floats rounded to 4 decimals

Constraints

NumPy only. Vectorise the distance computation: no loop over positions or codes.

Follow-up

• What is codebook collapse, and how do EMA codebook updates or dead-code restarts fix it?
• Memory of the (N, K) distance matrix for a 32×32 grid, batch 64 and K = 8192? (64·1024·8192·4 bytes ≈ 2 GB, so chunk it.)

Input: z_e = [[[[1, 1], [9, 8], [4, 6]]]] (shape (1, 1, 3, 2)), codebook = [[0, 0], [10, 10], [5, 5]]
Output: indices = [[[0, 1, 2]]], z_q = [[[[0, 0], [10, 10], [5, 5]]]]

Input: vq_loss(z_e, z_q, beta=0.25) for the example above
Output: (1.5, 0.375)

import numpy as np

def vq_quantize(z_e, codebook):
    """
    Vector quantization for VQ-VAE.

    Args:
        z_e: Encoder output, shape (batch, H, W, embedding_dim)
        codebook: Code vectors, shape (num_embeddings, embedding_dim)

    Returns:
        z_q: Nearest code for every position, shape (batch, H, W, embedding_dim)
        indices: Index of that code, int array of shape (batch, H, W);
                 ties go to the lowest index
    """
    # Your code here
    pass


def vq_loss(z_e, z_q, beta=0.25):
    """
    Codebook and commitment losses (forward values).

    Args:
        z_e: Encoder output, shape (batch, H, W, embedding_dim)
        z_q: Quantized vectors, same shape
        beta: Commitment loss weight (default 0.25)

    Returns:
        (codebook_loss, commitment_loss): Python floats rounded to 4 decimals,
        codebook_loss = mean((z_e - z_q)**2), commitment_loss = beta * that
    """
    # Your code here
    pass