Variational Autoencoders: Probabilistic Generative Models

Rangoli, Reimagined

Every Diwali, in front of homes across India, someone sits down with a tray of coloured powder and draws a rangoli. No two rangolis are ever pixel-for-pixel identical, yet every one is unmistakably a rangoli — symmetric, built on a grid of dots, drawn from a familiar vocabulary of petals, loops, and borders. The artist did not memorise thousands of exact patterns. Somewhere along the way, they internalised a small set of adjustable "knobs": how many-fold symmetry to use, how tightly to pack the dots, which colours to place next to each other, how ornate the border should be. Any reasonable setting of those knobs produces something that looks right. Nudge the knobs a little, and a new but equally valid design appears.

This is close to the problem a machine faces when it is asked to generate new images, new molecules, or new sounds: can it study thousands of examples and discover a small number of continuous, well-behaved "knobs" on its own, such that any setting of those knobs — including combinations it has never seen together before — decodes into something realistic? A Variational Autoencoder (VAE), introduced by Diederik Kingma and Max Welling in a 2013 paper titled "Auto-Encoding Variational Bayes" (published at the ICLR 2014 conference, with a closely related approach proposed independently around the same time by Danilo Rezende and colleagues at DeepMind), is one of the earliest and most influential neural network architectures built to do exactly this. It does not just compress data — it learns to draw.

Why a Plain Autoencoder Cannot Generate

To see what makes a VAE different, start with an ordinary autoencoder. It has two parts: an encoder that squeezes an input x down through a narrow bottleneck layer into a short code z, and a decoder that expands z back out into a reconstruction of x. The whole system is trained end to end to minimise reconstruction error — how different the decoder's output is from the original input. Once trained, the encoder is an excellent compressor and the decoder an excellent decompressor, for the specific kinds of input the network has seen.

The trouble starts the moment you try to use it as a generator. To create something new, you would want to invent a plausible code z — without any input x — and hand it straight to the decoder. But an ordinary autoencoder gives no guidance on what a "plausible" z looks like. Nothing during training ever told the encoder where in code-space to place each input; it only had to keep codes far enough apart to reconstruct every input correctly. The result is a latent space full of isolated islands and unpredictable stretching: one region might correspond to real, recognisable outputs, the neighbouring region might decode into meaningless noise, and there is no way to tell them apart without already knowing the answer. A vanilla autoencoder is a superb librarian for cataloguing existing books; it has no idea how to write a new one.

The Generative Story: From Latent Code to Data

A VAE fixes this by starting from the opposite direction and asking a probabilistic question. Imagine every rangoli photograph in a dataset was produced by, first, drawing a short vector of "knob settings" z from a simple, well-known prior distribution p(z) — typically a standard Gaussian, z ~ N(0, I) — and then feeding z through a decoder network to produce the image, described by a generative model p(x|z). If a decoder can be trained so that this story becomes plausible for real data, generating new examples becomes trivial: sample a fresh z from the simple prior and decode it.

The difficulty lies in training such a decoder. Fitting p(x|z) to real data by maximum likelihood would require computing p(x) = ∫ p(x|z) p(z) dz — averaging the decoder's output over every conceivable z. For a neural-network decoder, this integral has no closed form and cannot be evaluated directly. Nor can we easily ask, "given this real photograph x, which z produced it?" The true posterior distribution, p(z|x) = p(x|z)p(z) / p(x), needs that same intractable p(x) in its denominator.

The VAE's answer is to stop trying to compute the true posterior and instead learn an approximate posterior q(z|x) with a second neural network — the encoder, sometimes called a recognition network. Given an input x, the encoder does not output a single code; it outputs the parameters of a distribution over plausible codes, typically a diagonal Gaussian described by a mean and a variance for each latent dimension. Because the same encoder is reused for every input rather than solving a fresh optimisation problem for each data point, this is called amortised inference: the cost of learning to infer is paid once, during training, and every future input gets its posterior estimate from a single fast forward pass.

The Loss Function: The Evidence Lower Bound

With an encoder q(z|x) and decoder p(x|z) in place, training is framed as maximising a quantity called the Evidence Lower Bound, or ELBO. Using Jensen's inequality, it can be shown that for any choice of q(z|x):

log p(x)  ≥  E[log p(x|z)] − KL(q(z|x) || p(z)),   where z ~ q(z|x)

The right-hand side is tractable to compute and differentiate even though log p(x) itself is not. Maximising the ELBO pushes it as close as possible to the true, unreachable log p(x), while simultaneously pulling q(z|x) toward the true posterior. Flip the sign, and the quantity actually minimised during training splits into two familiar ingredients:

Loss(x)  =  −E[log p(x|z)]  +  KL(q(z|x) || p(z))
         =  Reconstruction loss  +  KL regularisation term

The first term is exactly the reconstruction loss from an ordinary autoencoder: how well can the decoder rebuild x from a sampled code z? For pixel intensities in [0, 1], this is usually implemented as binary cross-entropy; for continuous data it is often mean squared error. The second term is the Kullback-Leibler (KL) divergence, a standard measure of how much one probability distribution differs from another. It penalises the encoder for producing a q(z|x) that strays from the simple prior p(z) = N(0, I). When both q(z|x) and p(z) are Gaussian, this term has a clean closed form that needs no sampling to compute. For a diagonal Gaussian with mean μ and variance σ² in each of J latent dimensions:

KL(q(z|x) || p(z))  =  −½ · Σ_j (1 + log σ_j² − μ_j² − σ_j²),   j = 1 … J

This one formula is doing a lot of work: it equals zero exactly when μ = 0 and σ² = 1 in every dimension — the encoder's output matching the prior perfectly — and grows whenever the mean drifts from zero or the variance drifts from one in either direction.

The Reparameterization Trick

There is a subtle problem hiding inside the ELBO. The reconstruction term is an expectation over z ~ q(z|x), so computing its gradient requires sampling z from a distribution whose own parameters, μ(x) and σ(x), are outputs of the encoder we are trying to train. Standard backpropagation cannot send a gradient through a random sampling step: "sample z from N(μ, σ²)" is not a differentiable operation with respect to μ and σ.

The reparameterization trick fixes this with a small but powerful rewrite. Instead of sampling z directly from N(μ, σ²), sample a fixed-distribution noise variable ε ~ N(0, I) that depends on no network parameter, and compute z deterministically from it:

z  =  μ + σ · ε,   where ε ~ N(0, I)

It helps to think of this as separating a recipe from its randomness. Instead of a chef rolling dice in the middle of cooking — impossible to analyse cleanly — the dice are rolled once, outside the kitchen, and the result, ε, is handed in as an ordinary ingredient. Everything from that point on, including how ε combines with μ and σ to form z, is ordinary, differentiable arithmetic. Gradients can now flow from the reconstruction loss, through z, and into μ(x) and σ(x), exactly as they flow through any other layer. An older alternative, the score-function or REINFORCE estimator, avoids this rewrite but produces gradient estimates with much higher variance, which slows training considerably. The reparameterization trick is the reason VAEs can be trained end to end with ordinary stochastic gradient descent.

The shift from a plain autoencoder to a VAE can be summarised in three changes:

• Encoder output: a single point z in a plain autoencoder; the parameters of a distribution, μ and σ², in a VAE.
• Training objective: reconstruction error alone in a plain autoencoder; reconstruction error plus KL regularisation toward N(0, I) in a VAE.
• Latent space: unpredictable and full of gaps in a plain autoencoder; smooth and densely organised around the prior in a VAE.

A Fully Worked Example: Encoding One Rangoli Tile

To see every step with real numbers, shrink the problem down to a size that can be computed by hand. Instead of a full rangoli photograph, take one tiny 2×2 tile from a black-and-white sketch — four numbers, each either 1 (filled) or 0 (empty) — and use a latent space of just two dimensions instead of the dozens or hundreds a real system would use.

Suppose the target tile, read row by row, is x = [1, 0, 1, 0]: a filled cell, an empty cell, a filled cell, an empty cell. Suppose the encoder, on seeing this tile, outputs:

μ       =  [ 0.00,  0.40]
log σ²  =  [−0.20,  0.10]

Step 1 — convert log-variance to standard deviation. The encoder outputs log σ² rather than σ² directly because exponentiating guarantees a positive variance no matter what real number the network produces. Using σ = exp(0.5 · log σ²):

σ₁ = exp(0.5 × −0.20) = exp(−0.10) ≈ 0.9048
σ₂ = exp(0.5 ×  0.10) = exp( 0.05) ≈ 1.0513

Step 2 — apply the reparameterization trick. Draw one noise sample, ε = [0.50, −1.00], from N(0, I), and compute z = μ + σ · ε elementwise:

z₁ = 0.00 + 0.9048 × 0.50    =  0.4524
z₂ = 0.40 + 1.0513 × (−1.00) = −0.6513

Step 3 — decode. Passing z = [0.4524, −0.6513] through the decoder produces reconstruction probabilities for the four cells, each the output of a final sigmoid unit and readable as "the network's confidence this cell is filled":

x̂ = [0.88, 0.15, 0.77, 0.22]

Step 4 — reconstruction loss. Compare x̂ against the true tile x = [1, 0, 1, 0] with binary cross-entropy, −[x · log x̂ + (1 − x) · log(1 − x̂)], summed over the four cells:

Cell 1: x=1, x̂=0.88 → −log(0.88)     ≈ 0.128
Cell 2: x=0, x̂=0.15 → −log(1−0.15)   ≈ 0.163
Cell 3: x=1, x̂=0.77 → −log(0.77)     ≈ 0.261
Cell 4: x=0, x̂=0.22 → −log(1−0.22)   ≈ 0.248
                              Total  ≈ 0.800

Step 5 — KL divergence. Apply the closed-form formula to each latent dimension:

j=1: 1 + (−0.20) − 0.00² − exp(−0.20) = 1 − 0.20 − 0 − 0.819 = −0.019
j=2: 1 +   0.10  − 0.40² − exp( 0.10) = 1.10 − 0.16 − 1.105 = −0.165

KL = −½ × (−0.019 + −0.165) = −½ × (−0.184) ≈ 0.092

Step 6 — total loss. Add the two terms:

Loss  =  Reconstruction + KL  ≈  0.800 + 0.092  ≈  0.892

Every one of these numbers now has a clear direction to move in. Backpropagating this loss nudges the encoder's weights so that this tile earns a tighter, more accurate μ and σ, and nudges the decoder's weights so that a z near [0.4524, −0.6513] decodes closer to [1, 0, 1, 0]. Repeat this across every tile in the dataset, many times over, and the encoder learns to place similar tiles near each other in latent space while the decoder learns to turn any nearby point into something tile-like.

Implementing a VAE in PyTorch

The architecture from the worked example — an encoder producing μ and log σ², the reparameterization trick, a decoder, and the two-term loss — translates directly into code. Here is a complete, minimal VAE for 28×28 grayscale images (784 pixels), the size used by the classic MNIST handwritten-digit dataset:

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

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)         # encoder hidden layer
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)       # outputs mean
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)   # outputs log-variance
        self.fc2 = nn.Linear(latent_dim, hidden_dim)         # decoder hidden layer
        self.fc3 = nn.Linear(hidden_dim, input_dim)          # outputs reconstruction

    def encode(self, x):
        h = F.relu(self.fc1(x))
        return self.fc_mu(h), self.fc_logvar(h)

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

    def decode(self, z):
        h = F.relu(self.fc2(z))
        return torch.sigmoid(self.fc3(h))

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

def vae_loss(x_hat, x, mu, logvar):
    recon = F.binary_cross_entropy(x_hat, x, reduction='sum')
    kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon + kl

# one training step
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
x_hat, mu, logvar = model(x_batch)         # x_batch: shape (batch, 784), values in [0, 1]
loss = vae_loss(x_hat, x_batch, mu, logvar)
optimizer.zero_grad()
loss.backward()
optimizer.step()

Notice how directly the code mirrors the mathematics: reparameterize is the line z = μ + σ·ε, unchanged; vae_loss is the reconstruction-plus-KL formula from the previous section, unchanged. Everything from encode to loss.backward() forms a single differentiable chain — exactly what the reparameterization trick made possible.

Why the KL Term Matters: Taming the Latent Space

It is worth asking what would happen if the KL term were simply deleted, leaving only reconstruction loss. The network would still train, and reconstructions might even look sharper — but the model would quietly collapse back into an ordinary autoencoder. Nothing would stop the encoder from shrinking σ toward zero, producing near-deterministic codes, and scattering different inputs' μ values wherever is most convenient for reconstruction, with no obligation to stay near the origin or near each other. The latent space would fill again with gaps that were never decoded during training.

The KL term prevents exactly this. By constantly pulling every input's q(z|x) back toward the shared prior N(0, I), it forces the probability "clouds" belonging to different inputs to overlap rather than scatter, and keeps every cloud centred near the origin with roughly unit spread. The result is a latent space where the region that matters — everything within a few standard deviations of the origin — has been visited by the decoder during training and therefore decodes into something plausible. This is fundamentally a tug-of-war: reconstruction loss wants codes as precise and as far apart as needed to tell inputs apart perfectly; the KL term wants every code squeezed into the same standard-normal-shaped neighbourhood. The balance between the two is what produces a latent space that is both meaningful and generative. Researchers have even experimented with increasing the weight on the KL term, an approach called β-VAE, proposed by Irina Higgins and colleagues at DeepMind, to encourage each latent dimension to capture a single, more human-interpretable factor of variation.

Generating Brand-New Patterns

Once training is complete, the encoder can be set aside entirely for generation. To create a rangoli tile that never existed in the training data, simply sample z directly from the prior, z ~ N(0, I), and pass it through the decoder alone — no input photograph required. Because the KL term has made the whole neighbourhood around the origin decodable, almost any z drawn this way produces something recognisable.

A second, equally telling capability is latent space interpolation: encode two real tiles to get their mean codes z_a and z_b, walk along a straight line between them, and decode several points along the way. In a well-trained VAE, this produces a smooth, gradual morph from one pattern into the other, because nearby points in latent space genuinely correspond to similar outputs. A vanilla autoencoder, whose latent space was never organised this way, typically produces nonsense partway along such a walk. This same smoothness is also why VAE reconstructions tend to look a little soft compared with a rival family of generative models, Generative Adversarial Networks (GANs): a VAE explicitly optimises a likelihood-based, averaging objective, which favours safe, smooth outputs, while a GAN is trained against an adversarial discriminator that rewards sharp, realistic detail at the cost of a less stable training process.

From Rangoli to Real-World Generative AI

The same idea — an encoder-decoder pair with a regularised latent space in between — now shows up well beyond toy images:

• Drug discovery: VAEs trained on molecular structures let chemists pick a point in a smooth "chemical latent space" and decode a candidate molecule with a desired property, rather than searching a discrete database by trial and error.
• Recommendation systems: a VAE can model which items a user is likely to want by treating their interaction history as the observed x and their underlying taste as the latent z.
• Anomaly detection: in manufacturing and cybersecurity, a VAE trained only on normal behaviour reconstructs normal inputs well and unfamiliar ones poorly, so unusually high reconstruction error is itself the useful output, flagging defects or intrusions.

Return, finally, to the rangoli artist. Everything a trained VAE does mirrors what that artist has already done without ever writing an equation: look at enough real examples to internalise a compact set of continuous, meaningful controls (the encoder); make sure any reasonable setting of those controls corresponds to something valid rather than a small memorised catalogue (the KL-regularised latent space); and turn a fresh setting of those controls into a finished, coherent piece of work (the decoder). The rangoli artist never needs to have seen this exact pattern before to draw it convincingly. Neither, now, does the machine.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where variational autoencoders: probabilistic generative models is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting variational autoencoders: probabilistic generative models to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind variational autoencoders: probabilistic generative models, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.