E2E: Vector Quantized VAE - End-to-End Implementations

Implement E2E: Vector Quantized VAE in Python with NumPy. hard exercise with examples, hints, and instant feedback. No setup or signup required.

End-to-End VQ-VAE

What you'll build

A small but complete VQ-VAE: a vector-quantization layer with its losses, an encoder and decoder around it, the total loss, and codebook-usage statistics. It extends VQ-VAE Vector Quantization (Generative Models) into a model.

Background

A VQ-VAE autoencodes through a discrete bottleneck: each encoder vector is replaced by its nearest codebook entry. The indices are a compressed, tokenized version of the input that a Transformer can model later (this is how VQGAN and many audio and image tokenizers work). Unlike a VAE there is no KL term; the codebook and commitment losses keep encoder outputs and codes close.

Architecture (simplified to MLPs on 4×4 inputs)

Stage	Operation	Shape
input	`x`	`(batch, C, 4, 4)`
flatten	`x.reshape(batch, -1)`	`(batch, 16·C)`
encoder	`relu(· @ enc_w1 + enc_b1) @ enc_w2 + enc_b2`	`(batch, D)`
as a 1×1 latent grid	reshape	`z_e: (batch, 1, 1, D)`
quantize	nearest row of `vq.embedding`	`z_q: (batch, 1, 1, D)`, `indices: (batch, 1, 1)`
decoder	`relu(flat(z_q) @ dec_w1 + dec_b1) @ dec_w2 + dec_b2`	`(batch, 16·C)`
reshape		`x_recon: (batch, C, 4, 4)` (no output activation)

D = embedding_dim, K = num_embeddings. Both constructors are given: they seed np.random.seed(42) and create the weights listed above (the codebook vq.embedding is N(0, 0.1²) of shape (K, D)).

Losses

sg[·] is stop-gradient; in a NumPy forward pass the two VQ terms have the same value:

recon_loss      = mean((x − x_recon)²)
codebook_loss   = mean((sg[z_e] − z_q)²)        # trains the codebook
commitment_loss = β · mean((z_e − sg[z_q])²)     # trains the encoder; β is already included here
total_loss      = recon_loss + codebook_loss + commitment_loss

With autograd, the decoder's gradient reaches the encoder through the straight-through estimator z_q_st = z_e + sg[z_q − z_e] (forward value z_q); the PyTorch track implements it.

Method contract

Method	Returns
`VectorQuantizer.quantize(z_e)`	tuple `(z_q, indices, distances)`: `z_q` like `z_e`; `indices` int `(batch, H, W)` (ties → lowest index); `distances` = squared distance to the chosen code, `(batch, H, W)`
`VectorQuantizer.compute_loss(z_e, z_q, beta=0.25)`	dict `{'codebook_loss', 'commitment_loss'}` (floats)
`VQVAE.encode(x)`	`z_e (batch, 1, 1, D)`
`VQVAE.decode(z_q)`	`x_recon (batch, C, 4, 4)`
`VQVAE.forward(x)`	dict `{'z_e', 'z_q', 'indices', 'x_recon'}` (decode `z_q`, not `z_e`)
`VQVAE.compute_loss(x, x_recon, z_e, z_q, beta=0.25)`	dict `{'total_loss', 'recon_loss', 'codebook_loss', 'commitment_loss'}`
`VQVAE.get_codebook_usage(indices)`	int array `(K,)`: how many times each code was chosen (zeros for unused codes)

Constraints

NumPy only; vectorised nearest-neighbour search (no loop over positions or codes).

Follow-up

Codebook collapse: few codes get all the assignments. Usage statistics reveal it; dead-code restarts fix it.
EMA codebook updates replace the codebook loss: N_k ← γN_k + (1−γ)n_k, m_k ← γm_k + (1−γ)Σ z_e assigned to k, e_k = m_k / N_k.

Examples

Input: vq.embedding = [[0, 0], [10, 10], [5, 5]]; vq.quantize(np.array([[[[1, 1], [9, 8], [4, 6]]]]))
Output: indices = [[[0, 1, 2]]], distances = [[[2.0, 5.0, 2.0]]]

[9, 8] is 1² + 2² = 5 from [10, 10]; the others are 2 from their nearest code

Input: vq.compute_loss(z_e, z_q) for the example above
Output: {'codebook_loss': 1.5, 'commitment_loss': 0.375}

Mean squared error 9 / 6 = 1.5, and β = 0.25 times that

Input: model.get_codebook_usage(np.array([[[0]], [[2]], [[2]]])) with K = 6
Output: array([1, 0, 2, 0, 0, 0])

Code 0 used once, code 2 twice, the rest never

Practice in Python

import numpy as np

class VectorQuantizer:
    def __init__(self, num_embeddings: int, embedding_dim: int):
        """
        Vector quantization layer (given).

        Args:
            num_embeddings: Size of codebook (K)
            embedding_dim: Dimension of each embedding (D)
        """
        np.random.seed(42)
        self.embedding = np.random.randn(num_embeddings, embedding_dim) * 0.1   # (K, D)
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim

    def quantize(self, z_e: np.ndarray) -> tuple:
        """
        Replace every vector of z_e by its nearest codebook entry.

        Args:
            z_e: Encoder output (batch, H, W, D)

        Returns:
            z_q: Quantized vectors (same shape as z_e)
            indices: Codebook indices (batch, H, W); ties go to the lowest index
            distances: Squared distance to the chosen code (batch, H, W)
        """
        # Your code here
        pass

    def compute_loss(self, z_e: np.ndarray, z_q: np.ndarray, beta: float = 0.25) -> dict:
        """
        Returns:
            {'codebook_loss': mean((z_e - z_q)**2),
             'commitment_loss': beta * mean((z_e - z_q)**2)}
        """
        # Your code here
        pass


class VQVAE:
    def __init__(self, input_channels: int, hidden_dim: int,
                 num_embeddings: int, embedding_dim: int):
        """
        VQ-VAE for (batch, input_channels, 4, 4) inputs, with MLP encoder/decoder (given).
        """
        np.random.seed(42)
        self.input_channels = input_channels
        self.hidden_dim = hidden_dim
        self.embedding_dim = embedding_dim

        # Encoder: (C*16) -> hidden_dim -> embedding_dim
        self.enc_w1 = np.random.randn(input_channels * 16, hidden_dim) * 0.1
        self.enc_b1 = np.zeros(hidden_dim)
        self.enc_w2 = np.random.randn(hidden_dim, embedding_dim) * 0.1
        self.enc_b2 = np.zeros(embedding_dim)

        # Vector quantizer (re-seeds, so its codebook does not depend on the encoder size)
        self.vq = VectorQuantizer(num_embeddings, embedding_dim)

        # Decoder: embedding_dim -> hidden_dim -> (C*16)
        self.dec_w1 = np.random.randn(embedding_dim, hidden_dim) * 0.1
        self.dec_b1 = np.zeros(hidden_dim)
        self.dec_w2 = np.random.randn(hidden_dim, input_channels * 16) * 0.1
        self.dec_b2 = np.zeros(input_channels * 16)

    def encode(self, x: np.ndarray) -> np.ndarray:
        """
        x: (batch, C, 4, 4) -> flatten -> relu(x @ enc_w1 + enc_b1) @ enc_w2 + enc_b2

        Returns:
            z_e: (batch, 1, 1, embedding_dim)
        """
        # Your code here
        pass

    def decode(self, z_q: np.ndarray) -> np.ndarray:
        """
        z_q: (batch, 1, 1, D) -> flatten -> relu(z @ dec_w1 + dec_b1) @ dec_w2 + dec_b2

        Returns:
            x_recon: (batch, C, 4, 4)
        """
        # Your code here
        pass

    def forward(self, x: np.ndarray) -> dict:
        """
        Encode, quantize, decode z_q.

        Returns:
            {'z_e', 'z_q', 'indices', 'x_recon'}
        """
        # Your code here
        pass

    def compute_loss(self, x: np.ndarray, x_recon: np.ndarray,
                     z_e: np.ndarray, z_q: np.ndarray, beta: float = 0.25) -> dict:
        """
        Returns:
            {'total_loss', 'recon_loss', 'codebook_loss', 'commitment_loss'},
            total = recon + codebook + commitment (commitment already includes beta)
        """
        # Your code here
        pass

    def get_codebook_usage(self, indices: np.ndarray) -> np.ndarray:
        """
        Returns:
            usage: int array (num_embeddings,), count of each code in indices
        """
        # Your code here
        pass

The interactive editor loads with Python and NumPy. Run your implementation against 7 tests, or use the progressive hints.

More exercises in this section →