Variational Inference: Why the ELBO Trades Accuracy for Tractability
Bayesian inference is easy to write down and hard to compute. Given data and latent variables , Bayes' rule gives the posterior:
The numerator is usually easy: a prior and a likelihood . The denominator, the evidence , is the problem. It is an integral over every possible value of :
For any realistic model, that integral has no closed form and is too high-dimensional to compute numerically. Markov chain Monte Carlo (MCMC) attacks the problem by sampling, and given enough time it can be asymptotically exact. But sampling is slow, hard to scale to large data sets, and awkward to plug into a gradient-based training loop. Variational inference (VI) takes the opposite bargain: give up exactness, gain speed, differentiability, and scale. The price of that bargain, and how to manage it, is the whole subject of this article.
Inference as optimisation
Instead of computing the posterior exactly, VI picks a family of tractable distributions , indexed by parameters , and searches for the member of that family closest to the true posterior. Closeness is measured by the Kullback-Leibler (KL) divergence:
This is a clean objective, but it cannot be optimised directly, because it contains the very posterior we cannot compute. The trick of VI is to rewrite it into something computable. Expand the posterior using Bayes' rule, and the log evidence separates out:
where ELBO stands for the evidence lower bound:
Since the KL term is never negative, the ELBO is always a lower bound on . And since does not depend on , maximising the ELBO is exactly equivalent to minimising the KL divergence to the true posterior. That is the central identity of VI, and it is worth pausing on: we never evaluate the intractable posterior, yet optimising a tractable bound moves us as close to it as our chosen family allows.
The ELBO has a second, more interpretable form. Split the joint :
Now the trade-off is visible. The first term rewards reconstructions: samples from should explain the data well under the likelihood. The second term penalises deviation from the prior: should not wander too far from what we believed before seeing data. VI is therefore a regularised fit. Drop the KL term and you overfit the latent representation; drop the reconstruction term and you just return the prior. Every variational model lives somewhere on the line between those two failures.
The mean-field choice, and what it costs
The most common family is the mean-field family: assume the latent variables are independent under , so . If each factor is Gaussian, the KL term against a Gaussian prior has a closed form, and optimisation is straightforward. This is why mean-field VI is so popular in practice.
But independence is a strong assumption, and its cost is systematic. The KL divergence used here, , is mode-seeking. Where the true posterior is multimodal, a unimodal will typically lock onto one mode and ignore the rest, because putting mass where is tiny is heavily penalised, while missing a whole mode costs relatively little. Where the true posterior has correlated variables, a factorised cannot represent the correlation at all, and it compensates by shrinking its variances. The result is the classic VI symptom: posterior uncertainty that is too small. Credible intervals from mean-field VI should be treated as optimistic, not calibrated.
This is not a reason to avoid VI. It is a reason to know which question you are asking. If you need a point estimate plus a rough sense of uncertainty, or a latent representation for a downstream task, VI is often excellent. If you need accurate tail probabilities or a faithful picture of a multimodal posterior, MCMC or a richer variational family is the honest choice.
The reparameterisation trick
There is one more obstacle. The ELBO contains an expectation over , and appears both in the distribution being sampled from and in the function being averaged. A naive gradient estimator through samples has high variance or is not even defined for discrete variables.
For continuous latents, the reparameterisation trick solves this cleanly. Instead of sampling directly from , write as a deterministic function of a fixed noise variable . For a Gaussian , that is:
Now the randomness comes from , which does not depend on , and gradients flow through and exactly as in ordinary backpropagation. This single idea is what made variational autoencoders (VAEs) trainable, and it is why VI scales to neural networks at all.
A minimal VAE training step in PyTorch-style pseudocode makes the pieces concrete:
# encoder outputs mean and log-variance of q(z | x)
mu, logvar = encoder(x)
std = torch.exp(0.5 * logvar)
# reparameterisation trick
eps = torch.randn_like(std)
z = mu + eps * std
# decoder gives the likelihood p(x | z)
x_recon = decoder(z)
# ELBO: reconstruction minus KL to a standard normal prior
recon_loss = ((x - x_recon) ** 2).sum(dim=1).mean()
kl = -0.5 * (1 + logvar - mu**2 - logvar.exp()).sum(dim=1).mean()
loss = recon_loss + kl
loss.backward()
Note the structure: the reconstruction term and the KL term are computed separately, and the loss is their sum with a sign flip because we minimise the negative ELBO. If you ever tune a VAE and wonder why the latents look unused or the reconstructions look blurry, the first place to look is the balance between these two terms.
Amortised inference, and the amortisation gap
In classical VI, you optimise separately for each data set. In a VAE, you amortise: a single encoder network maps any input to the parameters of . This is enormously convenient at test time, since inference becomes one forward pass, but it introduces a second source of error on top of the restricted family. Even if the family contains a good approximation to the posterior for a given , the encoder may not find it, because one network must serve all inputs. That shortfall is the amortisation gap, and it is distinct from the approximation gap caused by the family itself. When a VAE underperforms, it is worth asking which gap is to blame: a richer family will not fix an encoder that is too weak, and a bigger encoder will not fix a family that cannot express the posterior.
Limitations and common pitfalls
First, posterior collapse. In a VAE with a powerful decoder, the model can achieve a good ELBO by ignoring the latents entirely: collapses to the prior, the KL term goes to zero, and the decoder reconstructs from its own autoregressive power alone. Symptoms include a KL term stuck near zero and latents that carry no information about the input. Remedies include weakening the decoder, annealing the KL weight from zero up to one during training (so the model learns to use latents before being penalised for it), or using a structured prior.
Second, the ELBO is not a reliable model comparison score across different families. A higher ELBO means a tighter bound for that family, but the gap between the ELBO and the true log evidence depends on the family, so comparing ELBOs across very different models can mislead. For model selection, held-out likelihood estimates or predictive checks are safer.
Third, discrete latents break the reparameterisation trick. Gradients through discrete samples need score-function estimators, which have high variance, or continuous relaxations such as the Gumbel-Softmax, which introduce bias. If your latents are naturally discrete, budget for this difficulty up front rather than discovering it mid-project.
Fourth, optimisation matters more than the theory suggests. The ELBO is generally non-convex, and VI optimisation can stall in poor local optima, especially with a bad initialisation of the variational parameters. Multiple restarts, sensible initialisation, and monitoring both ELBO terms separately (not just their sum) catch most of these failures early.
When to reach for variational inference
Use VI when you need fast, differentiable, approximate posteriors at scale: large data sets, neural latent variable models, or pipelines where inference must run in a single forward pass. Use a richer family, such as a normalizing flow for , when the mean-field approximation is visibly too crude but you still need VI's speed. Reach for MCMC when accuracy of the posterior itself is the deliverable, when the posterior is strongly multimodal, or when you need trustworthy uncertainty quantification and can afford the compute.
The enduring lesson of VI is not the ELBO formula. It is the discipline of naming your approximation: a family, a divergence, and an optimisation procedure. Once those three choices are explicit, the failures of VI become diagnosable rather than mysterious, and the trade it makes, accuracy for tractability, becomes a decision you control rather than a surprise you discover.
Takeaway
Variational inference turns an intractable integral into an optimisation problem by maximising the ELBO, which is equivalent to minimising the KL divergence to the true posterior within a chosen family. The reparameterisation trick makes that optimisation trainable with gradients, and amortisation makes it fast, but both steps add their own gaps: the family may be too simple, the encoder may be too weak, and the KL term may shrink uncertainty or collapse the latents. Name the family, watch the two ELBO terms separately, and choose VI when tractability matters more than exactness.
Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial
Top comments (0)