VAE

Overview

A VAE, or variational autoencoder, is a neural network that learns to compress data into a small set of hidden variables and to rebuild the data from them. Unlike a plain autoencoder, it does not map each input to a single point. It maps each input to a probability distribution over the latent space. A second part of the loss pushes those distributions toward a simple prior, usually a standard Gaussian. The result is a smooth latent space from which you can draw new samples.

This makes the VAE both a representation learner and a generative model. You can use the encoder to get compact features, the decoder to generate new data, and the reconstruction error to flag unusual inputs. For tables, VAEs are a common choice for synthetic rows, missing value imputation and anomaly scoring. This article explains the idea with everyday analogies, walks through the loss, gives code, and compares the VAE with autoencoders, GANs, DDPM and normalizing flows.

What it is

A VAE has three parts.

• The encoder takes an input x and outputs the parameters of a distribution over a latent vector z. Usually these are a mean and a variance for each latent dimension.
• The latent space is where z lives. It has fewer dimensions than the input in the usual setup.
• The decoder takes a sample z and outputs the parameters of a distribution over x. For numbers, this is often a Gaussian. For categories, a categorical distribution.

It is trained to maximise a lower bound on the log likelihood of the data, called the evidence lower bound, or ELBO.

An everyday analogy

Imagine describing faces to a sketch artist over the phone. A plain autoencoder is like giving exact measurements: the artist redraws the one face. A VAE is like describing a face with some fuzziness, such as "nose about this long, give or take". Because the descriptions are fuzzy and tied to a common vocabulary, nearby descriptions give similar faces, and an invented description gives a believable new face. The pull toward a shared simple prior is like making all painters use the same vocabulary, so any description from that vocabulary makes sense.

What it is not

• It is not a plain autoencoder. The latent space is probabilistic and regularised.
• It is not a GAN. There is no discriminator, and training uses a likelihood bound.
• It is not an exact-likelihood model. The ELBO is a bound, not the true likelihood.

Where it came from

Autoencoders existed for decades as a way to learn compressed codes. They were good at reconstruction but their latent spaces had gaps, so sampling from them gave poor results.

In 2013 and 2014 two groups introduced the variational autoencoder. The paper "Auto-Encoding Variational Bayes" by Kingma and Welling, and a closely related paper on stochastic backpropagation by Rezende, Mohamed and Wierstra, showed how to train a latent-variable model with neural networks using a gradient-friendly trick. See variational inference for the background on approximate posteriors.

Key milestones that followed:

• The beta-VAE weighted the KL term more heavily to encourage disentangled factors.
• Conditional VAEs added labels or other inputs to both encoder and decoder.
• Importance weighted autoencoders tightened the bound with multiple samples.
• Hierarchical VAEs stacked several layers of latents.
• VQ-VAE replaced the continuous latent with a discrete codebook.
• Latent diffusion systems used VAE-style autoencoders to compress images before applying a diffusion model, as in Stable Diffusion.
• For tables, models such as TVAE adapted the VAE to mixed numeric and categorical columns, and VAEs became a standard baseline for tabular synthesis.

See /model/research for how this line of work relates to modern table models.

How it works

Step 1: encode to a distribution

For an input row x, the encoder network outputs two vectors: a mean mu and a log variance. Together they define a Gaussian over z. Using the log variance keeps the variance positive and numerically stable.

Step 2: sample with the reparameterization trick

We need to draw z from that Gaussian and still send gradients back through the encoder. Sampling is a random operation, so gradients do not pass through it directly. The trick is to write the sample as:

• z = mu + sigma * eps, where eps is drawn from a standard Gaussian

Now the randomness sits in eps, which does not depend on the network weights. Gradients flow through mu and sigma as usual.

An analogy: instead of asking a dice-rolling machine to produce a number in your range, you roll a standard die separately and then scale and shift the result. The scaling and shifting are things you can adjust and learn.

Step 3: decode

The decoder takes z and outputs the parameters of the data distribution. For a numeric column this is a mean (and sometimes a variance). For a categorical column it is a set of probabilities. A reconstruction of the input is obtained from these parameters.

Step 4: compute the loss

The loss has two terms:

• Reconstruction term: the negative log likelihood of x under the decoder. It rewards faithful rebuilding.
• KL term: the Kullback-Leibler divergence between the encoder's distribution and the prior. It rewards a latent space that matches the simple prior.

The total is the negative ELBO:

• loss = reconstruction loss + KL(q(z | x) || p(z))

For a Gaussian encoder and a standard Gaussian prior, the KL term has a closed form:

• KL = 0.5 * sum( mu^2 + sigma^2 - log(sigma^2) - 1 )

Step 5: balance the two forces

The two terms pull in opposite directions.

• Too much weight on reconstruction yields sharp reconstructions but a patchy latent space.
• Too much weight on KL yields a tidy latent space but blurry or generic outputs. In the extreme, the decoder ignores z. This is called posterior collapse.

A common remedy is KL annealing: start with a small KL weight and raise it slowly. Another is the free bits trick, which stops the KL pressure on a dimension once it carries a minimum amount of information.

Step 6: generate or analyse

• Generate: draw z from the prior, then decode.
• Embed: run the encoder and use mu as a feature vector.
• Reconstruct: encode then decode, and compare with the original.
• Interpolate: move between two latent codes and decode points along the line.
• Score: compute the ELBO or reconstruction error for a row to judge how typical it is.

Why the latent space is smooth

Because each input is mapped to a cloud rather than a point, clouds of similar inputs overlap. The decoder must handle every point in an overlap, so it learns to give similar outputs there. The KL pressure keeps all clouds near the same region, so no big empty gaps remain. A smooth space is what makes sampling and interpolation meaningful.

Key concepts and terms

A worked example

A numeric check of the KL term

Suppose a one-dimensional latent with mu = 1.0 and sigma = 0.5 for some input.

• mu squared is 1.0
• sigma squared is 0.25
• log of sigma squared is about -1.386
• Inside the sum: 1.0 + 0.25 - (-1.386) - 1 = 1.636
• KL = 0.5 * 1.636 = 0.818

If the encoder instead returned mu = 0 and sigma = 1, the KL would be exactly 0. The KL term measures how far the encoder's cloud is from the standard cloud.

A compact PyTorch VAE for numeric table columns

See PyTorch for the framework. This sketch uses a Gaussian decoder with unit variance for simplicity.

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

class TableVAE(nn.Module):
    def __init__(self, d_in, d_latent=8, hidden=128):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(d_in, hidden), nn.ReLU(),
                                 nn.Linear(hidden, hidden), nn.ReLU())
        self.mu = nn.Linear(hidden, d_latent)
        self.logvar = nn.Linear(hidden, d_latent)
        self.dec = nn.Sequential(nn.Linear(d_latent, hidden), nn.ReLU(),
                                 nn.Linear(hidden, hidden), nn.ReLU(),
                                 nn.Linear(hidden, d_in))

    def encode(self, x):
        h = self.enc(x)
        return self.mu(h), self.logvar(h)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + std * eps

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.dec(z), mu, logvar

def vae_loss(x, x_hat, mu, logvar, beta=1.0):
    recon = F.mse_loss(x_hat, x, reduction="none").sum(dim=1)
    kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(dim=1)
    return (recon + beta * kl).mean(), recon.mean(), kl.mean()

Training loop and use

model = TableVAE(d_in=12)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(100):
    beta = min(1.0, epoch / 30)          # KL annealing over 30 epochs
    for (x,) in loader:
        x_hat, mu, logvar = model(x)
        loss, r, k = vae_loss(x, x_hat, mu, logvar, beta)
        opt.zero_grad(); loss.backward(); opt.step()

# generate new rows
with torch.no_grad():
    z = torch.randn(1000, 8)
    new_rows = model.dec(z)

# anomaly score: reconstruction error per row
with torch.no_grad():
    x_hat, mu, logvar = model(test_x)
    score = ((x_hat - test_x) ** 2).sum(dim=1)

Handling mixed columns

For a table with numbers and categories, a typical recipe is:

1. Standardise numeric columns, or transform them with a quantile or mixture-of-Gaussians encoding.
2. One-hot encode categories.
3. Use a Gaussian output head for numbers and a softmax head per categorical column.
4. Sum the per-column losses. Use mean squared error or Gaussian negative log likelihood for numbers, and cross entropy for categories.
5. For missing cells, multiply each column's loss by an observed mask so unobserved cells contribute nothing.

Imputation with a trained VAE

def impute(model, x, mask, iters=20):
    # mask is 1 for observed, 0 for missing; x has any placeholder in missing cells
    x_fill = x.clone()
    for _ in range(iters):
        mu, logvar = model.encode(x_fill)
        x_hat = model.dec(mu)
        x_fill = mask * x + (1 - mask) * x_hat   # keep observed, update missing
    return x_fill

This simple loop encodes the current guess, decodes, and refreshes only the missing cells. More careful variants sample z and use importance weighting.

Strengths and limitations

Strengths

• Stable training. A single likelihood-based loss, no adversarial game.
• Fast sampling. One pass through the decoder creates a sample.
• Useful encoder. You get compact, smooth features for downstream tasks.
• Probabilistic. It handles noise and uncertainty and supports missing data naturally.
• Simple to extend. Conditions, mixed types and masks are easy to add.
• Reconstruction-based anomaly scores come for free.
• Small and cheap. Works well on modest hardware and moderately sized tables.

Limitations

• Blurry outputs. With simple decoders the average over plausible outputs is blurry, which for tables shows as over-smoothed values and weak tails.
• Posterior collapse. Powerful decoders may ignore the latent code.
• Loose bound. The ELBO may be far from the true likelihood.
• Prior mismatch. If aggregated posteriors do not match the Gaussian prior, samples from the prior land in poor regions. This is called the prior hole problem.
• Mixed types need engineering. Skewed numbers and many-level categories need care.
• Disentanglement is not guaranteed. Latent dimensions may not match meaningful factors.
• Sample quality can lag behind diffusion models on complex data.

How it compares with alternatives

Simple rules of thumb:

• Need features and a quick generator? Use a VAE.
• Need exact densities? Use flows.
• Need the sharpest samples and have compute? Use diffusion.
• Need discrete tokens for a sequence model? Use a VQ-VAE.

Role in tabular foundation models

Tabular foundation models are pre-trained networks that work directly on raw tables. A table has rows as examples, columns of mixed type, missing cells, and often no useful column names. These models rely on pre-training over very many tables, frequently synthetic, and predict on new tables by in-context learning. A VAE is not a foundation model, but its ideas appear in many parts of the table pipeline. See /model/architecture for how table-native models organise rows and columns, and /model/benchmarks for how they are measured.

Learning from rows with mixed types

A VAE for tables must deal with the same issues as a foundation model: numbers on different scales, categories with many levels, dates, and missing cells. The standard answers carry over: per-column encoders, a shared latent space, per-column output heads, and masks. Studying how a VAE handles these shows what any table model needs from its input layer.

Missing values and imputation

Because the decoder defines a distribution, a VAE can treat missing cells as latent. Training with a mask on the loss lets the model learn from incomplete rows. At test time, the encoder can use only observed cells, and the decoder fills in the rest. Sampling several z values gives multiple imputations. This ties directly to imputation networks, and is a cousin of masked cell modeling, where a transformer predicts hidden cells from visible ones.

Anomaly detection

A VAE trained on normal rows gives each new row a reconstruction error and an ELBO. Rows far from the training distribution tend to have high error or low ELBO. This is a standard reconstruction-based method in anomaly detection. It is easy to run and combines well with isolation forest. Two cautions apply: a flexible decoder may reconstruct anomalies well, and likelihoods from deep generative models can be misleading for out-of-distribution data. Validate scores on labelled examples.

Synthetic data for pre-training

Pre-training a table model on synthetic tables requires diverse, realistic generators. A VAE fitted to real tables can contribute extra rows with realistic column dependencies. It is faster to sample than diffusion, though less sharp. See synthetic data generation for the broader design space.

Embeddings, in-context learning and continual learning

The encoder gives row embeddings that can feed a classifier or a nearest-neighbour search. In-context learners such as TabPFN learn their own representations during pre-training and do not rely on a VAE, but a VAE can offer a compact summary of large context tables. For continual learning, a VAE can act as a generative memory: sample pseudo-rows of old data and mix them with new data to reduce forgetting. This is called generative replay, and its success depends on how well the VAE covers the old distribution.

Sector and semantic classification

Latent codes of rows or columns can be clustered or classified to group tables by sector or by the semantic meaning of columns. A conditional VAE can generate rows conditioned on a sector label, which can help balance rare classes. This is a practical technique but not a solved problem, so check it on your own data.

What a VAE does not give you

• No direct in-context prediction from a handful of examples.
• No guarantee of exact likelihoods.
• No automatic handling of unnamed columns. If column names are absent, the model sees only values and positions, so column order or learned column embeddings must carry identity.

Practical guidance and best practices

Design choices

Training tips

• Watch reconstruction and KL separately. A KL near zero signals collapse.
• Use early stopping on validation ELBO.
• Try free bits if some latent dimensions go unused.
• Use batch norm or layer norm sparingly, because they can interact with the sampling noise.
• Keep a fixed seed and log the settings.

Evaluation checklist

1. Compare column marginals of generated rows to real ones.
2. Compare correlation matrices.
3. Train a model on generated rows and test it on real rows.
4. Check the share of generated rows that closely copy training rows.
5. For anomaly scoring, evaluate AUROC or precision at k on labelled data.
6. For imputation, hide known cells and compare the filled values with the truth.

Choosing between VAE variants

• Plain VAE: a good first baseline.
• Beta-VAE: when you need more structured latents.
• Conditional VAE: when you need control over labels or segments.
• Importance weighted VAE: when the likelihood bound matters.
• Hierarchical VAE: for complex data where one latent layer is not enough.
• VQ-VAE: when you need discrete tokens.

Variants in more detail

Conditional VAE

• Both encoder and decoder receive a condition c, such as a class label or a segment code.
• The latent then only needs to explain what c does not.
• For tables, c can be a sector, a region or a customer type.
• At generation time you fix c and sample z to get rows for that group.

Beta-VAE

• Multiplies the KL term by beta.
• Beta above one pushes toward factorised, more interpretable latents, at the cost of reconstruction.
• Beta below one favours sharp reconstructions and weaker sampling from the prior.
• For tables, beta below one is often better when fidelity matters more than structure.

Importance weighted autoencoder

• Draws several samples of z per input and combines their weights.
• Gives a tighter bound than the plain ELBO.
• Costs more compute per step.
• Helpful when you use the likelihood for scoring, such as anomaly detection.

Hierarchical and ladder VAEs

• Stack several levels of latents, from coarse to fine.
• Higher levels capture global structure, lower levels capture local detail.
• Useful for complex data. For most tables a single level suffices.

Mixture priors

• Replace the single Gaussian prior with a mixture of Gaussians or a learned prior.
• Helps when the data has clear clusters, such as customer segments.
• Reduces the prior hole problem.

Handling skewed and multimodal numeric columns

Real table columns are rarely Gaussian. Income, counts and durations are skewed. Some columns have several peaks. Options:

1. Quantile transform to a normal shape. Simple and robust.
2. Log transform for positive skewed data, with care for zeros.
3. Mode-specific normalisation. Fit a mixture of Gaussians to the column, pick the mode for each value, and encode the value relative to that mode plus a one-hot mode indicator.
4. Binning to turn a number into a category when exact value does not matter.

Whichever you choose, store the inverse transform with the model so that generated rows return to the original units.

Many-level categorical columns

• One-hot columns grow wide when a category has thousands of levels.
• Use an embedding layer for the input and a softmax over the top levels for the output.
• Group rare levels into an "other" level.
• Treat high-cardinality identifiers as noise and remove them unless they carry meaning.

A short evaluation recipe

Latent space diagnostics

• Active units: count latent dimensions whose mean varies across inputs. Dimensions with no variation are unused.
• Aggregate posterior check: compare the spread of encoded means with the prior. A large mismatch means sampling from the prior will be poor.
• Traversals: vary one latent dimension while fixing the others, decode, and see what changes in the row.
• Nearest neighbours in latent space: check that close codes belong to similar rows.

Example configuration

model:
  type: vae
  latent_dim: 16
  hidden: [256, 256]
  numeric_head: gaussian
  categorical_head: softmax
train:
  epochs: 150
  batch_size: 256
  lr: 0.001
  beta_schedule:
    start: 0.0
    end: 1.0
    warmup_epochs: 40
data:
  source: example.internal/tables/customers.parquet
  numeric_transform: quantile
  mask_missing: true

Common mistakes

• Ignoring the KL term and then wondering why sampling from the prior gives nonsense.
• Letting the KL dominate so the decoder outputs the average row.
• Using a plain mean squared error loss on categorical data. Use cross entropy.
• Feeding raw skewed values and getting poor tails.
• Training on rows with leaked test information, for example by fitting transforms on all data.
• Trusting reconstruction error blindly for anomaly detection when the decoder is very powerful.
• Reading latent dimensions as meaningful without tests.
• Using a latent size that is too large, which invites unused dimensions and collapse.
• Forgetting the mask and teaching the model that missing values are zeros.
• Comparing against a weak baseline. Always include a tuned gradient boosting model or a plain autoencoder as a reference.

Frequently asked questions

What is a VAE in simple terms?

A VAE is a neural network that squeezes data into a small hidden code and rebuilds it, while forcing the codes to follow a simple known distribution. Because of that, you can draw new random codes and decode them into new, realistic samples.

How is a VAE different from an autoencoder?

A plain autoencoder maps each input to one point and is trained only to reconstruct. A VAE maps each input to a distribution and adds a penalty that keeps those distributions near a prior, which gives a smooth space you can sample from.

What is the ELBO?

The evidence lower bound is the training objective of a VAE. It equals the expected reconstruction quality minus the KL divergence between the encoder's distribution and the prior, and maximising it pushes up a lower bound on the data likelihood.

Why is the reparameterization trick needed?

Sampling is random, and gradients cannot flow through a random draw directly. Writing the sample as a mean plus a scaled standard noise term moves the randomness outside the network, so gradients can reach the encoder weights.

What is posterior collapse?

It happens when the decoder learns to ignore the latent code and the encoder outputs the prior for every input. The KL term goes to zero and the latent carries no information. KL annealing, free bits and weaker decoders can help.

Can a VAE handle tabular data?

Yes. With suitable encodings for numbers and categories, per-column output heads and masks for missing cells, a VAE works well as a baseline for synthetic data, imputation and anomaly scoring on tables.

Is a VAE good for anomaly detection?

It can be. Rows with high reconstruction error or low ELBO are candidates for anomalies. Results depend on the decoder strength and on the data, so validate with labelled examples and compare against simple detectors.

How does a VAE compare with a diffusion model?

A VAE samples in one pass and trains very stably, but its outputs are often smoother. A DDPM needs many steps to sample but usually gives sharper, more diverse samples. Many systems combine them by running diffusion in a VAE latent space.

What size should the latent space be?

There is no fixed rule. For tables, a few to a few dozen dimensions is typical. Start small, check reconstruction, and look for unused dimensions, which show a KL near zero.

Is a VAE a tabular foundation model?

No. It is a general generative method trained on one data set at a time. It can supply embeddings, synthetic rows, imputations and anomaly scores around such models, but a tabular foundation model is pre-trained across many tables and predicts through in-context learning.

Glossary

• Amortised inference: using one network to infer latent variables for any input.
• Approximate posterior: the encoder's distribution that stands in for the true posterior.
• Beta-VAE: a VAE with a weight on the KL term.
• Conditional VAE: a VAE that also receives a label or other condition.
• Decoder: the network that maps latent codes back to data.
• Disentanglement: separate latent dimensions matching separate factors.
• ELBO: evidence lower bound, the VAE objective.
• Encoder: the network that maps data to latent distributions.
• Free bits: a trick that guarantees a minimum information content per latent dimension.
• KL divergence: a measure of the difference between two distributions.
• Latent variable: a hidden variable that explains the observed data.
• Posterior collapse: a failure where the latent is ignored.
• Prior: the assumed distribution of the latent variable.
• Reparameterization trick: writing a random sample as a deterministic function of parameters and noise.
• Reconstruction error: the difference between input and output.

Summary

A VAE learns a smooth, probabilistic latent space by balancing two goals: rebuild the input well, and keep the latent distributions close to a simple prior. The reparameterization trick makes training possible with ordinary gradient methods. The result is a fast generator, a useful encoder and a source of anomaly scores. For tables, VAEs handle mixed types through per-column heads, treat missing cells with masks, and serve as practical baselines for synthetic rows, imputation and anomaly detection. They are less sharp than diffusion models and give only a bound on the likelihood, but their simplicity and speed keep them useful. They support, and do 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: