Combination of variational methods and empirical Bayes

Question

Suppose I have a posterior $p(z, \theta | y, \eta)$ with $y$ observed data, $z$ are hidden variables and $\theta$ are parameters, and $\eta$ is a vector of hyperparameters. I construct a mean field approximation to the posterior using coordinate ascent, i.e. $$ q(z) \gets \exp\left\{E_q(\log p(z, \theta | y, \eta) | z) + \mbox{const}\right\} \\ q(\theta) \gets \exp\left\{E_q(\log p(z, \theta | y, \eta) | \theta) + \mbox{const}\right\} $$ where $\mbox{const}$ makes each integrate to $1$. Iterate until convergence. My question is, if I want to do empirical Bayes, and (approximately) profile out $\eta$, is it valid to just augment this with an EM step $$ \eta \gets \arg \max_\eta E_q \log(p(z, \theta | y, \eta))? $$ Based on the evidence lower bound $$ \log p(y | \eta) \ge E_q \log(p(z, \theta | y, \eta)) - E_q \log(q(z, \theta)) $$ it seems like I should be able to get away with this; optimizing the lower bound over $\eta$ is exactly the proposed EM step, so I'm still doing coordinate ascent on the lower bound.

I suspect that, if this works, it is the obvious thing to do for people who do variational inference regularly. Since variational methods are fast, my thought is that maybe I could do this to set hyperparameters before launching into a Gibbs sampler for exact (up to MCMC error) inference.

Update: I've toyed around with this a bit and found that adding the EM step can work okay, but in some situations it seems to increase sensitivity to bad local optima in the variational algorithm. Initializations of variational parameters that normally seem to work okay don't. If I instead only throw in the EM step every few iterations it works better, but then of course I slow down convergence. One approach that works okay is to "burn in" the variational algorithm using a method that I trust, and then add the EM step every iteration thereafter. I'm adding a bounty because I'm still looking for some good general advice.

Answer 1

One way of deciding how to run variational MLE is to look at how the experts do it.

In Blei's LDA code (http://www.cs.princeton.edu/~blei/lda-c/lda-c-dist.tgz), within the "run_em" function, the "lda_inference" function (inside "doc_e_step") repeatedly maximizes with respect to each $q$ distribution until convergence. After the $q$'s converge, the algorithm maximizes with respect to the parameters in "lda_mle".

The justification for this order is that by maximizing with respect to the $q$'s until convergence you get a better estimate of the expectations of hidden variables (or marginalized parameters) needed to maximize with respect to the parameters.

In standard EM, of course, the expectations you are computing are exact - which is the main difference between standard and variational EM - so this is not a concern.

From the perspective of EM as a maximization algorithm over the function $F(q,\theta)$ (www.cs.toronto.edu/~radford/ftp/emk.pdf) or from the perspective of maximizing the evidence lower bound, it is not clear that maximizing over the q's until convergence is best in terms of computationally efficiency because the algorithm will reach a local maximum no matter the order of maximization steps.

Answer 2

Generally, in empirical Bayes, you maximise the marginal likelihood (also called model evidence, or the normalising constant of the posterior) with respect to the hyperparameters and plug this estimate of the hyperparameters into the posterior.

In Casella (2001) there is a derivation of an EM for empirical Bayes. Casella first writes the marginal likelihood $m(\eta ; y)$ as:

\begin{align} m(\eta ; y) = \frac{p(y, \theta, z | \eta) }{p(\theta, z| y, \eta)}. \end{align}

So the marginal likelihood is the joint distribution of model parameters $\theta$, latent parameters $z$ and data $y$ ($p(y,\theta,z| \eta)$) divided by the joint posterior distribution of the model and latent parameters given the data ($p(\theta,z| y, \eta$).

Take the logarithm and expectation with respect to the posterior $p(\theta, z| y, \eta^{(0)})$, for some starting value $\eta^{(0)}$:

$$ E [\log m(\eta ; y)| \eta^{(0)}] = E [\log p(y, \theta, z | \eta)|\eta^{(0)}] - E [\log p(\theta, z| y, \eta) | \eta^{(0)}]. $$

By Gibbs' inequality, $E [\log p(\theta, z| y, \eta) | \eta^{(0)}] \leq E [\log p(\theta, z| y, \eta^{(0)}) | \eta^{(0)}]$. So, maximising $E [\log p(\theta, z| y, \eta) | \eta^{(0)}]$ with respect to $\eta$ beyond $\eta^{(0)}$ increases the marginal likelihood such that the sequence

$$ \eta^{(k+1)} = \underset{\eta}{\arg \max} \, E [\log p(\theta, z| y, \eta) | \eta^{(k)}] $$

converges to a local maximum of the marginal likelihood. In general, this expectation is not available in closed form. In Casella (2001), the expectation is approximated using $M$ Monte Carlo samples from the posterior $p(\theta, z| y, \eta^{(k)})$:

\begin{align} E [\log p(\theta, z| y, \eta) | \eta^{(k)}] & \approx M^{-1} \sum_{m=1}^M \log p(y, \theta^{(m)}, z^{(m)} | \eta), \\ \theta^{(m)}, z^{(m)} & \sim p(\theta, z| y, \eta^{(k)}). \end{align}

However, you could of course use other approximating methods, such as variational Bayes. The expectation then becomes:

$$ E [\log p(\theta, z| y, \eta) | \eta^{(k)}] \approx E_Q [\log p(\theta, z| y, \eta) | \eta^{(k)}], $$

where the right-hand side expectation is now with respect to the variational posterior.

So, in the EM algorithm to find the hyperparameters that maximise the marginal likelihood, the E step is now a variational Bayes EM itself, so we have an EM within an EM. Which is basically what you describe.

That also explains why it doesn't work so well when you update the hyperparameters at every iteration: the expectation with respect to the variational Bayes posterior is not very accurate after one iteration. Updating the hyperparameters after the variational Bayes algorithm has converged gives better performance, since the expectation is more accurate.