How does maximizing ELBO in Bayesian neural networks give us the correct posterior predictive distribution?

Question

In Bayesian/variational neural networks one often uses the Evidence Lower BOund (ELBO) as the objective function to optimize with respect to the model parameters. That is if $D=\{y_i,x_i\}_{1\dots n}$ is the set of our training data, $\{y,x\}$ is a new data point, and $w$ represents our network weights/parameters, then the posterior predictive distribution is: \begin{array}{rl} p(y|x,D) =& \int p(y|x,w,D)p(w|x,D)dw \\ =& \int p(y|x,w)p(w|D)dw \\ \end{array} We don't know $p(w|D)$ so we instead approximate it with some known parameteric function (always seemingly chosen to be a Gaussian density) which we call $q_\phi(w)$. We want $q_\phi(w)$ to be close to $p(w|D)$ so we minimize the KL divergence between them as part of our optimization procedure. That is we end up with $d_{KL}(q_\phi(w)||p(w|D))$ in our objective function, essentially, in addition to the usual likelihood component. Thus the ELBO is:

$$E_{q_\phi(w)}[log(p(D|w))] - d_{KL}(q_\phi(w)||p(w|D))$$

Here is my question: how does maximizing ELBO lead to a good/correct posterior predictive distribution $p(y|x,D)$? The first term $$E_{q_\phi(w)}[log(p(D|w))]$$ ensures that our model accurately models the mean of the posterior predictive distribution, while the second term ensures that the approximate posterior of the weight (note this is not the posterior of $y$!) distribution is close to its true distribution. Why do we care that the approximate posterior of the weight distribution is good? I don't see how using our objective function to approximate $p(w|D)$ well translates to $p(y|x,D)$ also being well approximated.

Answer 1

TL;DR to "how does maximizing ELBO lead to a good/correct posterior predictive distribution 𝑝(𝑦|𝑥,𝐷)?"

The initial formulation for making new predictions is correct, essentially you have a model $p(y|x,w)$ to predict the distribution of $y$ for a given new $x$, which requires parameters $w$. Which parameters to use? You can

• use a point estimate $w^*$ learned from the data $D$ (e.g. maximum likelihood estimate (MLE) or maximum a posteriori (MAP)). Predictions are given by $p(y|x,D) \approx p(y|x,w^*)$;
• or take a Bayesian approach by averaging over all possible values of $w$ weighted by $p(w|D)$, their probability given the observed data - a.k.a the posterior - learned from the data (e.g. by maximising the ELBO). Predictions are given by $p(y|x,D) = \int_w p(y|x,w)p(w|D)$.

Maximising the ELBO implicitly minimises $KL[q(w|D)\|p(w|D)]$ (see below). The ELBO is maximal when the KL term equals 0, meaning that the approximate posterior $q(w|D)$ equals the true posterior $p(w|D)$ (which is intractable to compute directly).

---

Longer answer: Your definition of the ELBO isn't quite right.

You want to approximate the posterior $p(w|D)$ with $q(w)$ (whether this is denoted $q(w|D)$ or $q(w)$ doesn't particularly matter, but the former makes clearer it approximates the posterior, not a prior over $w$, $p(w)$, so I use $q(w|D)$ below).

• To fit $q(w|D$ to $p(w|D)$, typically an intractable distribution, as you say, you maximise a lower bound (ELBO) on the log likelihood of the data, where $p(D) = \int_w p(D|w)p(w)$.

• The ELBO os derived:

  $\log p(D) = \int_w q(w|D)\log \tfrac{p(D,w)}{p(w|D)} \quad= \int_w q(w|D)\log \tfrac{p(D|w)p(w)}{q(w|D)}\tfrac{q(w|D)}{p(w|D)}$

  $= \int_w q(w|D)\log p(D|w) - \int_w q(w|D)\log \tfrac{q(w|D)}{p(w)} + \int_w q(w|D)\log \tfrac{q(w|D)}{p(w|D)}$

  $\geq \int_w q(w|D)\log p(D|w) - \int_w q(w|D)\log \tfrac{q(w|D)}{p(w)} $

  $= \mathbb{E}_{q(w|D)}[\log p(D|w)] - d_{KL}[q(w|D)\|p(w)]$

The difference to what you wrote is in the KL term. In particular the ELBO doesn't feature the true posterior $p(w|D)$ as you don't know it (if you did you wouldn't be trying to approximate it).

The dropped term leading to the inequality "gap" (leading to a lower bound), is the KL divergence between the approximate and true posterior, which is (implicitly) minimised as the ELBO is maximised. Subject to how well the various probabilities are modelled, the ELBO is fully maximised when the approximate posterior matches the true posterior

• (i.e. the KL term between the approximate and true posteriors goes to zero, hence the inequality "gap" disappears and you have an equality).

So maximising the ELBO is an elaborate indirect way of getting the approximate posterior $q(w|D)$ to fit the true posterior $p(w|D)$, which can then be used for predictions etc.

• It requires (i) an assumed prior over the parameters $p(w)$, (ii) the the data likelihood for a given parameter choice $p(D|w)$ to be computable, and (iii) a parametric family of distributions $q(w|D)$.
• Here, the likelihood $p(D|w) = \sum_{(x,y)\in D} p(y|x,w)$ is the same likelihood function as used for making predictions "at test time" (as you refer to at the outset).
• Given sufficient data, the posterior $q(w|D)$ may concentrate on a particular value $w^*$ (the maximum a posteriori estimate), allowing $w^*$ to be plugged directly into $p(y|x,w)$ for a good approximation to having to integrate over the whole posterior distribution.

Answer 2

With Bayesian methods, in general, you extract information from your data $D$ which, together with some prior belief, leads to a posterior for your parameters $w$. That means this posterior $p(w|D)$ is your best guess for a distribution over $w$.

And if you would now generate new data $(x, y)$ by sampling a $w$ from $q_\phi(w)$, and then sampling for a given $x$ a $y$ from $p(y|x, w)$ you would get pairs $(x, y)$ that, if all your prior belief is right, would belong to the same population as your previous data $D$.