ELBO

Evidence lower bound, variational free energy etc

October 2, 2020 — August 3, 2022

approximation
optimization
probability
statistics

\[ \renewcommand{\Ex}{\mathbb{E}} \renewcommand{\dd}{\mathrm{d}} \renewcommand{\kl}{\operatorname{KL}} \renewcommand{\H}{\mathbb{H}} \renewcommand{\vv}[1]{\boldsymbol{#1}} \renewcommand{\pd}{\partial} \]

Figure 1

On using the most convenient probability metric (i.e. KL divergence) to do variational inference.

There is nothing novel here. But everyone who is doing variational inference has to work through this just once, and I’m doing so here.

Yuge Shi’s introduction is the best short intro that gets to state-of-the-art. The canonical intro is Matthews (2017) who did a thesis on it. Murphy (2012) sec 21.2 is also pretty good.

We often want a variational approximation to the marginal (log-)likelihood \(\log p_{\theta}(\vv{x})\) (a.k.a. “evidence”) for some probabilistic model with observations \(\vv{x},\) unobserved latent factors \(\vv{x}\), model parameters \(\mathbb{\theta}\) and variational parameters \(\phi\).

Here’s one. The steps are all elementary, although realizing you would want to take those steps is not, IMO.

For convenience, we assume everything has a density with respect to some unspecifed dominating measure over \(\vv{x}\) and \(\vv{z}\), which is usually an OK assumption.1

\[\begin{aligned} \log p_{\theta}(\vv{x}) &=\log (p_{\theta}(\vv{x})) \Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[ 1 \right] \\ &=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[ \log (p_{\theta}(\vv{x})) \right]\\ &=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[ \log \left( \frac{ p_{\theta}(\vv{x},\vv{z}) }{ p_{\theta}(\vv{z}|\vv{x}) } \right) \right]\\ &=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[ \log\left( \frac{ q_{\phi}(\vv{z}|\vv{x})p_{\theta}(\vv{x},\vv{z}) }{ p_{\theta}(\vv{z}|\vv{x})q_{\phi}(\vv{x}|\vv{x}) }\right) \right]\\ &=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[ \log \left(\frac{ q_{\phi}(\vv{z}|\vv{x}) }{ p_{\theta}(\vv{z}|\vv{x}) } + \frac{ p_{\theta}(\vv{x},\vv{z}) }{ q_{\phi}(\vv{z}|\vv{x}) }\right) \right]\\ &=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[\log\left(\frac{ q_{\phi}(\vv{z}|\vv{x}) }{ p_{\theta}(\vv{z}|\vv{x}) }\right) \right] \\ &\qquad+\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[ \log\left(\frac{ p_{\theta}(\vv{x},\vv{z}) }{ q_{\phi}(\vv{z}|\vv{x}) }\right) \right] \nonumber\\ &=\kl_{\vv{z}}( q_{\phi}(\vv{z}|\vv{x}) \| p_{\theta}(\vv{z}|\vv{x}) ) +\mathcal{L}(\vv{x},\theta,\phi) \end{aligned}\]

\[ \mathcal{L}(\vv{x},\theta,\phi):=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[\log\left(\frac{ p_{\theta}(\vv{x},\vv{z}) }{ q_{\phi}(\vv{z}|\vv{x})}\right) \right] \] is called the free energy or Evidence Lower Bound.

We can understand this as a decomposition of the total marginal evidence into two parts which are sort-of-interpretable. \(\kl_{\vv{z}}( q_{\phi}(\vv{z}|\vv{x}) \| p_{\theta}(\vv{z}|\vv{x}))\) represents the cost of approximating the exact \(\theta\) distribution over latents with the \(\phi\) distribution. We can’t actually evaluate this in general but for fancy enough \(q_{\phi}\) it could be small.

The other bit, \(\mathcal{L}(\vv{x},\theta,\phi)\), represents an objective we can actually maximise, which motivates a whole bunch of technology. We can stop there and think of it as just that, or we can break it up further. The traditional next step if we want to decompose further is to observe that

\[\begin{aligned} \mathcal{L}(\vv{x},\theta,\phi) &=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[\log\left(\frac{ p_{\theta}(\vv{x},\vv{z}) }{ q_{\phi}(\vv{z}|\vv{x}) } \right) \right]\\ &=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[\log\left(\frac{ p_{\theta}(\vv{x}|\vv{z})p_{\theta}(\vv{z}) }{ q_{\phi}(\vv{z}|\vv{x}) } \right) \right]\\ &=\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[\log(p_{\theta}(\vv{x}|\vv{z})) \right] -\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[\log\left(\frac{ q_{\phi}(\vv{z}|\vv{x}) }{ p_{\theta}(\vv{z}) } \right) \right]\\ &=\underbrace{\Ex_{ q_{\phi}(\vv{z}\sim\vv{z}|\vv{x})}\left[\log(p_{\theta}(\vv{x}|\vv{z})) \right]}_{\text{Expected log likelihood}} - \underbrace{\kl_{\vv{z}}( q_{\phi}(\vv{z}|\vv{x}) \| p_{\theta}(\vv{z}))}_{\text{KL of approx posterior update}} \end{aligned}\]

This suggests we might intuit maximising the ELBO as maximising the data-conditional likelihood of the latents whilst penalising them for diverging too far from the prior on those latents. I find a marginal prior on the latents to be a weird concept here and this particular formulation makes my head hurt.

Yuge Shi, summarising Hoffman and Johnson (2016) and Mathieu et al. (2019) observes that we can break this down in a way which makes the per-observation latents into a coding problem. Suppose we index our \(N\) observations by \(n\), and they are independent. Then we can write this bad boy using the marginal \(q(\vv{z})\).

\[\begin{aligned} \mathcal{L}(\vv{x},\theta,\phi) &= \underbrace{\left[ \frac{1}{N} \sum^N_{n=1} \Ex_{q(z_n\mid x_n)} [\log p(x_n \mid z_n)] \right]}_{\color{#4059AD}{\text{(1) Average reconstruction}}} - \underbrace{(\log N - \Ex_{q(z)}[\H[q(x_n\mid z)]])}_{\color{#EE6C4D}{\text{(2) Index-code mutual info}}} \notag \\ &\quad + \underbrace{\kl(q(z)\mid p(z))}_{\color{#86CD82}{\text{(3) KL between q and p}}} \notag \\ \end{aligned}\] \(\H\left[p(\vv{z}) \right] \triangleq-\Ex_{p(\vv{z})} \left[\log p(\vv{z})\right]\) is entropy.

Another way to rewrite the ELBO is as \[\mathcal{L}(\vv{x},\theta,\phi)=\Ex_{q_{\phi}(\vv{z} \mid \vv{x})}\left[\log p_{\theta}(\vv{z}, \vv{x})\right]+\H\left[q_{\phi}(\vv{z} \mid \vv{x})\right].\] The log joint \(\log p_{\theta}(z, x)\) what physics people call the “negative energy”. This version highlights that a good posterior approximation \(q_{\phi}(z \mid x)\) must assign most of its probability mass to regions of low energy (i.e. high joint probability). At the same time the entropy term in the ELBO prevents \(q_{\phi}(z \mid x)\) from collapsing to an atom, unlike in, say, a MAP estimate.

Next thing, Importance weighted sampling in variational inference. Also a recommendation from Yuge Shi, see Adam Kosiorek’s What is wrong with VAEs? which finds autoencoders via importance sampling.

1 Bethe free energy

Everything so far as been about Helmholtz Free Energy. In graphical models graphs we are concerned with a related, more general (?) free energy called Bethe free energy (J. S. Yedidia, Freeman, and Weiss 2005; Jonathan S. Yedidia, Freeman, and Weiss 2000; M. J. Wainwright and Jordan 2008).

2 Incoming

3 References

Blei, Kucukelbir, and McAuliffe. 2017. Variational Inference: A Review for Statisticians.” Journal of the American Statistical Association.
Burda, Grosse, and Salakhutdinov. 2016. Importance Weighted Autoencoders.” In arXiv:1509.00519 [Cs, Stat].
Cremer, Morris, and Duvenaud. 2017. Reinterpreting Importance-Weighted Autoencoders.” In ICLR 2017.
Gelman, Carlin, Stern, et al. 2013. Bayesian Data Analysis. Chapman & Hall/CRC texts in statistical science.
Hoffman, and Johnson. 2016. ELBO Surgery: Yet Another Way to Carve up the Variational Evidence Lower Bound.” In Advances In Neural Information Processing Systems.
Huggins, Kasprzak, Campbell, et al. 2019. Practical Posterior Error Bounds from Variational Objectives.” arXiv:1910.04102 [Cs, Math, Stat].
Jordan. 1999. Learning in Graphical Models.
Kingma, Diederik P. 2017. Variational Inference & Deep Learning: A New Synthesis.”
Kingma, Diederik P., and Welling. 2019. An Introduction to Variational Autoencoders. Foundations and Trends in Machine Learning.
Li, and Turner. 2016. Rényi Divergence Variational Inference.” In Advances in Neural Information Processing Systems.
MacKay. 2002. Information Theory, Inference & Learning Algorithms.
Mathieu, Rainforth, Siddharth, et al. 2019. Disentangling Disentanglement in Variational Autoencoders.” In International Conference on Machine Learning.
Matthews. 2017. Scalable Gaussian Process Inference Using Variational Methods.”
Murphy. 2012. Machine learning: a probabilistic perspective. Adaptive computation and machine learning series.
Rainforth, Kosiorek, Le, et al. 2019. Tighter Variational Bounds Are Not Necessarily Better.” arXiv:1802.04537 [Cs, Stat].
Riegler, Kirkelund, Manchón, et al. 2012. Merging Belief Propagation and the Mean Field Approximation: A Free Energy Approach.” arXiv:1112.0467 [Cs, Math, Stat].
Roeder, Wu, and Duvenaud. 2017. Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference.” arXiv:1703.09194 [Cs, Stat].
Tucker, Lawson, Gu, et al. 2018. Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives.” arXiv:1810.04152 [Cs, Stat].
Wainwright, Martin, and Jordan. 2005. “A Variational Principle for Graphical Models.” In New Directions in Statistical Signal Processing.
Wainwright, Martin J., and Jordan. 2008. Graphical Models, Exponential Families, and Variational Inference. Foundations and Trends® in Machine Learning.
Yedidia, Jonathan S., Freeman, and Weiss. 2000. “Generalized Belief Propagation.” In Proceedings of the 13th International Conference on Neural Information Processing Systems. NIPS’00.
Yedidia, J.S., Freeman, and Weiss. 2003. Understanding Belief Propagation and Its Generalizations.” In Exploring Artificial Intelligence in the New Millennium.
———. 2005. Constructing Free-Energy Approximations and Generalized Belief Propagation Algorithms.” IEEE Transactions on Information Theory.

Footnotes

  1. Although in machine learning we tend to assume the dominating measure is Lebesgue, which as Matthews (2017) shows, can get you in to trouble.↩︎