Variational autoencoder

Overview

A VAE has an encoder that maps an input to the mean and spread of a Gaussian over a latent code, a decoder that maps a latent back to the input space, and a loss that is the reconstruction error plus a KL divergence between the encoder's Gaussian and a prior, usually a standard normal. The loss is a bound on the data likelihood 2103.04922 §3. The latent is drawn with the reparameterisation trick, z = mu + sigma * eps, so gradients reach the encoder 2103.04922 §3 2103.14274 §3.1. At generation time the encoder is usually dropped, and a latent is drawn from the prior and decoded.

The cards on this page share that core and change one or more parts:

• the condition: a label, a class embedding, a previous pose or a point cloud is given to both encoder and decoder;
• the latent type: continuous Gaussian, or discrete codes from a codebook;
• the regulariser: the KL term is weighted differently, or replaced by an adversarial or moment-matching term;
• what sits on top: a policy, a planner, a transformer prior, or a diffusion model running in the latent space.

Blurry samples are a known property of the plain Gaussian form. The survey links them both to the likelihood objective and to the restrictiveness of the Gaussian posterior, and reports evidence for the second cause 2103.04922 §3.

Variants

Plain Gaussian VAE as a compressor

The VAE shrinks observations so that a small learner can work on them. World Models trains a convolutional VAE on 64x64 frames, with a 32-dimensional latent for car racing and 64 for VizDoom 1803.10122 §3.2 1809.01999 app. A.2. A single linear controller on the latent and the recurrent state reaches 906 ± 21 on CarRacing-v0 1803.10122 tab. 2. The VizDoom controller is trained inside the learned model and then run on the real game, with 1092 ± 556 at temperature 1.15 1809.01999 tab. 2. The two cards report the same numbers for these runs. The first is World Models and the second is Recurrent World Models.

For a locomotion policy, a CNN-VAE maps a 20x40 terrain height map to a 16-value latent 2512.12993 §II.B. The authors report that the latent size of 16 learns faster than 64 or the full map, and that policies stay above 80% success on the tested terrains in simulation 2512.12993 §IV.A 2512.12993 §IV.E. The card is Learning Terrain Aware Bipedal Locomotion.

Two VAEs with 64-dimensional latents can also map one image domain onto another through a mixture density network, but the domain-transfer result is only shown as examples 1812.01718 §3.2. See Deep Learning for Classical Japanese Literature.

Conditional VAE

The condition is concatenated to the encoder input and to the latent before decoding.

• Zero-shot features. A conditional VAE is trained on image features given a class embedding. For each unseen class it decodes 300 latent draws into pseudo-features, and an SVM is fitted on them 1709.00663 alg. 1 1709.00663 §5.3. It reaches 71.4 on AwA1 in the disjoint setting, above 65.6 for SJE 1709.00663 tab. 2, and 47.2 on generalized AwA1 against 27.5 for ALE 1709.00663 tab. 3. The model generates features, not images. See A Generative Model For Zero Shot Learning.
• Grasps. A grasp encoder and decoder are conditioned on a point-cloud latent, and the KL weight is annealed during training 2312.11243 §3.1 2312.11243 §3.3. See GraspLDM for its use with a latent prior.
• Audio goals. A conditional VAE maps a 512-d audio embedding to a 512-d visual goal in the space of a fixed policy 2412.00949 §3.4. The paper does not state the loss; the released code uses KL plus reconstruction 2412.00949 §3.4. See STEVE-Audio.
• Motion. The Motion VAE predicts the next pose from the current one. It uses a 32-dimensional latent, and a mixture-of-experts decoder with the latent fed to every layer 2103.14274 §3.1 2103.14274 §3.2. It is trained with scheduled sampling over 8-frame roll-outs 2103.14274 §3.4. On random walks it reports foot skate of 0.067 cm/frame against 0.10 for the motion capture data 2103.14274 tab. 1. See Character Controllers Using Motion VAEs.
• Design. In a 2-D study, a conditional VAE scores lower conditioning reconstruction error than a conditional GAN (0.007 against 0.026) 2302.02913 tab. 6. See Beyond Statistical Similarity.

VAE with a diffusion prior in latent space

A VAE compresses the data, and a diffusion model learns the distribution of the latent. The diffusion model can then be sampled with fewer steps.

• Motion. A transformer VAE encodes a motion clip into one 256-d latent token, and a transformer denoiser runs diffusion on that token with classifier-free guidance 2212.04048 §3.1 2212.04048 §5. It reports FID 0.473 on HumanML3D against 0.544 for MDM 2212.04048 tab. 1. For 2048 clips with 50 DDIM steps, it takes 10.24 s against 225.28 s for MDM 2212.04048 tab. 10. See Motion Diffusion in Latent Space.
• Grasps. In the same design, the grasp latent is sampled from the VAE prior and refined by a DDPM before decoding 2312.11243 §3.2. On 100 grasps for one scene, VAE-only sampling takes 0.02 s with median success 0.660, DDIM with 100 steps takes 0.75 s with 0.756, and DDPM with 1000 steps takes 7.39 s with 0.792 2312.11243 tab. 1.

Discrete VAE and VQ-VAE

The latent is a grid of indices into a learned codebook. The encoder output is mapped to the nearest code, or to a sample from a categorical distribution.

• dVAE with Gumbel-softmax. A 256x256 image becomes a 32x32 grid of tokens from 8192 codes. The gradient passes through a relaxed categorical sample, and the relaxation temperature is annealed 2102.12092 §2.1. A transformer then models caption tokens and image tokens as one sequence 2102.12092 §2.2. See Zero-Shot Text-to-Image Generation.
• VQ-VAE with an EMA codebook. VideoGPT uses 3D convolutions and axial attention, a codebook updated by exponential moving average with decay 0.99, and a straight-through gradient 2104.10157 §2.1. Its latent of 8x32x32 has the lowest FVD in the sweep, 103.6 2104.10157 tab. 5. The headline FVD of 103.3 on BAIR is in 2104.10157 tab. 1. See VideoGPT.

Adversarially regularised autoencoder

CAAE drops the KL term. A discriminator on the codes pushes them toward a uniform prior, and a second discriminator judges (image, age) pairs 1702.08423 §4.3 1702.08423 §4.4. It needs no paired images of one person at different ages, and a user study prefers its faces over prior work at 52.77% of votes 1702.08423 §1 1702.08423 §5.3. The card is listed under this concept, but its objective has no KL term. See Age Progression/Regression by Conditional Adversarial Autoencoder.

Disentangling VAEs

beta-VAE, FactorVAE, beta-TCVAE and DIP-VAE add penalties on the aggregate posterior to separate factors of variation 1905.12506 §2. Across 360 encoders trained on two toy datasets, FactorVAE and beta-VAE scores track the accuracy of a downstream reasoning model in the few-sample regime. After 100K steps, reconstruction error tracks accuracy better 1905.12506 §4.2.2. See Are Disentangled Representations Helpful for Abstract Visual Reasoning?.

VAE as a pretrained component at test time

A pretrained autoencoder need not be trained in the same pipeline. A depth method averages the latents of several exposures of one scene inside the Stable Diffusion VAE, decodes the average, and feeds it to a depth model 2408.06083 §3.3. The fusion step adds almost nothing on the transparent-and-mirror region: 0.859 to 0.860 2408.06083 tab. 3. The main gain comes from the training changes. See Towards Robust Monocular Depth Estimation.

Latent-variable guidance inside iterative reasoning

A recursive reasoner adds a learned Gaussian guidance term to each latent step. It is trained with an ELBO over the trajectory, using a posterior that sees the target and a prior that does not, and it uses truncated gradients 2605.19376 §2.2. At test time, N trajectories are drawn from the prior and the most frequent answer is returned 2605.19376 §2.3. Removing the guidance drops Sudoku accuracy to the deterministic baseline, and guidance without noise fails completely 2605.19376 tab. 3. The card reports 96.1% on Sudoku-Extreme at N = 20 2605.19376 tab. 8. See Generative Recursive Reasoning.

EEG-to-media survey

The EEG survey describes VAE decoders as one of several generators in EEG-to-image work, with no numeric comparison 2502.12048 §4.3. Its best text number, a BLEU-1 of 52.38 on ZuCo, comes from a different family 2502.12048 tab. 4. See A Survey on Bridging EEG Signals and Generative AI.

When to use which

The choice depends on what the downstream step needs and what the budget allows.

Implementation sketch

The core below is the common part. Each card adds its own condition, codebook or prior. Cards disagree on the details noted under the sketch.

# Core VAE objective shared by the cards  [arxiv:2103.04922 §3] [eq. 13]
# Input: x (B, D), optional condition c (B, C); latent size J
# Output: scalar loss to minimise

def vae_loss(encoder, decoder, x, c, beta):
    # Encoder gives mean and log-variance per latent unit  [arxiv:2103.04922 §3]
    mu, logvar = encoder(x, c)                        # (B, J) each
    # Reparameterised sample; gradients reach the encoder  [arxiv:2103.14274 §3.1]
    eps = normal_like(mu)
    z = mu + exp(0.5 * logvar) * eps                  # (B, J)
    # Decoder reads the sample and the condition  [arxiv:1709.00663 §3.2]
    x_hat = decoder(z, c)
    # Reconstruction term; the cards use squared error, L1 or a likelihood
    recon = sum((x - x_hat) ** 2)
    # Closed-form KL to N(0, I)  [arxiv:2103.04922 eq. 13]
    kl = -0.5 * sum(1 + logvar - mu**2 - exp(logvar), axis=1)
    # beta is a free weight; cards use values from 1e-4 to 1.0  [see below]
    return mean(recon + beta * kl)

# Generation: draw from the prior and decode  [arxiv:1803.10122 §3.2]
# [arxiv:1709.00663 alg. 1]
def generate(decoder, condition, n):
    # Standard normal draws, one per sample  [arxiv:1709.00663 alg. 1]
    z = normal(size=(n, J))
    # Decode with the condition appended, if there is one  [arxiv:1709.00663 §3.2]
    return decoder(z, condition)

Where the cards differ:

• KL weight. The cards use very different values. MLD's released code uses 1e-4 2212.04048 §3.1. GraspLDM anneals from 1e-7 to 0.1 over the first half of training 2312.11243 §3.3. The Motion VAE paper uses 0.2, and the released code uses 1.0 2103.14274 §3.4. GRAM uses between 0.015 and 0.5 depending on the task 2605.19376 app. B.2. Tune it and state the value used.
• Sum or mean. GraspLDM's code averages the reconstruction over pose dimensions and sums the KL over latent dimensions. That changes the effective KL weight, and the card warns that summing the reconstruction would weight the KL six times less 2312.11243 §3.3. The audio-prior code sums the KL over the batch and averages the MSE, so its weighting depends on batch size 2412.00949 §3.4.
• Std or variance. The terrain paper calls sigma a variance but uses it as a multiplier, so whether the encoder outputs a std or a log-variance is not settled 2512.12993 §II.B. Choose one and check the KL.
• Inference sample. Some cards sample at test time, some do not state it, and the audio prior draws fresh noise on every call 2412.00949 Pitfalls. Test the mean and the sample, and report which one you used.
• Discrete gradients. The dVAE relaxes the categorical sample and anneals the temperature with a cosine schedule. The card says linear annealing diverged 2102.12092 app. A.2. VQ-VAE uses a straight-through estimator instead 2103.04922 §3.2.
• Bound validity. The dVAE paper says its bound is exact only at beta = 1, and it trains with beta = 6.6 2102.12092 §2 2102.12092 app. A.2. The objective is then no longer a strict bound.

Contradictions

• Which family gives better samples. The survey's conclusion says GANs have led on sample quality, and its table puts a convolutional VAE at FID 106.37 on CIFAR-10 against 14.73 for BigGAN 2103.04922 §9 2103.04922 tab. 1. The design-metrics paper runs a 2-D six-mode test with one VAE and one GAN. The VAE scores higher on precision-recall F1 (0.475 against 0.386) and AUC, while the GAN scores better on nearest-datapoint distance and MMD 2302.02913 tab. 2. The paper runs one instance per model and cautions that the result may not carry over 2302.02913 §4.5. The two sources use different data and metrics, so neither settles the question.
• Size of the KL weight. Papers that use the same family choose weights that differ by orders of magnitude (see the sketch). No card gives a rule that carries over.
• Best latent size. VideoGPT's best BAIR latent is the 8x32x32 grid, and its largest latent gives worse samples, which it attributes to limited compute for the prior 2104.10157 tab. 5 2104.10157 §4.7. The terrain card picks 16 for policy learning while its loss curves favour 32 to 64 2512.12993 Pitfalls. These are optima for different downstream targets, not conflicting measurements of one quantity.
• Bound exactness. The survey writes the training objective as the bound without caveat 2103.04922 §3. The dVAE paper states that the bound is exact only at beta = 1 2102.12092 §2.

Papers

• World Models: VAE, recurrent model and controller; car racing 906 ± 21 and dream-trained VizDoom controller.
• Recurrent World Models Facilitate Policy Evolution: same pipeline with the VizDoom result 1092 ± 556 at temperature 1.15.
• A Generative Model For Zero Shot Learning Using Conditional Variational Autoencoders: conditional VAE on features, pseudo-features for unseen classes, SVM classifier.
• Age Progression/Regression by Conditional Adversarial Autoencoder: adversarial code regularisation in place of KL, conditioned on age.
• Deep Learning for Classical Japanese Literature: two VAEs with 64-d latents for a qualitative domain transfer.
• Are Disentangled Representations Helpful for Abstract Visual Reasoning?: 360 disentangling VAEs and their link to reasoning sample efficiency.
• Zero-Shot Text-to-Image Generation: discrete VAE with Gumbel-softmax tokens and a transformer prior.
• Deep Generative Modelling: A Comparative Review: survey of VAE, GAN, flow, energy and autoregressive models with the ELBO derivation.
• Character Controllers Using Motion VAEs: autoregressive conditional VAE for motion, with latent actions for RL.
• VideoGPT: VQ-VAE with axial attention and a GPT-style prior over video codes.
• Executing your Commands via Motion Diffusion in Latent Space: transformer VAE with a latent diffusion model for text and action to motion.
• Beyond Statistical Similarity: metrics for design generators, with a VAE vs GAN test on 2-D data.
• GraspLDM: conditional VAE for grasps with a DDPM prior over its latent.
• Towards Robust Monocular Depth Estimation in Non-Lambertian Surfaces: uses Stable Diffusion VAE latents to fuse several exposures at test time.
• STEVE-Audio: conditional VAE prior mapping audio into the goal space of a Minecraft agent.
• A Survey on Bridging EEG Signals and Generative AI: survey of EEG-to-image, text and audio generators, including VAE decoders.
• Learning Terrain Aware Bipedal Locomotion: CNN-VAE compresses terrain height maps to a 16-value latent for a biped policy.
• Generative Recursive Reasoning: latent-variable guidance with an ELBO inside a recursive reasoner.