The variational autoencoder: turning an autoencoder's holey latent into one continuous space you can sample and interpolate

A plain autoencoder squeezes each input to a single point in a latent code and decodes it back. That works for compression, but the latent space it learns is holey: training points sit on islands with dead gaps between them, so if you pick a random code — or a point between two examples — the decoder produces garbage. You can't sample from it, so you can't generate. A VAE fixes this with three small additions, and I built a live 2-D one where the reparameterization trick and the Gaussian KL are genuine math in the browser, not a canned animation. Here is how it works.

The encoder emits a distribution, not a point

A normal encoder ends in one linear layer → the code. A VAE encoder ends in two linear heads over the same hidden features: one produces the mean μ, the other the log-variance logσ². Predicting the log keeps σ² positive and unconstrained. Together they define a diagonal Gaussian posterior q(z|x) = N(μ, σ²) for this input — a whole little neighbourhood instead of a dot.

class Encoder(nn.Module):
    def __init__(self, in_dim, hidden, zdim):
        super().__init__()
        self.body    = nn.Sequential(nn.Linear(in_dim, hidden), nn.ReLU())
        self.fc_mu   = nn.Linear(hidden, zdim)   # mean of q(z|x)
        self.fc_logv = nn.Linear(hidden, zdim)   # LOG-variance (unconstrained)
    def forward(self, x):
        h = self.body(x)
        return self.fc_mu(h), self.fc_logv(h)

The reparameterization trick makes sampling differentiable

Now we must sample z ~ N(μ, σ²) — but drawing a random sample has no gradient with respect to μ and σ, so backprop would stop dead. The trick moves the randomness outside the parameters: sample ε ~ N(0,1), then form z = μ + σ⊙ε. Now z is a smooth, differentiable function of μ and σ, with ∂z/∂μ = 1 and ∂z/∂σ = ε, so gradients flow straight through.

def reparameterize(mu, logvar):
    std = torch.exp(0.5 * logvar)      # sigma = exp(0.5 * log sigma^2)
    eps = torch.randn_like(std)        # epsilon ~ N(0,1), the randomness
    return mu + std * eps              # z = mu + sigma (elementwise) eps

This is the one line that separates a VAE from a plain autoencoder plus noise. Sampling that code and decoding it is what forces a whole neighbourhood around μ to decode sensibly, not just the exact point.

The Gaussian KL is one closed-form line

The loss has two halves. The first is the usual reconstruction term — how far the decoded output is from the input (binary cross-entropy for pixels, MSE for continuous data). The second is what glues the space together: because both the posterior N(μ,σ²) and the prior N(0,1) are Gaussians, their KL divergence has an exact closed form, no sampling needed.

def kl_divergence(mu, logvar):
    # KL( N(mu, sigma^2) || N(0,1) ) = 0.5 * sum( mu^2 + sigma^2 - 1 - log sigma^2 )
    return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / mu.size(0)

It is zero only when every μ=0 and σ=1 — exactly the prior — and grows as the posterior drifts away. So minimizing it pulls every blob toward the origin with unit spread, merging the islands into one continuous region.

Recon + β·KL is the negative ELBO

Add the two: L = recon + β·KL. This is precisely the negative ELBO (evidence lower bound) — minimizing it maximizes a lower bound on log p(x). β=1 is the original VAE; β>1 (a β-VAE) weights the prior harder for a more continuous, disentangled space at some cost to sharpness; β=0 collapses back to a plain autoencoder. The demo lets you drag β and literally watch the encoded blobs get pulled toward the unit-circle prior, the shaded islands fusing.

The payoff: generate and interpolate

Because the KL pushed the aggregate posterior toward N(0,1), generation becomes trivial — sample z ~ N(0,1) from the prior and decode. No encoder, no input needed.

@torch.no_grad()
def generate(dec, n, zdim):
    z = torch.randn(n, zdim)   # z ~ prior N(0,1)
    return dec(z)              # decode -> new data

And interpolation works: walk a straight line between two encodings and every midpoint decodes to something sensible, morphing smoothly, because the KL term filled the space between the two blobs. In a plain autoencoder those middle codes decode to noise — that gap is the whole difference. The probabilistic-latent idea you build here is the ancestor of conditional VAEs, VQ-VAE (the discrete codebook behind modern image tokenizers), and latent diffusion — Stable Diffusion runs its diffusion inside a VAE's latent space.

Drag β and click the latent space to decode any point:
https://dev48v.infy.uk/dl/day42-vae.html