All the Ways to Carve Up the ELBO

Mar 10 2020

Despite the title, this is my by-no-means exhaustive collection of known ways to decompose the ELBO, focusing on the role of the aggregate KL regularizer, its relation to the aggregate posterior, and the mutual information between the data and the latent variable.


We assume the data distribution has a density, \(p^*(x)\), and we would like to estimate it by that of a latent variable model \(p(x) := \int p(x|z) p(z) dz\), via maximum likelihood. Everything works the same in the discrete setting by replacing the word "density" with "PMF", and integrals with sums. In reality, the data distribution may not have a probability density, as is likely the case for real-world data such as natural images; then maximum likelihood density estimation runs into severe trouble, as I explain here. However, let's just stick with the standard textbook setting, and suppose the data distribution really does have a density \(p^*(x)\) 1.

The ELBO (evidence lower bound) is a tractable lower bound on the model log density \(p(x)\):

$$ELBO := \mathbb{E}_{q(z|x)}[ \log p(x,z) - \log q(z|x) ] = \log p(x) - KL(q(z|x)\|p(z|x)) \leq \log p(x)$$

where \(q(z|x)\) is any distribution 2 over the latent space, and \(p(z|x) \propto p(z,x)\) is the Bayesian posterior of \(z\) given \(x\) under our latent variable model. Note that all the distributions without the * superscript (i.e., everything except the true data density \(p^*(x)\)) are ones we can potentially control and parameterize.

For maximum-likelihood training, ideally we would like to maximize the population log-likelihood \(\mathbb{E}_{p^*(x)}[ \log p(x)]\), or its sample approximation based on a training set. Since \(\log p(x)\) is intractable, we replace it with the ELBO, resulting in the aggregate ELBO objective,

$$\mathbb{E}_{p^*(x)} [\mathbb{E}_{q(z|x)}[ \log p(x,z) - \log q(z|x) ] ],$$

which can be decomposed into an aggregate conditional log-likelihood, minus an aggregate ELBO regularizer:

$$\mathbb{E}_{p^*(x)} [\mathbb{E}_{q(z|x)}[ \log p(x|z)] - \mathbb{E}_{p^*(x)} [KL(q(z|x)\| p(z)) ].$$

Of particular interest to me is the aggregate ELBO regularizer. Many of the decompositions of the aggregate ELBO are simply based on rewriting the aggregate regularizer term. Moreover, the aggregate ELBO regularizer also arises outside the context of latent variable modeling and density estimation, and has a nice connection to information theory.

Decomposing the Aggregate ELBO Regularizer

Besides the standard interpretation of the aggregate ELBO regularizer as penalizing the complexity of the approximate posterior \(q(z|x)\) by pulling it closer to the prior \(p(z)\), here are some other ways to decompose and interpret it:

1. As the cross entropy between the aggregate posterior and prior, plus the negative average entropy of approximate posteriors:

$$\mathbb{E}_{p^*(x)} [KL(q(z|x)\| p(z)) ] = CE(q(z)\| p(z)) + \mathbb{E}_{p^*(x)} [ - \mathbb{H}[q(z|x)]].$$

The aggregate posterior is defined as \(q(z) = \mathbb{E}_{p^*(x)}[ q(z|x) ]\), and the cross-entropy \(CE(q(z)\| p(z)) = \mathbb{E}_{q(z)} [-\log p(z)]\). This form is easily obtained by writing out \(KL(q(z|x)\| p(z)) = \mathbb{E}_{q(z|x)}[- \log p(z) + \log q(z|x) ]\) and simplifying on the LHS.

This form reveals that the optimal prior, in the sense of minimizing the aggregate KL regularizer (and hence the aggregate ELBO), should be equal to the aggregate posterior, i.e., \(p(z) = q(z)\), which is when the cross-entropy term is minimized. See the ELBO surgery paper and the VampPrior paper, which popularized this decomposition.

2. As the KL divergence between the aggregate posterior and prior, plus the mutual information of \((x, z) \sim p^*(x) q(z|x)\):

$$\mathbb{E}_{p^*(x)} [KL(q(z|x)\| p(z)) ] = KL(q(z)||p(z)) + I_{p^*(x) q(z|x)}[X; Z],$$

where the mutual information is equal to

$$I_{p^*(x) q(z|x)}[X; Z] := \int \int p^*(x) q(z|x) \log \frac{q(z|x)} {q(z)} dz dx.$$

This form can be derived from the first decomposition ("cross-entropy plus negative posterior entropy"), by adding and subtracting \(\mathbb{E}_{q(z)}[\log q(z)] = \mathbb{E}_{p^*(x)q(z|x)}[\log q(z)]\) from the two terms. As in the first decomposition, the prior \(p(z)\) only appears in one term -- its divergence to the aggregate posterior, which immediately reveals that the optimal \(p(z) = q(z)\). This decomposition also nicely reveals the minimum value of the ELBO regularizer (when \(p(z) = q(z)\)) as the mutual information.

Looked at another way, given any joint distribution \(p^*(x) q(z|x)\), the aggregate ELBO regularizer is a variational upper bound on its mutual information, with \(p(z)\) being any arbitrary variational distribution,

$$I_{p^*(x) q(z|x)}[X; Z] = \min_{p(z)} \mathbb{E}_{p^*(x)} [KL(q(z|x)\| p(z)) ].$$

This is quite a useful upper bound on mutual information, a quantity that often appears in machine learning but is typically challenging to compute; see Poole et. al 2019. This upper bound also plays a key role in the Blahut-Arimoto algorithm, a well-established algorithm from information theory for computing rate-distortion functions and channel capacity (see Section 19.8 of the Elements of Information Theory textbook).

3. As the KL divergence between joint distributions \(p^*(x) q(z|x)\) and \(p^*(x) p(z)\):

$$\mathbb{E}_{p^*(x)} [KL(q(z|x)\| p(z)) ] = \mathbb{E}_{p^*(x) q(z|x)} [ \log \frac{q(z|x)}{p(z)} + \log \frac{p^*(x)}{p^*(x)} ] = KL(p^*(x) q(z|x) \| p^*(x) p(z)).$$

This is an intuitive restatement of effect of the aggregate ELBO regularizer in pulling the approximate posteriors towards the prior, in an "aggregate" sense. This form in terms of a KL divergence between two distributions defined on the product space \(\mathcal{X} \times \mathcal{Z}\) has some technical advantages, especially in the convergence proof of the Blahut-Arimoto algorithm.

4 (bonus content). This a decomposition of the whole aggregate ELBO, rather than just the aggregate ELBO regularizer term, but I'm including it given its popularity in the literature:

$$\mathbb{E}_{p^*(x)} [\mathbb{E}_{q(z|x)}[ \log p(x,z) - \log q(z|x) ] ] = - KL(q(z|x)p^*(x) \| p(x,z)) - \mathbb{H}[p^*]$$

which is very much analogous to how the (true) maximum likelihood objective decomposes into the negative KL between the data and the model distributions, minus the data entropy:

$$\mathbb{E}_{p^*(x)} [ \log p(x) ] = - KL(p^*(x) \| p(x)) - \mathbb{H}[p^*]$$

The connection follows quickly from the fact that \(ELBO = \log p(x) - KL(q(z|x)\|p(z|x))\), and the chain rule of KL divergence,

$$KL(q(z|x)p^*(x) \| p(x,z)) = KL(p^*(x) \| p(x)) + \mathbb{E}_{p^*(x)} [ KL(q(z|x)\|p(z|x)) ].$$

This form is discussed at length in, e.g., Structured Disentangled Representations (note that in their definition, \(p^*(x)\) is the empirical distribution defined by training samples, rather the true data distribution).

  1. It's possible to relax this assumption if we consider a more general notion of ``density'' than the standard Lebesgue density on \(\mathbb{R}^n\). i.e., if we know in advance the "manifold" on which the data is concentrated, then the data distribution has a density function (Radon–Nikodym derivative) w.r.t. the reference measure of that manifold, which we can still model/learn with a latent variable model. 

  2. Throughout this note I use the words "distribution" and "density" interchangeably.