VQ-VAE

The Vector-Quantised Variational Auto-Encoder replaces the continuous Gaussian bottleneck of a VAE with a discrete one. An encoder $z_e(x)$ maps an image to a spatial grid of $D$-dimensional vectors, and each grid position is snapped to its nearest neighbour in a learned codebook $e \in \mathbb{R}^{K \times D}$:

The decoder then reconstructs from $z_q(x)$. Because the $\arg\min$ has zero gradient almost everywhere, training uses the straight-through estimator — the forward pass carries the quantised vector while the backward pass copies the decoder's gradient directly onto the encoder output, which is implemented as $z_q = z_e + \mathrm{sg}[z_q - z_e]$ with $\mathrm{sg}$ the stop-gradient operator.

The codebook receives no gradient from that path, so it is trained by two additional terms. The codebook loss pulls each selected entry toward the encoder output it was matched to, and the commitment loss pulls the encoder output toward the entry it chose, preventing the encoder's output space from growing without bound. The full objective is

with $\beta = 0.25$ in the paper's experiments. Note the asymmetry: the codebook term is unweighted because the codebook has no other learning signal, while $\beta$ scales how strongly the encoder is held to its current assignment.

Unlike a Gaussian VAE there is no KL term to anneal. The prior over the discrete latents is held uniform during training, which makes its KL the constant $\log K$ and removes posterior collapse as a failure mode entirely — the paper's central practical claim. A faithful generative prior is fit afterwards as a separate autoregressive model (PixelCNN over the latent grid), so the auto-encoder trained here is a representation learner and a discrete tokeniser rather than a stand-alone sampler. That tokeniser role is what later made the architecture load-bearing across discrete generative modelling.