VQ-VAE

Overview

A VQ-VAE, or vector-quantized variational autoencoder, is an autoencoder whose hidden code is made of discrete symbols picked from a learned dictionary. The encoder produces a continuous vector, the model snaps it to the nearest entry in a codebook, and the decoder rebuilds the input from that entry. The result is a compact sequence of integer ids that stand for the data, much like words stand for ideas.

This design solves two problems. It avoids the blurry outputs and the posterior collapse that often affect the continuous VAE, and it gives a ready-made vocabulary of tokens that other models can use. A second model, such as an autoregressive transformer or a diffusion model, can then be trained over the token sequences to generate new data. This article explains the method with everyday analogies, covers the loss and the training tricks, gives a worked example, compares VQ-VAE with related generators, and describes how discrete tokens connect to models that learn from raw tables.

What it is

A VQ-VAE has four parts.

• The encoder maps an input to a grid or list of continuous vectors.
• The codebook is a table of K learned vectors, each of dimension D. Each row of the table is a "code".
• The quantizer replaces every encoder vector with its nearest codebook vector, and records the index of that code.
• The decoder rebuilds the input from the quantized vectors.

The latent representation is the list of code indices. With K codes, each index needs only log2(K) bits.

An everyday analogy

Think of a paint shop with 512 pre-mixed colours. When you want to copy a painting, you do not mix a new shade for every spot. You look at each spot and write down the number of the closest tin. Another person can rebuild the painting from your list of numbers and the same set of tins. The tins are the codebook. The list of numbers is the discrete latent. The quality depends on having a good set of tins.

What it is not

• It is not a sampling model by itself. After training, you still need a prior over code sequences to generate new data.
• It is not a standard VAE. There is no Gaussian latent and no KL term in the usual sense.
• It is not clustering, although the codebook update resembles k-means.

Where it came from

Discrete representations were attractive because language, speech and many structured signals are naturally discrete, and because continuous VAEs often ignored their latents when paired with strong decoders.

In 2017 the paper "Neural Discrete Representation Learning" by van den Oord, Vinyals and Kavukcuoglu introduced the VQ-VAE. It used vector quantization with a straight-through gradient trick and showed that a discrete latent could be learned without posterior collapse. The authors then fit an autoregressive prior over the codes to generate images, audio and video.

Follow-up work built on the idea:

• A second version added a hierarchy of codebooks and a stronger prior, improving image quality.
• Adding adversarial and perceptual losses gave sharper reconstructions in VQ-GAN style models.
• Residual and product quantization used several codes per position to raise capacity.
• Exponential moving average updates for the codebook improved stability.
• Text-to-image and speech systems used learned discrete tokens as the bridge between modalities, and latent diffusion systems such as Stable Diffusion used related autoencoders to compress images.
• Finite scalar quantization offered a simpler alternative with no explicit codebook.

For tables, research has explored tokenising numeric columns and rows into discrete codes. See /model/research for how this fits with other table model work.

How it works

Step 1: encode

The encoder maps an input x to continuous vectors z_e. For an image this is a grid of vectors. For a table row it might be a short list of vectors, one per column group or a single vector.

Step 2: quantize by nearest neighbour

For each vector z_e, find the nearest codebook entry:

• k = argmin over j of || z_e - e_j ||
• z_q = e_k

The index k is the discrete code. The vector z_q is passed to the decoder.

An analogy: it is like rounding a price to the nearest coin value. Any amount between two coins gets mapped to the same coin, so many slightly different inputs share a code.

Step 3: decode

The decoder takes the quantized vectors and rebuilds x. For numeric data, the output is a mean (and maybe a variance). For categories, probabilities.

Step 4: pass gradients with the straight-through trick

The nearest-neighbour step is not differentiable. The trick is to copy the gradient from the decoder input straight to the encoder output, as if the quantizer were the identity function. In code this is written as:

• z_q_st = z_e + (z_q - z_e).detach()

Forward, the value is z_q. Backward, the gradient goes to z_e.

An analogy: a teacher grades the final essay and passes the same comments back to the student's draft, ignoring the editor in between.

Step 5: compute the three loss terms

The VQ-VAE loss has three parts:

• Reconstruction loss. How well the decoder rebuilds x. This trains the encoder and decoder.
• Codebook loss. Moves the chosen code toward the encoder output. This trains the codebook. Written as || sg[z_e] - e ||^2 where sg stops gradients.
• Commitment loss. Keeps the encoder output close to the chosen code so it does not drift. Written as beta * || z_e - sg[e] ||^2.

Together:

• loss = recon + || sg[z_e] - e ||^2 + beta * || z_e - sg[e] ||^2

A common alternative replaces the codebook loss by an exponential moving average update, where each code becomes a running average of the encoder outputs assigned to it.

Step 6: avoid codebook collapse

A frequent failure is that only a few codes are ever used while the rest sit idle. This is called codebook collapse. Remedies:

• Initialise the codebook from k-means on early encoder outputs.
• Reset dead codes to random recent encoder outputs.
• Use a lower code dimension, or normalise vectors and use cosine distance.
• Use EMA updates.
• Monitor perplexity, which measures how many codes are effectively in use.

Step 7: train a prior over codes

After the autoencoder is trained, encode the training data into code sequences. Train a second model on those sequences:

• An autoregressive transformer predicts the next code given earlier ones.
• A masked model predicts hidden codes from visible ones.
• A discrete diffusion model denoises corrupted codes.

To generate, sample a code sequence from the prior and decode it. This split lets the autoencoder focus on detail and the prior focus on global structure. See transformers for the common prior architecture.

Key concepts and terms

A worked example

A numeric nearest-code lookup

Suppose a codebook with three codes in two dimensions:

• e0 = (0.0, 0.0)
• e1 = (1.0, 1.0)
• e2 = (2.0, 0.0)

An encoder outputs z_e = (1.2, 0.8).

• Distance squared to e0: 1.44 + 0.64 = 2.08
• Distance squared to e1: 0.04 + 0.04 = 0.08
• Distance squared to e2: 0.64 + 0.64 = 1.28

The nearest code is e1, so the code index is 1 and the decoder receives (1.0, 1.0). The commitment loss for this item is beta times 0.08. The codebook loss is 0.08, which pulls e1 slightly toward (1.2, 0.8).

A compact PyTorch quantizer

See PyTorch for the framework. This sketch shows the quantization layer with the three loss terms.

import torch
import torch.nn as nn
import torch.nn.functional as F

class VectorQuantizer(nn.Module):
    def __init__(self, num_codes=256, dim=32, beta=0.25):
        super().__init__()
        self.codebook = nn.Embedding(num_codes, dim)
        self.codebook.weight.data.uniform_(-1 / num_codes, 1 / num_codes)
        self.beta = beta

    def forward(self, z_e):
        # z_e: (batch, n_tokens, dim)
        flat = z_e.reshape(-1, z_e.shape[-1])
        d = (flat.pow(2).sum(1, keepdim=True)
             - 2 * flat @ self.codebook.weight.t()
             + self.codebook.weight.pow(2).sum(1))
        idx = d.argmin(dim=1)
        z_q = self.codebook(idx).view_as(z_e)
        codebook_loss = F.mse_loss(z_q, z_e.detach())
        commit_loss = F.mse_loss(z_e, z_q.detach())
        loss = codebook_loss + self.beta * commit_loss
        z_q = z_e + (z_q - z_e).detach()      # straight-through
        return z_q, loss, idx.view(z_e.shape[:-1])

A tiny VQ-VAE for table rows

class TableVQVAE(nn.Module):
    def __init__(self, d_in, n_tokens=4, dim=32, num_codes=128, hidden=256):
        super().__init__()
        self.n_tokens, self.dim = n_tokens, dim
        self.enc = nn.Sequential(nn.Linear(d_in, hidden), nn.ReLU(),
                                 nn.Linear(hidden, n_tokens * dim))
        self.vq = VectorQuantizer(num_codes, dim)
        self.dec = nn.Sequential(nn.Linear(n_tokens * dim, hidden), nn.ReLU(),
                                 nn.Linear(hidden, d_in))

    def forward(self, x):
        z_e = self.enc(x).view(-1, self.n_tokens, self.dim)
        z_q, vq_loss, idx = self.vq(z_e)
        x_hat = self.dec(z_q.reshape(x.shape[0], -1))
        recon = F.mse_loss(x_hat, x)
        return x_hat, recon + vq_loss, idx

Each row becomes four integer tokens. A downstream model can learn the joint distribution of those tokens.

Reading the result

1. Train until reconstruction error flattens.
2. Compute the code usage histogram and perplexity. Low perplexity signals collapse.
3. Encode the whole table into token sequences.
4. Train a small autoregressive model on the sequences.
5. Sample sequences, decode them, undo preprocessing, and compare with real rows.

Anomaly score from codes

@torch.no_grad()
def anomaly_score(model, x):
    x_hat, _, idx = model(x)
    recon_err = ((x_hat - x) ** 2).sum(dim=1)
    return recon_err, idx          # idx can also be scored by a prior model

Two signals are useful. The reconstruction error tells how well the codebook covers the row. The likelihood of the code sequence under the prior tells how typical the combination of codes is.

Strengths and limitations

Strengths

• No posterior collapse in the usual sense. The decoder must use the codes because they are its only input.
• Sharp reconstructions. The decoder is trained on a clean quantized signal, not on noisy samples.
• Discrete tokens. Other models, especially transformers, can work directly on the code sequences.
• Compression. Data become short lists of integers, which are cheap to store and index.
• Separation of concerns. The autoencoder models detail, the prior models structure.
• Stable objective. Training relies on regression-like losses with no adversary.

Limitations

• Two-stage training. You need the autoencoder first and a prior second, and mistakes in the first stage limit the second.
• Codebook collapse and dead codes. Needs monitoring and tricks.
• Information loss at quantization. Fine numeric precision is hard to keep with few codes.
• Not a density model. There is no direct likelihood of the data.
• Hyperparameter choices matter. Codebook size, code dimension, number of tokens and the commitment weight interact.
• Gradient approximation. The straight-through estimator is a heuristic.
• Table-specific difficulty. Mixed types, skewed values and exact integers do not always fit a nearest-code rule without careful design.

How it compares with alternatives

When to choose VQ-VAE:

• You want discrete tokens to feed a transformer.
• You want a compact, storable representation.
• You want sharper reconstructions than a plain VAE.

When to choose something else:

• You need exact densities, so use flows.
• You want a single-stage generator, so use a VAE or a GAN.
• You need top sample quality and can pay for it, so use diffusion.

Role in tabular foundation models

Tabular foundation models are networks pre-trained on very many tables, often synthetic ones, so that they can predict on a new table directly through in-context learning. Their inputs are raw tables: rows as examples, columns of mixed types, missing cells, and possibly no useful column names. VQ-VAE is not one of these models. It is an input representation and generation tool, and its ideas touch several design questions in table modelling. See /model/architecture for how table-native models represent rows and columns, and /model/benchmarks for how they are evaluated.

Tokenising cells and rows

Language models work on tokens. A table has numbers, categories, dates and gaps, which do not come as tokens. A VQ-VAE offers one way to produce them:

• Cell level: quantize each numeric value to one of K bins that are learned, not fixed. This is like learned bucketing.
• Row level: compress a whole row to a few codes, then model sequences of row codes.
• Column level: compress a column to a code that stands for its distribution, which can serve as a column "name" when real names are missing.

This relates to data tokenization and to numeric encoding choices in models such as FT-Transformer.

Mixed types and missing values

Discrete codes give a uniform interface for different types. A numeric column, a category and a missing marker can each be mapped to codes, and a sequence model can handle them together. A special "missing" code can be included in the codebook or reserved as a token. The decoder then learns to produce a plausible value when a missing code appears, which is a form of imputation.

Masked prediction on codes

Once rows are code sequences, a masked model can predict hidden codes from visible ones. This mirrors masked cell modeling, where cells are hidden and predicted. Working on codes rather than raw values can stabilise training for heavy-tailed numbers, because the targets are classes rather than unbounded reals.

Synthetic pre-training data

To produce synthetic tables for pre-training, a VQ-VAE with an autoregressive prior can learn the structure of real tables and generate new ones. The advantage is controllable diversity: sampling temperature and top-k choices on the code prior change how adventurous the rows are. See synthetic data generation.

Anomaly detection

Rows that map to rare code combinations, or that reconstruct poorly, are candidate anomalies. The prior's likelihood over code sequences gives a second view. These ideas fit beside anomaly detection methods like isolation forest. As with all deep scoring methods, validate on labelled cases.

In-context learning and continual learning

In-context learners such as TabPFN work on continuous inputs and do not require a codebook. Where context tables are large, a learned code could shrink them before use, though this is research, not standard practice. For continual learning, a codebook gives a stable vocabulary: new data can be encoded with old codes, which helps compare old and new tasks, and a prior can replay old code sequences.

Sector and semantic classification

Codes summarise columns and rows in a compact way. Frequency profiles of codes across a column can act as a fingerprint that a classifier uses to predict the sector of a table or the semantic type of a column, such as a date, a postal code or an amount. This is a plausible approach and an active research area, so confirm it against your own labelled tables.

Limits of the connection

• A VQ-VAE cannot predict a target from a few context rows by itself.
• Quantization loses precision, which matters for exact values.
• Foundation models that work on continuous values directly avoid this step and its tuning.

Practical guidance and best practices

Design choices

Training tips

1. Track reconstruction error, codebook perplexity and the share of dead codes every epoch.
2. Reset dead codes every few hundred steps early in training.
3. Warm up the encoder before turning on the full commitment loss if training is unstable.
4. Standardise inputs and apply quantile transforms to skewed numeric columns.
5. Keep the same preprocessing for the autoencoder and for later generation.

Evaluation checklist

• Reconstruction error on a held-out split.
• Perplexity close to the number of codes in use that you expect.
• Marginal and correlation fidelity of decoded samples.
• A downstream model trained on synthetic rows and tested on real rows.
• Distance of generated rows to the nearest training rows, to check for copying.
• For anomaly work, AUROC on labelled anomalies.

Configuration example

model:
  type: vq-vae
  codebook_size: 256
  code_dim: 32
  tokens_per_row: 8
  commitment_beta: 0.25
  codebook_update: ema
  ema_decay: 0.99
prior:
  type: transformer
  layers: 4
  heads: 4
  context: 8
data:
  source: example.internal/tables/orders.parquet
  numeric_transform: quantile
  missing_token: true
train:
  epochs: 100
  batch_size: 512
  dead_code_reset_every: 200

Variants and design alternatives

Exponential moving average codebook

• Each code keeps a running count of how many vectors it received and a running sum of those vectors.
• The code is set to sum divided by count.
• This avoids a separate optimiser step for the codebook and tends to be more stable.

Residual quantization

• Quantize the vector once, subtract the chosen code, and quantize the remainder with a second codebook.
• Repeat for several levels.
• Each position then has several codes, giving fine detail with small codebooks.

Product quantization

• Split the vector into equal parts.
• Give each part its own small codebook.
• The number of possible combinations grows as a power, so capacity is large while each codebook stays small.

Finite scalar quantization

• Round each dimension to one of a few fixed levels.
• No learned codebook, so no collapse.
• The implicit codebook is the product of the levels.

Choosing the number of tokens and codes

Checking the prior

After training the second stage, look at these items:

1. Negative log likelihood of held-out code sequences.
2. Samples decoded to rows and compared with real rows.
3. Effect of sampling temperature: low values give typical rows, high values give diverse but riskier rows.
4. Effect of top-k or nucleus sampling on rare categories.
5. Whether rare combinations of codes appear at about the right rate.

A worked view of dead codes

Suppose a codebook has 8 codes and, after one epoch, the counts of assigned vectors are:

Three codes are dead. A reset step would copy three recent encoder outputs into codes 4, 5 and 7. In the next epoch these codes can win assignments and learn useful positions. Perplexity would rise, which is the sign that capacity is being used.

Debugging guide

Common mistakes

• Ignoring codebook usage. A model with 512 codes may use ten of them.
• Using too large a codebook for a small table, which leaves most codes unused.
• Using too small a codebook and losing numeric precision.
• Skipping the commitment loss, which lets encoder outputs drift far from the codes.
• Forgetting the stop-gradient operations in the codebook and commitment terms.
• Training the prior on codes from a poor autoencoder, so the generator inherits its errors.
• Mixing preprocessing between the training and generation stages.
• Treating the codes as meaningful categories without checking. A code is a learned cluster, not a named concept.
• Quantizing exact integers or identifiers that must be preserved exactly. Handle them outside the codebook.
• Judging only by reconstruction error. Check the prior and the generated data too.

Frequently asked questions

What does VQ-VAE stand for?

VQ-VAE stands for vector-quantized variational autoencoder. It is an autoencoder whose latent code is made of discrete entries from a learned codebook, chosen by nearest-neighbour lookup.

How is a VQ-VAE different from a VAE?

A VAE uses a continuous Gaussian latent with a KL penalty, and you sample from a simple prior. A VQ-VAE uses a discrete codebook and a separate prior model over code sequences, so generation takes two stages.

What is a codebook?

A codebook is a learned table of vectors, each with an integer index. The encoder output is replaced by its nearest entry, so the whole latent is a list of indices into this table.

What is the straight-through estimator?

It is a trick for training through a non-differentiable step. In the forward pass the quantizer picks the nearest code, but in the backward pass the gradient is copied unchanged to the encoder output, as if no quantization had happened.

What is codebook collapse?

It is a failure where only a few codes are ever chosen and most of the codebook goes unused. This lowers the model's capacity. Fixes include EMA updates, dead code resets, k-means initialisation and monitoring perplexity.

Can a VQ-VAE generate new data?

Not by itself. It needs a prior over code sequences, such as an autoregressive transformer. You sample a code sequence from the prior and pass it to the decoder to get a new sample.

Can VQ-VAE be used for tabular data?

Yes. Rows or columns can be compressed to short code sequences, numeric values can be quantized with learned bins, and missing values can get their own token. Care is needed for exact integers, skewed numbers and mixed types.

Is VQ-VAE good for anomaly detection?

It can help. Reconstruction error and the prior's likelihood of the code sequence both indicate how typical a row is. Results depend on the data and the model, so test against labelled anomalies and simple baselines.

How does VQ-VAE relate to diffusion models?

Both are often combined. Autoencoders with quantized or continuous latents compress data, and a diffusion model then works in the smaller latent space. Discrete diffusion models can also operate directly on VQ code sequences. See DDPM for the diffusion side.

Is a VQ-VAE a tabular foundation model?

No. It is a representation and generation method trained on one data source at a time. It can provide tokens, compression and anomaly signals around table models, but a tabular foundation model is pre-trained across many tables and predicts through in-context learning.

Glossary

• Autoregressive model: a model that predicts each element from the earlier ones.
• Codebook: a learned dictionary of vectors.
• Codebook collapse: underuse of the codebook.
• Commitment loss: a penalty that keeps encoder outputs near their chosen codes.
• Dead code: a code that is never used.
• Discrete latent: a hidden representation made of integers from a finite set.
• EMA update: a moving-average update rule for codes.
• Finite scalar quantization: rounding each latent dimension to a few levels instead of using a codebook.
• Perplexity: a measure of the effective number of codes in use.
• Prior: the model of code sequences used for generation.
• Product quantization: splitting a vector into parts and quantizing each part separately.
• Quantization: replacing a continuous value by the nearest member of a finite set.
• Residual quantization: quantizing the remaining error with further codebooks.
• Straight-through estimator: a gradient approximation that bypasses a discrete step.
• Token: a discrete symbol from a vocabulary.

Summary

A VQ-VAE compresses data into short lists of integer codes taken from a learned codebook. Training combines reconstruction, a codebook update and a commitment penalty, with a straight-through trick to pass gradients across the nearest-neighbour step. The discrete latent avoids the collapse problems of a continuous VAE and gives other models a vocabulary to work with, but it adds a second stage for the prior and needs care to keep the codebook in use. For tables, it offers learned tokenisation of cells, rows and columns, a uniform way to handle mixed types and missing values, and anomaly signals from reconstruction and code likelihood. It supports, and does not replace, models that learn directly from raw tables. See also /model/research, /model/architecture and /model/benchmarks.

Related

It belongs to Architectures in my stack, under AI and machine learning. Others in the same group: