In this short note, I recall variational inference and the evidence lower bound (ELBO). In contrast to the introductions I have seen, I start by describing variational inference as a loss minimisation problem in its own right without reference to the ELBO, and then derive the ELBO afterwards.

§Bayesian inference

Let’s briefly recall the set-up of Bayesian inference.

• Let $\Theta$ be a parameter space, indexing a parametric statistical model class of interest. (In the following, I will refer to each $\theta \in \Theta$ as parameters or models interchangeably.)

• Let $\varphi \in \Delta(\Theta)$ be a prior distribution encoding our initial belief as to which parameter is the true data-generating model.

• Let $\ell : \Theta \to \mathbb{R}$ be the likelihood function, encoding the probability of observing the given data conditional on each parameter/model.

Bayesian inference says that, given this set-up, our updated posterior distribution over the parameters, denoted $p \in \Delta(\Theta)$, should be given by the conditional probability of the parameters given the data. That is, by Bayes’ rule, $$p(\theta) = \frac1Z \ell(\theta) \varphi(\theta),$$ where $Z \in \mathbb{R}$ is a normalisation constant.

§Evidence

The normalisation constant $Z$ can easily be expressed as the thing that makes $p$ integrate to 1, $$Z = \int_{\Theta} \ell(\theta) \varphi(\theta) \,d\theta.$$

This quantity is known to different communities by several different names.

• In probabilistic terms, $Z$ represents the marginal probability of the data under the given model class (with the given prior). Accordingly, it is called the marginal likelihood.

• In the context of model comparison, this marginal likelihood would be a key factor in selecting which of two separate model classes best explains the data. This earns $Z$ the alternative name of model evidence (or, simply evidence, which is the term I adopt in this note).

• The use of the symbol $Z$ derives from an analogous quantity in statistical mechanics, which in the original German was called the Zustandssumme (literally, “state sum”). In this analogy, the parameters are the potential states or configurations of the statistical problem.

While it’s easy to express the evidence as an integral, computing that integral for high-dimensional parameter spaces with non-trivial likelihood functions is often intractable.

This is the main barrier to using exact Bayesian inference in practice. Without normalisation, we can’t make any absolute statements about our posterior beliefs about particular parameters in our model class, or comparisons between the suitability of different model classes (we can only make comparisons between parameters within the one model class).

§Variational inference

That’s where variational inference comes in: it’s a method of constructing an approximate posterior that is as similar as possible to the Bayesian posterior, avoiding having to normalise the Bayesian posterior itself.

At a high level, the method works as follows:

1. We first fix some variational family, $\mathcal{Q}\subset \Delta(\Theta)$, a set of distributions over parameter space which we have efficient access to (for example, we should know how to compute their normalised densities, and how to sample from them).

2. Then, we somehow pick from this family a single distribution, $q_\star \in \mathcal{Q}$, of “best fit” to the posterior. That is, $q_\star$ should be maximally similar to the Bayesian posterior $p$ according to some measure of divergence between distributions.

This distribution $q_\star$ becomes our variational posterior and serves as our approximation for the Bayesian posterior $p$. The process of going from a prior to a variational posterior is called variational inference.

In detail, the standard measure of divergence used is the Kullback–Leibler (KL) divergence from the Bayesian posterior to the candidate variational posterior, $$\operatorname{KL}\!\left[\,q\,\middle\|\,p\,\right] = \mathbb{E}_{\theta \sim q}\!\left[\log\frac{q(\theta)}{p(\theta)}\right] .$$ This KL divergence is lower if the approximations are intuitively more similar, and zero if and only if they are equal.

Computing the KL divergence itself without knowing the Bayesian posterior $p$ would be challenging, but it turns out we can get away without it for the purposes of optimisation. Observe: $$\begin{align*} \operatorname{KL}\!\left[\,q\,\middle\|\,p\,\right] &= \mathbb{E}_{\theta \sim q}\!\left[\log\frac{q(\theta)}{p(\theta)}\right] \\ &= \mathbb{E}_{\theta \sim q}\!\left[ \log\frac{q(\theta) Z}{\ell(\theta) \varphi(\theta)} \right] \\ &= \mathbb{E}_{\theta \sim q}\!\left[ - \log \ell(\theta) \right] + \mathbb{E}_{\theta \sim q}\!\left[ \log \frac{q(\theta)}{\varphi(\theta)} \right] + \log Z. \end{align*}$$

Take a look at the three terms we ended up with:

1. The first term here is the expected negative log likelihood of the data under the candidate variational posterior $q$. This is easy enough to optimise using stochastic gradient descent as long as we can sample parameters from $q$ (some care may be needed to autodiff through the sampler if you use PyTorch).

2. The second term is just the KL divergence of the candidate variational posterior $q$ from the prior $\varphi$. This is likely easy enough to optimise assuming we can sample parameters from $q$. Or, if we choose $\varphi$ and $\mathcal{Q}$ carefully, we might even be able to write the expectation in closed form.

3. The third term $\log Z$ is the log evidence, which is hard to compute. But notice it doesn’t depend on the candidate variational posterior $q$. So, it’s actually trivial to optimise since the gradient is always zero.

Equivalently, the objective we end up minimising is $\operatorname{KL}\!\left[\,q\,\middle\|\,p\,\right]$ shifted down by the log evidence, $$\operatorname{KL}\!\left[\,q\,\middle\|\,p\,\right] - \log Z = \mathbb{E}_{\theta \sim q}\!\left[ - \log \ell(\theta) \right] + \operatorname{KL}\!\left[\,q\,\middle\|\,\varphi\,\right].$$ The right-hand side of the above equation is a tractable objective that can be optimised at scale using stochastic gradient descent.

§Evidence lower bound

A property of the KL divergence between any two distributions is that it is always non-negative. We can therefore rearrange the above decomposition of the KL divergence from the Bayesian posterior to produce the following inequality: $$\begin{align*} \log Z &= \operatorname{KL}\!\left[\,q\,\middle\|\,p\,\right] - \mathbb{E}_{\theta \sim q}\!\left[ - \log \ell(\theta) \right] - \operatorname{KL}\!\left[\,q\,\middle\|\,\varphi\,\right]. \\ &\geq \mathbb{E}_{\theta \sim q}\!\left[ \log \ell(\theta) \right] - \operatorname{KL}\!\left[\,q\,\middle\|\,\varphi\,\right]. \end{align*}$$

The final expression, $\mathbb{E}_{\theta \sim q}\!\left[ \log \ell(\theta) \right] - \operatorname{KL}\!\left[\,q\,\middle\|\,\varphi\,\right]$, is named the evidence lower bound (ELBO). Formally, define $\text{ELBO} : \mathcal{Q}\to \mathbb{R}$ such that for $q \in \mathcal{Q}$ we have $$\text{ELBO}(q) = \mathbb{E}_{\theta \sim q}\!\left[ \log \ell(\theta) \right] - \operatorname{KL}\!\left[\,q\,\middle\|\,\varphi\,\right].$$

To spell it out, we have, for all $q \in \mathcal{Q}$, $$\log Z \geq \text{ELBO}(q).$$ This explains the name “evidence lower bound,” except that we have produced a bound on the logarithm of the evidence (rather than the evidence itself).

The gap in the bound is the term we dropped, $\operatorname{KL}\!\left[\,q\,\middle\|\,p\,\right]$, which is smaller the closer we make $q$ to $p$, and zero if and only if $q = p$.

Notice that the bound holds regardless of which variational posterior distribution we plug in. This suggests the following methodology for approximating $\log Z$ by trying to make the bound as tight as possible.

1. First, fix some variational family, $\mathcal{Q}\subset \Delta(\Theta)$, a set of distributions over parameter space which we have efficient access to (for example, we should know how to compute their normalised densities, and how to sample from them).

2. Then, we somehow pick from this family a single distribution, $q_\star \in \mathcal{Q}$, giving the best bound. That is, $q_\star$ should be chosen to maximise the evidence lower bound $\text{ELBO}(\cdot)$.

We can apply the bound to reveal that $\log Z \geq \text{ELBO}(q_\star)$, this being the tightest possible bound available using our chosen variational family. Assuming we have a sufficiently expressive variational family, we could even use $\text{ELBO}(q_\star)$ as an approximation of $\log Z$ for purposes of model selection.

How do we maximise the ELBO over a variational family? Well, that’s just the same as minimising the negative ELBO over the variational family. That is, finding $q \in \mathcal{Q}$ so as to minimise $\mathbb{E}_{\theta \sim q}\!\left[- \log \ell(\theta) \right] + \operatorname{KL}\!\left[\,q\,\middle\|\,\varphi\,\right]$. Wait a second, that’s just the same as the variational inference objective. So, we’re done.

§Conclusion

In both variational inference and evidence lower bound maximisation, we discovered equivalent optimisation problems: optimising an expectation over a class of distributions. That we are optimising over a space of distributions is the “variational” part of the method.

In both cases, the quality of the approximation we get (how good is the variational posterior as an approximation of the Bayesian posterior, or how tight is the maximised evidence lower bound) comes down to how small we can get $\operatorname{KL}\!\left[\,q\,\middle\|\,p\,\right]$. In theory, this is just about having a sufficiently expressive variational family, so that somewhere there is a distribution that is close to $p$. In practice, if we make the variational family very complex, we also have to worry about whether we can find a good fit using practical optimisation methods such as stochastic gradient descent and friends.

If our parametric statistical model class is based on a neural network, the posteriors can be very thorny indeed, even in simple cases. These posteriors defy approximation by unimodal Gaussian posteriors that are commonly studied in the literature. It remains to be seen whether we can find variational families that can better approximate these neural posteriors while still being tractable to work with.