A hospital network wants to train a sepsis early-warning model on ICU vitals: heart rate, SpO2, systolic BP, respiration rate, sampled every minute. The model needs thousands of patient-hours to learn the subtle drift that precedes septic shock. A startup building that model does not have access to the hospital's raw records, and under India's Digital Personal Data Protection Act, 2023, it should not: patient vitals are personal data, and moving them off the hospital's systems requires consent and safeguards the startup cannot easily provide. The hospital's data science team has an idea: train a generative model on the real vitals inside the hospital's own infrastructure, then hand the startup only synthetic patient-hours, statistically realistic time series that were never attached to a real person. If the synthetic data preserves the patterns that precede sepsis, the startup can train on it without ever touching a real record.
This is exactly the problem a variational autoencoder is built to solve, and it exposes precisely why an ordinary autoencoder cannot do the job. This chapter builds the VAE mechanism from first principles: the probabilistic reframing that turns "generate data" into a tractable optimization problem, the derivation of the objective (the ELBO) that makes that optimization possible, the encoder-decoder architecture that implements it, and the reparameterization trick, the one piece of engineering that makes the whole thing trainable with ordinary backpropagation.
Why a plain autoencoder cannot dream
A standard autoencoder learns two functions: an encoder that maps an input vitals-vector x to a compressed code z, and a decoder that maps z back to a reconstruction x̂, trained so that x̂ matches x. Once trained, the encoder maps every real patient-hour in the training set to some specific point in latent space. Nothing in the training objective says anything about the points between those specific codes, or about what a decoder should output there. The hospital's team could sample a random vector from that latent space and feed it to the decoder, but nothing constrains that region to decode into a physiologically plausible vitals trace. It might decode into noise, or into an average of two unrelated patients that itself corresponds to no real physiological state. A plain autoencoder is a compression scheme, not a model of how vitals-data is distributed. It has no notion of a probability density over the space of possible patients, so it has nothing to sample from.
What is missing is a generative model in the formal sense: a joint distribution p(x, z) over data and a latent cause, factored as p(x, z) = p(x | z) p(z), where p(z) is a simple prior (say, a standard normal) that we know how to sample from, and p(x | z) is a decoder that turns a sampled z into a distribution over plausible vitals traces. If we can fit such a model to the hospital's real data, generating a synthetic patient-hour is just: sample z from the prior, sample x from p(x | z). The entire difficulty of a VAE is fitting this model without ever being able to compute the one quantity that would make fitting it easy.
The intractable posterior, and the trick that sidesteps it
Fitting a generative model by maximum likelihood means maximizing log p(x) for the training data, where p(x) = ∫ p(x | z) p(z) dz. That integral runs over every possible latent code, weighted by how well it explains x. For anything but the simplest decoder, it has no closed form and cannot be evaluated numerically at the resolution needed for training. The natural fix from Bayesian statistics is to work with the posterior p(z | x), the distribution over latent causes given that we observed this particular x. But p(z | x) = p(x | z) p(z) / p(x) requires exactly the intractable p(x) in its denominator. We are stuck in a circle: to train the decoder we need the posterior, and the posterior needs the very quantity training is supposed to produce.
Kingma and Welling's 2013 paper "Auto-Encoding Variational Bayes" (submitted to arXiv in December 2013, presented at ICLR 2014) breaks the circle by giving up on the exact posterior and fitting an approximation to it instead. Introduce a second network, the encoder, that outputs a distribution q(z | x) meant to approximate p(z | x). Because q is a distribution we choose the form of (a diagonal Gaussian, in the standard VAE), we can evaluate it, sample from it, and differentiate through it, none of which is true of the real posterior. Danilo Rezende, Shakir Mohamed and Daan Wierstra published essentially the same mechanism independently the same year, under the name "stochastic backpropagation," at ICML 2014. Both papers converge on the same two ideas this chapter builds: an objective that only needs the approximate posterior, and a trick for getting gradients through a sampling step.
Deriving the ELBO: what we can actually optimize
Start from the log-likelihood of a single data point and introduce the approximate posterior q(z | x) by multiplying and dividing by it inside the integral:
log p(x) = log ∫ p(x, z) dz
= log ∫ q(z|x) · [p(x, z) / q(z|x)] dz
= log E_q[ p(x, z) / q(z|x) ]The logarithm is a concave function, so Jensen's inequality gives log E[Y] ≥ E[log Y] for any random variable Y. Applying it here:
log p(x) ≥ E_q[ log p(x, z) − log q(z|x) ] ≜ ELBO(x)This lower bound is the Evidence Lower BOund, ELBO for short: "evidence" is Bayesian jargon for p(x). Expanding p(x, z) = p(x|z) p(z) inside the expectation and regrouping:
ELBO(x) = E_q[ log p(x|z) + log p(z) − log q(z|x) ]
= E_q[ log p(x|z) ] − E_q[ log q(z|x) − log p(z) ]
= E_q[ log p(x|z) ] − KL( q(z|x) ‖ p(z) )Every term on the last line is computable: the first is an expected reconstruction log-likelihood under samples from the encoder, the second is a KL divergence between two Gaussians we will derive a closed form for shortly. It is worth checking exactly how loose the bound is, because the gap tells us what optimizing the ELBO is really doing. Write the KL divergence between the approximate posterior and the true, intractable posterior:
KL( q(z|x) ‖ p(z|x) ) = E_q[ log q(z|x) − log p(z|x) ]Substitute p(z|x) = p(x, z) / p(x), so log p(z|x) = log p(x,z) − log p(x):
KL( q(z|x) ‖ p(z|x) ) = E_q[ log q(z|x) − log p(x,z) ] + log p(x)
= −ELBO(x) + log p(x)Rearranged, log p(x) = ELBO(x) + KL( q(z|x) ‖ p(z|x) ). Since a KL divergence is never negative, this confirms log p(x) ≥ ELBO(x) directly, without invoking Jensen a second time, and it says something sharper: the gap between the true log-likelihood and the bound we are optimizing is exactly the divergence between our approximate posterior and the true one. Maximizing the ELBO with respect to the encoder's parameters, holding the decoder fixed, is mathematically identical to minimizing that gap. The encoder is not a heuristic bolted onto a generative model; it is doing genuine approximate Bayesian inference, and the ELBO is the one objective that trains encoder and decoder together toward the same goal: encoder parameters that make q(z|x) a tight approximation to the real posterior, and decoder parameters that make p(x|z) assign high probability to the real data.
Two networks, three outputs
In practice, both p(x|z) and q(z|x) are implemented as neural networks whose outputs parameterize a chosen distribution family rather than a single number. For a diagonal Gaussian posterior in a latent space of dimension d, the encoder takes x and produces two vectors of length d: a mean μ(x) and, critically, not the variance directly but log σ²(x). This choice is deliberate, not incidental. A plain linear output layer can produce any real number, positive or negative, with no constraint. Variance must be strictly positive, so predicting σ² directly would need an extra positivity-enforcing activation (softplus, exp) layered on afterward, and gradients near zero variance become poorly behaved. Predicting log σ² sidesteps this entirely: it is unconstrained, matches the term that already appears inside the KL-divergence formula (derived below), and recovering σ is one exp(0.5 · log σ²) away.
The decoder takes a sampled latent code z and produces the parameters of p(x|z). For continuous data such as normalized vitals, that is usually a mean vector for a Gaussian with fixed variance, which makes the reconstruction term reduce to mean-squared error up to a constant. For binary or [0,1]-scaled pixel data, it is Bernoulli parameters via a sigmoid, which turns the reconstruction term into binary cross-entropy. Both are the same E_q[log p(x|z)] term in the ELBO under a different choice of output distribution; only the loss formula on the page changes.
The reparameterization trick
Training either network by gradient descent requires ∂ELBO / ∂(encoder weights). The reconstruction term is an expectation over z ~ q(z|x), and the sampling step z ~ N(μ(x), σ²(x)) is the problem: sampling from a distribution is not a differentiable operation with respect to that distribution's parameters. There is no way to write "how does this random draw change if I nudge μ by a small amount" using ordinary calculus, because the draw is a single stochastic outcome, not a smooth function of μ.
The fix is to rewrite the sampling step so that all the randomness comes from a source that does not depend on the parameters we are trying to learn. Instead of drawing z directly from N(μ, σ²), draw ε from a fixed, parameter-free distribution N(0, 1), and compute:
z = μ(x) + σ(x) · ε, ε ~ N(0, I)This z has exactly the same distribution as before (a location-scale transform of a standard normal is itself normal with that location and scale), so nothing about what the encoder represents has changed. What has changed is where the randomness sits in the computational graph. ε is now an external input, sampled once per forward pass and then held fixed for that pass, exactly like a fixed constant would be. The path from μ and σ to z is now pure arithmetic: multiplication and addition, both differentiable. Backpropagation can flow through z to μ and σ using ordinary chain rule, because from the graph's point of view ε is just a number that happened to be 0.8 this time, not a source of non-differentiability.
The KL term in closed form
The second half of the (negative) ELBO, the KL divergence between the encoder's q(z|x) = N(μ, σ²) and the prior p(z) = N(0, 1), has an exact closed form for Gaussians, so it needs no sampling at all. Derive it directly from the definition, for one latent dimension:
KL(q‖p) = ∫ q(z) [log q(z) − log p(z)] dz
log q(z) = −½ log(2πσ²) − (z−μ)²/(2σ²)
log p(z) = −½ log(2π) − z²/2
log q(z) − log p(z) = −½ log σ² − (z−μ)²/(2σ²) + z²/2Taking the expectation under q, and using E_q[(z−μ)²] = σ² and E_q[z²] = σ² + μ² (the second moment of a Gaussian with mean μ and variance σ²):
KL(q‖p) = −½ log σ² − σ²/(2σ²) + (σ²+μ²)/2
= ½(σ² + μ² − 1 − log σ²)For a d-dimensional diagonal Gaussian, sum this expression across dimensions, since the joint KL of independent factors is the sum of the per-dimension KLs.
Worked example: one complete forward and backward pass
Take a 2-dimensional latent space and a single, simplified decoder so every number can be checked by hand. Suppose the encoder, looking at one patient-hour vector, outputs μ = [0.5, −0.3] and log σ² = [−0.2, 0.1]. Recovering the standard deviations uses σ = exp(½ log σ²):
σ₁ = exp(0.5 × (−0.2)) = exp(−0.1) = 0.9048
σ₂ = exp(0.5 × 0.1) = exp(0.05) = 1.0513Draw noise for this forward pass, say ε = [0.8, −0.5], and reparameterize:
z₁ = μ₁ + σ₁·ε₁ = 0.5 + 0.9048 × 0.8 = 1.2239
z₂ = μ₂ + σ₂·ε₂ = −0.3 + 1.0513 × (−0.5) = −0.8256Feed z through a toy linear decoder, x̂ = w₁z₁ + w₂z₂ + b, with w = [0.4, −0.6], b = 0.1 (a linear decoder keeps the arithmetic checkable; a real decoder is a multi-layer network but every gradient step below composes the same way through it):
x̂ = 0.4 × 1.2239 + (−0.6) × (−0.8256) + 0.1
= 0.4896 + 0.4954 + 0.1 = 1.0849If the true (normalized) target for this vitals feature is x = 1.0, the reconstruction loss (squared error) is (1.0849 − 1.0)² = (0.0849)² = 0.0072. The KL term, per dimension, using the closed form derived above:
KL₁ = ½(0.8187 + 0.25 − 1 − (−0.2)) = ½(0.2687) = 0.1344
KL₂ = ½(1.1052 + 0.09 − 1 − 0.1) = ½(0.0952) = 0.0476
KL_total = 0.1344 + 0.0476 = 0.1820Total loss (reconstruction plus KL, unweighted) is 0.0072 + 0.1820 = 0.1892. Now trace the gradient back to the encoder's four outputs, which is the entire point of the reparameterization trick: d(recon)/d(x̂) = 2 × 0.08493 = 0.1699. Through the decoder weights, d(recon)/dz₁ = 0.1699 × 0.4 = 0.0679 and d(recon)/dz₂ = 0.1699 × (−0.6) = −0.1019. Because z = μ + σε with ε held fixed, dz/dμ = 1 and dz/dσ = ε, so:
d(recon)/dμ₁ = 0.0679 × 1 = 0.0679
d(recon)/dμ₂ = −0.1019 × 1 = −0.1019
d(recon)/dσ₁ = 0.0679 × 0.8 = 0.0543; dσ₁/d(logσ²₁) = 0.5σ₁ = 0.4524
⇒ d(recon)/d(logσ²₁) = 0.0543 × 0.4524 = 0.0246
d(recon)/dσ₂ = −0.1019 × (−0.5) = 0.0510; dσ₂/d(logσ²₂) = 0.5σ₂ = 0.5256
⇒ d(recon)/d(logσ²₂) = 0.0510 × 0.5256 = 0.0268Add the KL term's own gradient, using d(KL)/dμ = μ and d(KL)/d(logσ²) = ½(σ² − 1), both read straight off the closed form derived earlier:
d(KL)/dμ₁ = 0.5; d(KL)/d(logσ²₁) = ½(0.8187−1) = −0.0906
d(KL)/dμ₂ = −0.3; d(KL)/d(logσ²₂) = ½(1.1052−1) = 0.0526
Total gradients:
d(loss)/dμ₁ = 0.0679 + 0.5 = 0.5679
d(loss)/dμ₂ = −0.1019 − 0.3 = −0.4019
d(loss)/d(logσ²₁) = 0.0246 − 0.0906 = −0.0660
d(loss)/d(logσ²₂) = 0.0268 + 0.0526 = 0.0794Every one of these four numbers exists only because ε was treated as a constant rather than something to sample and forget. Before reparameterization, there was no path from a loss computed on x̂ back to μ and logσ² at all; the sampling step was a wall gradients could not cross.
In code
The same four quantities, structured as a trainable module (a real decoder would have hidden layers rather than the single linear step used above, but the reparameterization line is identical):
import torch
import torch.nn as nn
import torch.nn.functional as F
class VAE(nn.Module):
def __init__(self, input_dim, hidden_dim, latent_dim):
super().__init__()
self.enc_hidden = nn.Linear(input_dim, hidden_dim)
self.enc_mu = nn.Linear(hidden_dim, latent_dim)
self.enc_logvar = nn.Linear(hidden_dim, latent_dim)
self.dec_hidden = nn.Linear(latent_dim, hidden_dim)
self.dec_out = nn.Linear(hidden_dim, input_dim)
def encode(self, x):
h = F.relu(self.enc_hidden(x))
return self.enc_mu(h), self.enc_logvar(h)
def reparameterize(self, mu, logvar):
sigma = torch.exp(0.5 * logvar)
epsilon = torch.randn_like(sigma)
return mu + sigma * epsilon
def decode(self, z):
h = F.relu(self.dec_hidden(z))
return self.dec_out(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.mse_loss(x_hat, x, reduction="sum")
kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return recon + klvae_loss is the negative ELBO for a Gaussian decoder: recon is (up to a constant) −E_q[log p(x|z)], and the KL line is exactly ½(σ² + μ² − 1 − log σ²) summed over every latent dimension and every example in the batch, since -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()) is algebraically identical to that closed form. In practice this sum is divided by the batch size before logging, so the reported number is comparable across batch sizes; the code above sums to match the derivation exactly. Once trained, generating a synthetic patient-hour for the hospital's use case needs only the decoder: sample z from N(0, I) directly (no encoder involved, since there is no real x to condition on) and call model.decode(z). Whether the resulting latent space is smooth enough to interpolate between two patients, and how this generation quality holds up specifically for image data, are the questions the other chapters in this unit take further; the mechanism that makes generation possible at all is everything above.
Diagram: the mechanism end to end
The misconception worth correcting
Looking at z = μ + σ · ε for the first time, students commonly conclude that the reparameterization trick removes the randomness from a VAE, turning the latent code into a deterministic function of the input, just like an ordinary autoencoder's bottleneck. This is backwards, and it matters, because if it were true the whole scheme would collapse back into the plain autoencoder from the opening section. ε is drawn fresh from N(0, I) on every forward pass during training; the same input x produces a different z almost every time it is encoded, exactly as it should for z to be a genuine sample from q(z|x) rather than a single fixed point. What the trick removes is not the randomness itself but its position in the computational graph: instead of the parameters μ and σ feeding into an opaque sampling operation that blocks gradients, they now feed into an ordinary multiply-and-add whose other input, ε, happens to be random but is treated as a constant for the purpose of differentiation, the same way a fixed data point is a constant with respect to a weight. The randomness that makes E_q[log p(x|z)] a proper Monte Carlo estimate of an expectation is fully intact; only the non-differentiable node has moved to a place where it no longer needs a gradient.
Active recall
Attempt each question before reading its answer.
1. Why does maximizing the ELBO also improve the approximate posterior, not just the decoder?
2. A one-dimensional encoder outputs μ = 1.0 and log σ² = 0. Compute the KL divergence to N(0,1) by hand.
3. If the KL weight is set to zero during training (optimizing only reconstruction), what happens to σ, and why does that break generation?
4. In the worked example, suppose the encoder had instead output log σ²₁ = 0.4 (μ, ε, and dimension 2 all unchanged). Recompute σ₁, z₁, x̂, the reconstruction loss, the total loss, and every one of the four gradients d(loss)/dμ₁, d(loss)/d(logσ²₁), d(loss)/dμ₂, d(loss)/d(logσ²₂). Which of these change, and why might that be surprising?
5. Why does the encoder output log σ² instead of σ or σ² directly?
6. What does the reparameterization trick actually make differentiable: the sample z, or the distribution q(z|x) itself?
Answers
1. The exact identity log p(x) = ELBO(x) + KL(q(z|x) ‖ p(z|x)) holds for any choice of q. Since log p(x) does not depend on the encoder's parameters at all (it depends only on the decoder and the true data distribution), any increase in ELBO(x) achieved by changing the encoder must come entirely out of the KL term on the right, which can only decrease. So gradient ascent on the ELBO with respect to encoder parameters is, term for term, gradient descent on the distance between the approximate and true posteriors.
2. KL = ½(σ² + μ² − 1 − log σ²) = ½(1 + 1 − 1 − 0) = ½(1) = 0.5 nats. (σ² = exp(0) = 1 since log σ² = 0.)
3. With no KL pressure, the objective only rewards accurate reconstruction, and any nonzero σ injects noise into z that can only hurt reconstruction accuracy (noise moves z away from the single best-reconstructing point). Gradient descent therefore drives σ toward zero for every input, collapsing q(z|x) into a near-deterministic spike at whatever point best reconstructs that particular x. The model degenerates into an ordinary autoencoder: each training point gets its own isolated code, with no pressure for those codes to tile the space smoothly or to overlap with the region the prior N(0,I) actually samples from, so decoding a fresh z ~ N(0,I) at generation time lands in territory the decoder was never trained on.
4. New σ₁ = exp(0.5 × 0.4) = exp(0.2) = 1.2214 (σ₁² = exp(0.4) = 1.4918). New z₁ = 0.5 + 1.2214 × 0.8 = 1.4771; z₂ is untouched at −0.8256. New x̂ = 0.4 × 1.4771 − 0.6 × (−0.8256) + 0.1 = 0.5908 + 0.4954 + 0.1 = 1.1862. Reconstruction loss = (0.1862)² = 0.0347, versus 0.0072 before: worse, because z₁ moved further from the value that minimized error. KL₁ = ½(1.4918 + 0.25 − 1 − 0.4) = ½(0.3418) = 0.1709, versus 0.1344 before: also worse, since σ₁² has moved further from the prior's variance of 1. Total loss = 0.0347 + 0.1709 + 0.0476 = 0.2532, up from 0.1892.
Gradients: d(recon)/d(x̂) = 2 × 0.1862 = 0.3725. d(recon)/dz₁ = 0.3725×0.4 = 0.1490, d(recon)/dz₂ = 0.3725×(−0.6) = −0.2235. d(loss)/dμ₁ = 0.1490 + 0.5 (KL term unchanged since μ₁ unchanged) = 0.6490. d(loss)/d(logσ²₁): d(recon)/dσ₁ = 0.1490×0.8 = 0.1192, ×0.5σ₁ (=0.6107) = 0.0728; KL part = ½(1.4918−1) = 0.2459; total = 0.3187.
The part students usually miss: μ₂ and log σ²₂ were never touched, and z₂ itself is numerically identical to before, yet their gradients change anyway. d(loss)/dμ₂ = d(recon)/dz₂ + d(KL)/dμ₂ = −0.2235 + (−0.3) = −0.5235 (was −0.4019). d(loss)/d(logσ²₂) = (−0.2235 × ε₂ × 0.5σ₂) + 0.0526 = (−0.2235 × (−0.5) × 0.5256) + 0.0526 = 0.0587 + 0.0526 = 0.1113 (was 0.0794). Both shifted because d(recon)/dz₂ depends on the shared error signal (x̂ − x), which changed once z₁ moved the decoder's output. Every latent dimension that feeds the same decoder output shares that error term, so a change confined to one encoder output ripples into the gradients of every other dimension even when their own values never move.
5. A linear output layer is naturally unconstrained, matching log σ², whereas σ and σ² are both restricted to positive values and would need an extra activation (softplus or exp) to enforce that, with worse-behaved gradients as the variance approaches zero. log σ² also plugs directly into the closed-form KL expression derived above without any extra transformation.
6. Neither the sample nor the distribution becomes "differentiable" in isolation; a single random draw is still random, and a distribution is not a function that has a derivative in the input-output sense. What the trick makes differentiable is the map from the encoder's parameters (μ, σ) to the sampled value z, by rewriting that map as deterministic arithmetic applied to an independent noise source, so that gradients of any downstream loss with respect to μ and σ can be computed by ordinary backpropagation through that arithmetic.
Think About It
Think about this: How would you explain variational autoencoders: teaching machines to dream to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind variational autoencoders: teaching machines to dream, 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.