Variational Autoencoders (VAEs)

TL;DR. A VAE is an autoencoder whose bottleneck is probabilistic instead of a single point. The encoder maps each input to a small cloud (a Gaussian) in a low-dimensional latent space; the decoder reconstructs the input from a sample of that cloud. You train it by maximizing one objective — the ELBO — which simultaneously says "reconstruct the data well" and "keep the latent clouds shaped like a simple prior." Because sampling is non-differentiable, the reparameterization trick routes the randomness around the gradient path so you can train with ordinary backprop. The payoff: a smooth, samplable latent space you can generate from, interpolate through, and compress into. VAEs are also the silent workhorse under modern image generators — the "VAE" in Stable Diffusion is a VAE-style compressor, and VQ-VAE → latent diffusion is the lineage that made high-resolution generation affordable.

What it is (intuition first)

Start from a plain autoencoder (see the autoencoders note in this library — if absent, read the [[embeddings]] note, which covers the same compress-to-a-vector idea). A plain autoencoder is two neural nets glued at the waist:

• an encoder squeezes an input x (an image, say) into a small vector z (the "code" / "latent"),
• a decoder tries to rebuild x from z.

Train it to minimize reconstruction error and the middle vector z becomes a compressed summary. That is great for compression and denoising, but it has a problem if you want to generate new data: the latent space is full of holes. Pick a random z and decode it, and you usually get garbage, because the encoder only ever populated a few scattered islands and never agreed on how the space should be laid out.

The VAE fixes this with one conceptual move: make the bottleneck a distribution, not a point. Instead of emitting a single z, the encoder emits the parameters of a little Gaussian — a mean μ(x) and a spread σ(x). The actual z is sampled from that Gaussian. Then a second pressure is added to the loss: every input's little Gaussian is pulled toward a shared, simple prior (a unit Gaussian centered at the origin).

Two things fall out of that:

1. The space fills in. Because each input now claims a small region rather than a point, and all regions are pushed toward the same prior, neighboring codes start meaning similar things. The space becomes smooth and continuous — interpolate between two faces' codes and you get plausible in-between faces.
2. You can generate. To make a brand-new sample, you don't need an input at all: draw z straight from the prior N(0, I) and run the decoder. Because training forced the encoder's outputs to look like the prior, samples from the prior decode into plausible data.

So a VAE is "an autoencoder you can sample from," and the price of that superpower is a slightly blurrier reconstruction and a more involved loss. The rest of this note is about why that loss has the form it does.

Why it matters

• It is a principled generative model. A VAE is not a heuristic — it is a tractable approximation to maximum-likelihood learning of a latent-variable model p(x) = ∫ p(x|z) p(z) dz. The ELBO it maximizes is a genuine lower bound on log p(x). That grounding is rare and valuable.
• It gives you a *usable* latent space. Unlike a GAN (which has a latent space but no encoder, so you can't easily map a real image into it), a VAE comes with an encoder for free. That makes it the natural choice when you need both directions: data→code (compression, anomaly detection, representation learning) and code→data (generation).
• It is the compression layer under modern diffusion. Latent Diffusion Models / Stable Diffusion don't diffuse in pixel space — they diffuse in the latent space of a VAE-style autoencoder. The VAE makes a 512×512×3 image into a ~64×64×4 latent, an ~48× spatial compression, after which diffusion is cheap. See [[diffusion-models]].
• VQ-VAE seeded the "everything is tokens" era. By making the latent discrete, VQ-VAE turned images, audio, and video into sequences of integer codes — which is exactly what an autoregressive transformer wants to predict. That idea runs straight into [[multimodal-llms]] and modern image/video tokenizers.

How it works (real mechanics)

Setup and notation

We posit a generative story: a latent z is drawn from a prior p(z), and then data is drawn from a decoder p_θ(x|z). We want to fit θ. The trouble is the true posterior p_θ(z|x) — "given this image, which latent produced it?" — is intractable (the integral over z has no closed form). So we introduce an approximate posterior q_φ(z|x), a neural network with its own parameters φ. In Kingma & Welling's vocabulary:

• q_φ(z|x) is the recognition model / probabilistic encoder,
• p_θ(x|z) is the probabilistic decoder,
• φ (encoder weights) and θ (decoder weights) are learned jointly.

The ELBO (the whole objective in one line)

The marginal log-likelihood of one datapoint decomposes exactly (eq. 1 in the paper):

log p_θ(x) = D_KL( q_φ(z|x) ‖ p_θ(z|x) )  +  L(θ, φ; x)

The first term is the KL divergence between our approximate posterior and the true posterior — we can't compute it, but we know it's ≥ 0. Therefore L is a lower bound on log p_θ(x) — the Evidence Lower BOund (ELBO). Maximizing L does double duty: it pushes up the data likelihood and (because the total is fixed for given θ) it squeezes the approximate posterior toward the true one. The ELBO has two equivalent forms; the useful one (eq. 3) is:

L(θ, φ; x) =  E_{q_φ(z|x)} [ log p_θ(x | z) ]   −   D_KL( q_φ(z|x) ‖ p(z) )
              \_______________________________/     \_______________________/
                    reconstruction term                  regularizer

Read it in plain English:

• Reconstruction term — sample a z from the encoder, decode it, and reward high probability of reproducing the original x. (For Gaussian decoders this is essentially negative MSE; for Bernoulli pixels it's binary cross-entropy.)
• KL term — punish the encoder for letting its per-input Gaussian q_φ(z|x) drift away from the prior p(z) = N(0, I). This is the regularizer that makes the space smooth and samplable.

A VAE minimizes −L. That's it — one loss, two terms, fighting each other.

The reparameterization trick (why this trains at all)

There is a landmine in the reconstruction term: it's an expectation over z ∼ q_φ(z|x), and z depends on φ. To get ∂L/∂φ you'd differentiate through a sampling operation, which has no gradient. The naïve fix (the score-function / REINFORCE estimator) works but has crippling variance:

∇_φ E_{q_φ(z)}[f(z)] = E_{q_φ(z)}[ f(z) ∇_φ log q_φ(z) ]      # high variance, impractical

The reparameterization trick (eq. 4) sidesteps it. Instead of sampling z from a φ-dependent distribution, you sample a fixed noise ε from a φ-free distribution and push it through a deterministic, differentiable function g_φ:

z = g_φ(ε, x),     ε ∼ p(ε)

For the standard diagonal-Gaussian encoder this is concretely:

z = μ_φ(x) + σ_φ(x) ⊙ ε,     ε ∼ N(0, I)      # ⊙ = elementwise product

Now the randomness (ε) sits outside the gradient path, and μ and σ are ordinary differentiable network outputs. Backprop flows through μ and σ cleanly. The resulting SGVB (Stochastic Gradient Variational Bayes) estimator (eq. 5–6) is low-variance and trains with plain SGD/Adam. This single trick is what turned variational inference from an MCMC-flavored chore into a one-GPU autoencoder you train in an afternoon.

The Gaussian KL has a closed form

When both q_φ(z|x) = N(μ, σ²) and the prior p(z) = N(0, I) are Gaussian, the KL term needs no sampling at all — it integrates analytically (Appendix B of the paper). For latent dimension J:

−D_KL( q_φ(z|x) ‖ p(z) )  =  ½ · Σ_{j=1}^{J} ( 1 + log(σ_j²) − μ_j² − σ_j² )

So in practice only the reconstruction term is Monte-Carlo estimated (usually with a single sample, L=1, per datapoint), and the KL is computed in closed form. The encoder typically outputs log σ² rather than σ for numerical stability.

Minimal training loop (pseudocode)

# x: a minibatch of inputs
mu, logvar = encoder(x)                 # two heads: mean and log-variance
std        = exp(0.5 * logvar)
eps        = randn_like(std)            # noise sampled OUTSIDE the gradient path
z          = mu + std * eps             # reparameterization trick

x_hat      = decoder(z)

recon = reconstruction_loss(x_hat, x)   # BCE for {0,1} pixels, or MSE / Gaussian NLL
kl    = -0.5 * sum(1 + logvar - mu**2 - exp(logvar))   # closed-form Gaussian KL

loss  = recon + beta * kl               # beta = 1 is the vanilla ELBO
loss.backward(); opt.step()

To generate afterwards: z = randn(...) from the prior, x_new = decoder(z). No encoder needed.

Key ideas & tradeoffs

VAE vs plain Autoencoder (AE)

The AE optimizes only "rebuild it." The VAE adds "...and keep the codes organized like the prior." That organization is the entire reason a VAE can dream up new data while an AE mostly can't.

VAE vs GAN

GANs (a generator vs. a discriminator in a minimax game) and VAEs are the two classic deep generative families, with near-opposite tradeoffs:

• Sample sharpness. GANs win — they produce crisper, more photorealistic images. VAEs tend to be blurry, partly because the Gaussian-decoder likelihood and the averaging effect of the KL smear high-frequency detail.
• Training stability. VAEs win — they minimize a single well-defined loss with stable gradients. GANs are adversarial, prone to mode collapse (generator emits only a few sample types) and finicky to balance.
• Likelihood & encoder. VAEs win — they give a principled likelihood bound and a built-in encoder (data→latent). Vanilla GANs have neither; you can't easily invert a real image into a GAN's latent.
• Coverage. VAEs tend to cover the data distribution more completely (they're trained to explain all the data); GANs trade coverage for fidelity.

In practice the field largely moved past the VAE-vs-GAN dichotomy: diffusion models (see [[diffusion-models]]) now dominate image generation on quality, coverage, and training stability — and they frequently run on top of a VAE compressor. So the modern role of the VAE is less "the generator" and more "the latent space the generator lives in."

VQ-VAE: discrete latents

VQ-VAE (van den Oord et al., 2017) replaces the continuous Gaussian latent with a discrete codebook. You keep a learnable table of K embedding vectors e ∈ R^{K×D}. The encoder produces a continuous vector z_e(x); quantization snaps it to its nearest codebook entry (eq. 1–2):

k = argmin_j ‖ z_e(x) − e_j ‖₂ ,     z_q(x) = e_k        # nearest-neighbor lookup

That argmin has no gradient. VQ-VAE uses the straight-through estimator: in the forward pass use the quantized z_q(x); in the backward pass copy the decoder-input gradient straight back to the encoder as if quantization were the identity. The full loss (eq. 3) has three terms:

L = log p(x | z_q(x))  +  ‖ sg[z_e(x)] − e ‖₂²  +  β · ‖ z_e(x) − sg[e] ‖₂²
    \_______________/     \_____________________/     \__________________________/
      reconstruction          codebook loss               commitment loss

where sg[·] is the stop-gradient operator (identity forward, zero gradient backward). The middle term pulls codebook vectors toward the encoder outputs they're matched to (a dictionary-learning / k-means-style update); the last term, weighted by β (the paper uses β = 0.25), keeps the encoder from outrunning the codebook. With a uniform prior over codes the KL term reduces to a constant log K, so it drops out of training. Crucially, the discrete bottleneck sidesteps posterior collapse — the failure mode where a too-powerful autoregressive decoder learns to ignore the latents entirely.

VQ-VAE → latent diffusion (the lineage that matters)

This is the through-line to today's image generators:

1. VQ-VAE / VQGAN showed you can compress an image into a small grid of discrete (or low-dim continuous) latents and decode it back with high fidelity. The latents carry the semantics; the decoder restores pixel-level texture.
2. Latent Diffusion Models (Rombach et al., CVPR 2022 — the Stable Diffusion paper) made the decisive observation: pixel space is mostly high-frequency detail with little semantic content, so it's wasteful to run an expensive diffusion model there. Instead, train a VAE-style autoencoder once to map images into a compact perceptually-equivalent latent space, then run the diffusion process entirely in that latent space, and only decode to pixels at the very end. The reported result was a large compute reduction at training and inference time while retaining quality — which is precisely what put high-resolution text-to-image generation within reach of consumer GPUs.

So when people say "the VAE in Stable Diffusion," they mean exactly this autoencoder — the compress/decompress sandwich around the diffusion core. The VAE isn't the generator; it's the coordinate system the generator works in. For the diffusion half of this story see [[diffusion-models]]; for how compressed latents become discrete tokens a transformer can predict, see [[multimodal-llms]].

Honest caveats

• VAE samples are blurry. The classic complaint. Continuous-VAE reconstructions and samples lack high-frequency detail compared to GANs or diffusion. This is intrinsic to the Gaussian-likelihood + KL-smoothing setup, and it's the main reason vanilla VAEs aren't used as standalone photorealistic generators today.
• Posterior collapse is real. If the decoder is powerful enough (e.g. an autoregressive PixelCNN/transformer), the model can drive the KL term to ~0 and ignore the latent entirely, degenerating into "decoder-only." Mitigations include KL warm-up/annealing, free-bits, β-scheduling, weaker decoders, and — the structural fix — discrete latents (VQ-VAE).
• The ELBO is a *lower* bound, and the gap can be large. A high ELBO doesn't guarantee a tight fit to log p(x); the slack is the KL between approximate and true posterior. Tighter-bound variants exist (IWAE / importance-weighted, normalizing-flow posteriors, inverse-autoregressive flows) precisely because the simple diagonal-Gaussian q is often too crude.
• The β knob trades off two things you both want. β-VAE (upweighting the KL) buys more disentangled, prior-conforming latents at the cost of worse reconstructions; β<1 buys fidelity at the cost of a less samplable space. There is no free lunch — the reconstruction and regularizer terms genuinely pull against each other.
• "Disentanglement" claims are shaky. β-VAE-style disentanglement is sensitive to seeds, hyperparameters, and unidentifiable up to rotation/relabeling; it's not the robust, free property early papers suggested.
• VQ-VAE has its own gremlins. Codebook collapse (most entries go unused), sensitivity to codebook size and β, and the bias of the straight-through estimator are all live issues. Production tokenizers add tricks: EMA codebook updates, codebook reset, k-means init, factorized/low-dim codes (as in later VQGAN/SD tokenizers).
• Reproduction note for this article. The arXiv abstract pages do not contain equations; the equations quoted here were extracted directly from the PDFs of arXiv:1312.6114 (AEVB) and arXiv:1711.00937 (VQ-VAE) via local text extraction, and cross-checked against the LDM paper (arXiv:2112.10752). Equation numbers refer to those papers. The latent-diffusion compute/quality claims are as reported by Rombach et al.; treat headline numbers as paper-reported rather than independently re-measured.

How it connects to OpenAlice + the Academy ladder

Where VAEs sit in the OpenAlice picture. OpenAlice is text/agent-first, so VAEs aren't on the critical path of Alice's core loop. But the concepts are load-bearing across the ecosystem:

• Latent spaces are the lingua franca of representation. The encoder-to-a-distribution idea is the same family as [[embeddings]] — the difference is that a VAE adds a generative prior and a samplable geometry on top of "compress to a vector." Anywhere OpenAlice reasons over embedding spaces (retrieval, memory, [[graphrag]]), the VAE is the generative-model end of that same spectrum.
• The VAE-as-compressor pattern underlies any image/video/audio generation OpenAlice products touch. If a streaming/persona product ever generates or edits visual media, the relevant stack is "VAE/VQ-VAE tokenizer + [[diffusion-models]]" — exactly the LDM lineage above. Knowing the VAE is the cheap-compute trick that makes that affordable is the practically useful takeaway.
• VQ-VAE's discrete tokens are the bridge to "everything is a sequence." The argmin-into-a-codebook move is why images and audio can be fed to the same autoregressive transformer machinery OpenAlice already uses for text — directly relevant to [[multimodal-llms]] and to [[world-models]], where learned discrete latent dynamics are a recurring design.

Academy ladder placement. This note is a mid-tier generative-models rung. Suggested approach order:

1. Foundations first — probability, expectation, KL divergence, and the basic autoencoder (the autoencoders note when present, else [[embeddings]]). The ELBO and the Gaussian KL closed form are unintelligible without comfort here.
2. This note (VAEs) — the canonical first encounter with amortized variational inference and the reparameterization trick. The reparameterization trick in particular is a reusable idea that shows up far beyond VAEs (it's a general low-variance gradient estimator for stochastic nodes).
3. Then [[diffusion-models]] — the natural successor; read VAEs first so that "diffusion in the latent space of a pretrained autoencoder" lands as an obvious optimization rather than a mystery.
4. Sideways to [[embeddings]] and [[multimodal-llms]] — to see the encoder/latent and discrete-token ideas reused in retrieval and cross-modal modeling.

If you remember exactly one thing: a VAE is an autoencoder whose bottleneck is a distribution, trained by maximizing the ELBO, made differentiable by the reparameterization trick — and that single design is the reason we have smooth, samplable latent spaces to do generation and compression in.