Variational Autoencoders

Why VAEs exist (motivation)

In many problems we want a model that can generate realistic data and also explain data in terms of hidden factors. Think:

• •Images explained by pose, lighting, identity
• •Audio explained by speaker, phoneme content
• •Text explained by topic, style

A standard (deterministic) autoencoder learns an encoder f(x) → z and decoder g(z) → x̂, but it does not define a probability distribution over data. You can reconstruct, but “sampling” z and decoding often produces arbitrary garbage because the latent space has no probabilistic structure.

A VAE fixes this by making the model explicitly probabilistic. It’s an instance of a latent-variable generative model:

• •Sample latent z from a prior p(z)
• •Sample data x from a likelihood/decoder p_θ(x|z)

This defines a joint distribution:

p_θ(x, z) = p_θ(x|z) p(z)

If we can learn θ well, then we can generate new data by sampling z ∼ p(z) and then x ∼ p_θ(x|z).

The core obstacle: posterior inference

Given an observed x, the Bayesian posterior over latents is

p_θ(z|x) = p_θ(x|z) p(z) / p_θ(x)

where the marginal likelihood (evidence) is

p_θ(x) = ∫ p_θ(x|z) p(z) dz

In deep models, that integral is typically intractable.

But to learn the model, we’d like to maximize log p_θ(x) over θ for the dataset. And to do inference (encode x), we want p_θ(z|x). Both are blocked by the same intractable evidence integral.

The VAE idea in one sentence

Introduce a tractable approximation q_φ(z|x) (the variational posterior / encoder) and optimize a lower bound on log p_θ(x) that is differentiable and scalable.

What makes it an “autoencoder”?

The VAE has two neural networks:

• •Encoder q_φ(z|x): maps x to parameters of a distribution over z
• •Decoder p_θ(x|z): maps z to parameters of a distribution over x

Unlike a deterministic autoencoder, the encoder outputs a distribution (often Gaussian) and the decoder defines a likelihood (often Gaussian for real-valued data, Bernoulli for binary pixels, categorical for discrete tokens, etc.).

Typical choice of distributions

A common (and very useful) baseline is:

• •Prior: p(z) = 𝒩(0, I)
• •Encoder: q_φ(z|x) = 𝒩(μ_φ(x), diag(σ²_φ(x)))
• •Decoder:
• •For real-valued x: p_θ(x|z) = 𝒩(μ_θ(z), σ² I) (often fixed σ)
• •For binary pixels: p_θ(x|z) = Bernoulli(π_θ(z))

This is not required, but it’s a common starting point because (1) sampling is easy, (2) KL terms often have closed form, and (3) reparameterization is straightforward.

Mental model

Think of training a VAE as simultaneously:

1. 1)Learning a generator that can map simple noise z into data space.
2. 2)Learning an inference network that can map data x back to a distribution over plausible z.
3. 3)Ensuring these two agree via a variational objective.

Why we need a bound at all

The quantity we would like to maximize for each datapoint x is log p_θ(x). But:

log p_θ(x) = log ∫ p_θ(x|z) p(z) dz

The log of an integral of a neural-network-defined density is generally not tractable.

Variational inference gives a workaround: introduce a distribution q_φ(z|x) that we can sample from and evaluate.

Deriving the ELBO (showing the work)

Start with the log evidence and multiply inside by q_φ(z|x) / q_φ(z|x):

log p_θ(x)

= log ∫ p_θ(x, z) dz

= log ∫ q_φ(z|x) · [p_θ(x, z) / q_φ(z|x)] dz

Now apply Jensen’s inequality to log 𝔼[·] (log is concave):

log ∫ q_φ(z|x) · [p_θ(x, z) / q_φ(z|x)] dz

= log 𝔼_{q_φ(z|x)} [ p_θ(x, z) / q_φ(z|x) ]

≥ 𝔼_{q_φ(z|x)} [ log p_θ(x, z) − log q_φ(z|x) ]

Define the ELBO:

ELBO(θ, φ; x) = 𝔼_{q_φ(z|x)} [ log p_θ(x, z) − log q_φ(z|x) ]

So we have the bound:

log p_θ(x) ≥ ELBO(θ, φ; x)

Unpacking the ELBO into “reconstruction − regularization”

Use p_θ(x, z) = p_θ(x|z) p(z):

ELBO

= 𝔼_{q_φ(z|x)}[ log p_θ(x|z) + log p(z) − log q_φ(z|x) ]

Group the last two terms as a KL divergence:

KL(q_φ(z|x) ‖ p(z))

= 𝔼_{q_φ(z|x)}[ log q_φ(z|x) − log p(z) ]

So:

ELBO

= 𝔼_{q_φ(z|x)}[ log p_θ(x|z) ] − KL(q_φ(z|x) ‖ p(z))

This is the form you implement.

Term 1: expected log-likelihood (reconstruction)

𝔼_{q_φ(z|x)}[ log p_θ(x|z) ]

• •Encourages z sampled from the encoder to decode into something that assigns high probability to the observed x.
• •With Gaussian likelihood and fixed variance, this becomes (up to constants) a negative squared error.
• •With Bernoulli likelihood, it becomes cross-entropy.

Term 2: KL to the prior (regularization)

KL(q_φ(z|x) ‖ p(z))

• •Encourages the encoded distribution to stay close to the prior.
• •This makes sampling from p(z) produce meaningful decodes.
• •Prevents the encoder from using arbitrarily “spiky” posteriors just to reconstruct perfectly.

The tightness of the bound

A key identity connects ELBO and the true posterior:

log p_θ(x) = ELBO(θ, φ; x) + KL(q_φ(z|x) ‖ p_θ(z|x))

Derivation sketch (showing the work):

KL(q ‖ p_θ(z|x))

= 𝔼_q[ log q(z|x) − log p_θ(z|x) ]

= 𝔼_q[ log q(z|x) − log (p_θ(x, z) / p_θ(x)) ]

= 𝔼_q[ log q(z|x) − log p_θ(x, z) + log p_θ(x) ]

= log p_θ(x) − 𝔼_q[ log p_θ(x, z) − log q(z|x) ]

= log p_θ(x) − ELBO

Rearrange:

log p_θ(x) = ELBO + KL(q ‖ p_θ(z|x))

Because KL ≥ 0, ELBO is a lower bound. It becomes tight when q_φ(z|x) matches the true posterior.

Dataset objective

For a dataset {xᵢ}ᵢ₌₁ᴺ, maximize:

∑ᵢ ELBO(θ, φ; xᵢ)

This trains:

• •θ to make the decoder a good likelihood model
• •φ to make the encoder approximate the posterior under the current decoder

A practical view: what gradients do we need?

We need gradients of

𝔼_{q_φ(z|x)}[ log p_θ(x|z) ]

with respect to both θ and φ, plus gradients of the KL term.

• •Gradients w.r.t. θ are usually straightforward: sample z and backprop through the decoder.
• •Gradients w.r.t. φ are subtle because z is sampled from q_φ, so the sampling operation depends on φ.

That’s exactly why the reparameterization trick matters.

Why reparameterization is needed

Suppose we approximate the expectation with Monte Carlo:

𝔼_{q_φ(z|x)}[ f(z) ] ≈ (1/L) ∑_{ℓ=1}^L f(z^{(ℓ)}), where z^{(ℓ)} ∼ q_φ(z|x)

If z^{(ℓ)} is produced by a sampling step that depends on φ, naive backprop gets stuck: the computational graph has a stochastic node.

One option is the score-function (REINFORCE) estimator:

∇_φ 𝔼_{q_φ}[f(z)] = 𝔼_{q_φ}[ f(z) ∇_φ log q_φ(z) ]

It’s unbiased but typically high-variance.

Reparameterization gives a lower-variance, pathwise gradient by rewriting the random variable as a deterministic function of φ and external noise.

The key idea

If you can write

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

where p(ε) does not depend on φ, then

𝔼_{q_φ(z|x)}[ f(z) ] = 𝔼_{p(ε)}[ f(g_φ(ε, x)) ]

Now the randomness is in ε, not in the parameters. Gradients can flow through g_φ.

Gaussian case (most common)

Let

q_φ(z|x) = 𝒩(μ_φ(x), diag(σ²_φ(x)))

Sample ε ∼ 𝒩(0, I) and define:

z = μ_φ(x) + σ_φ(x) ⊙ ε

Here ⊙ is elementwise multiplication.

This produces z distributed exactly as q_φ(z|x). And μ_φ, σ_φ are outputs of a neural net.

Backprop through the expectation

Consider the reconstruction term for a single x:

ℒ_rec(θ, φ; x) = 𝔼_{q_φ(z|x)}[ log p_θ(x|z) ]

Using reparameterization:

ℒ_rec = 𝔼_{ε∼𝒩(0,I)}[ log p_θ(x| μ_φ(x) + σ_φ(x) ⊙ ε ) ]

Approximate with L samples:

ℒ_rec ≈ (1/L) ∑_{ℓ=1}^L log p_θ(x| μ_φ(x) + σ_φ(x) ⊙ ε^{(ℓ)})

Now ∇_φ is just ordinary backprop through μ_φ and σ_φ.

The KL term and closed form

With a standard normal prior p(z) = 𝒩(0, I) and diagonal Gaussian q_φ, the KL has a closed form.

Let q = 𝒩(μ, diag(σ²)) and p = 𝒩(0, I). Then:

KL(q ‖ p) = (1/2) ∑ⱼ ( μⱼ² + σⱼ² − log σⱼ² − 1 )

This is extremely useful: it’s exact, differentiable, and cheap.

Show the structure (intuition)

Each latent dimension j pays a penalty:

• •μⱼ²: pushes means toward 0
• •σⱼ²: pushes variances toward 1 (too large is penalized)
• •−log σⱼ²: penalizes tiny variances (too confident)

So the encoder is encouraged to produce a distribution “not too far” from 𝒩(0,1).

Practical parameterization: log-variance

To keep σ positive, we usually output log σ² (call it s) and compute:

σ² = exp(s), σ = exp(0.5 s)

This avoids invalid (negative) variances and tends to be numerically stable.

Summary table: what you compute in a basic VAE

A note on discrete latents

Reparameterization is straightforward for continuous distributions like Gaussians. For discrete latents, you need alternatives (Gumbel-Softmax / Concrete distributions, score-function estimators, or other variational relaxations). Many VAE lessons stop at Gaussians because they cover the most common and useful case.

Generation

After training, generation is simple:

1. 1)Sample z ∼ p(z) = 𝒩(0, I)
2. 2)Sample x ∼ p_θ(x|z) or take the decoder mean as a “typical” sample

Because the KL term kept q_φ(z|x) near p(z), the decoder has been trained on z values that look like prior samples.

Representation learning

The latent z can be used as a learned feature representation. Common uses:

• •Clustering in latent space
• •Interpolation: decode points along (1−t)z₁ + tz₂
• •Downstream supervised tasks using z as input

Caution: VAEs trade off reconstruction fidelity vs. latent regularity. If the KL term dominates, representations can become less informative.

Conditional VAEs (cVAEs)

If you want generation conditioned on labels or attributes y:

• •Prior p(z|y)
• •Decoder p_θ(x|z, y)
• •Encoder q_φ(z|x, y)

This lets you generate samples with a chosen condition.

Anomaly detection

A common heuristic: points with low ELBO (or high reconstruction error) are considered anomalous. This works best when the model class fits normal data well.

The “posterior collapse” problem (important failure mode)

In powerful decoders (e.g., autoregressive text decoders), the model may learn to ignore z entirely:

• •q_φ(z|x) ≈ p(z) (KL goes to 0)
• •Decoder models p_θ(x) well without needing z

Then z carries little information about x.

Common mitigations

β-VAE: a small change with big effects

Objective:

ELBO_β = 𝔼_q[log p_θ(x|z)] − β KL(q ‖ p)

• •β > 1: stronger pressure to match the prior → often more “disentangled” latents but blurrier samples
• •β < 1: weaker pressure → better reconstructions but latent space may be less smooth for sampling

Importance-weighted autoencoders (IWAE)

ELBO uses one sample (or few) and is a lower bound. IWAE tightens the bound using multiple samples:

log p_θ(x) ≥ 𝔼_{z₁:K ∼ q}[ log ( (1/K) ∑ₖ p_θ(x, zₖ) / q(zₖ|x) ) ]

This can improve generative modeling but changes optimization dynamics.

Connection to diffusion models (why this node unlocks them)

Diffusion models also involve:

• •Latent/noise variables (a sequence of noisy states)
• •Learning to reverse a corruption process
• •Using tractable training objectives that avoid directly optimizing log p(x) in closed form

Conceptually, VAEs train a generator with latent variables via a variational bound; diffusion trains a generator via denoising/score objectives. Understanding:

• •latent-variable modeling
• •KL terms and approximate inference
• •reparameterization and sampling-based gradients

makes diffusion objectives feel much less mysterious.