Variational Bayes: Understanding Mean field approximation

Question

I am looking at the mean field approximation as used in Variational Bayes inference and I looked at this section on wikipedia with the factorised approximation as described here: https://en.wikipedia.org/wiki/Variational_Bayesian_methods#Factorized_approximation

I can understand the approximation we make which is that the parameters are independent a-posteriori but what I do not understand is the update that follows. For example, the derivation of $\ln q(\mu)$ that follows takes the expectation wrt to the other variables. This is stated as a fact here (https://en.wikipedia.org/wiki/Variational_Bayesian_methods#Mathematical_derivation_of_the_mean-field_approximation). Where the "In practice" section says that using calculus of variations we can show that. I could not derive it myself and neither could I find any proof for it in any of the books I own.

Can someone explain me why the log posterior in this factorized approximation is given as an expectation of the log joint model with the expectation taken wrt to the other parameters?

Answer 1

Taking the equation from Wikipedia
$D_{KL}(Q||P) = \sum_\limits{z}Q(Z)\log\frac{Q(Z)}{P(Z,X)} +\log P(X)$

What we want is to minimize KL distance wrt $Q$ distribution.

Since $P(X)$ is independent of $Q$ we need to care only about the first term.

Substituting the factored approximation, $Q=\prod\limits_{i=1}^M q(Z_i|X)$

$$\sum_\limits{z}Q(Z)\log\frac{Q(Z)}{P(Z,X)} = \sum_\limits{z}Q(Z)\log\{ \prod\limits_{i=1}^M q(Z_i|X)\}-\sum_\limits{z}Q(Z)\log P(Z,X)\\$$ Considering only terms dependent only on a particular factor $q(z_j)$

$$= \sum\limits_{z_j}q(z_j|X)\log q(z_j|X)-\sum_\limits{z_j}q(z_j|X) \left\lbrace \sum\limits_{z_{i\neq j}}\left(\prod\limits_{i\neq j}q(z_i|X)\right)\log P(Z,X)\right\rbrace+ \text{const}$$ In the above equation the terms inside the curly braces is an expectation of the function $\log P(Z,X)$. Where $\mathbb{E}_{i\neq j}$ denotes expectation under $Q(Z)$ distribution wrt all variables except $z_j$. This is where the expectation wrt other variables of log joint model comes into equation. Defining $ \log\tilde{P}(Z_j,X) \triangleq \mathbb{E}_{i\neq j}\left[\log P(Z,X)\right]$ (Note that once you do the expectation it is a function of $z_j$).

$$\begin{equation}= \sum\limits_{z_j}q(z_j|X)\log q(z_j|X)-\sum_\limits{z_j}q(z_j|X)\mathbb{E}_{i\neq j}\left[\log P(Z,X)\right]+ \text{const} \end{equation}$$

$$\begin{equation} =\sum\limits_{z_j}q(z_j|X)\left\lbrace \log q(z_j|X)-\log\tilde{P}(Z_j,X) \right\rbrace+\text{const}\\ =\sum\limits_{z_j}q(z_j|X)\log \left(\frac{q(z_j|X)}{\tilde P(Zj,X)}\right)+\text{const} \end{equation}$$

Now if we make $\tilde{P}(Z_j,X)$ a probability mass function by including a normalization constant, (such that $\tilde{P}(Z_j,X)$ sums to 1 over all possible $Z_j$), this normalizing constant can be absorbed in to the constant term. That will make this again a KL distance between $q(Z_j)$ and the the new distribution which is proportional to $\tilde{P}(Z_j,X)$. If the two distributions are equal, KL distance becomes zero and gives the minimum of our original objective.

Therefore, in order to minimize $D_{KL}(Q||P)$ wrt each $q(z_j)$ we want

$ \begin{equation} q(z_j|X) \propto \tilde{P}(Z_j,X) \end{equation}$

Since $\tilde{P}(Z_j,X) = \exp\{\mathbb{E}_{i\neq j}\left[\log P(Z,X)\right]\}$

$q(Z_j|X) = \dfrac{\exp\{\mathbb{E}_{i\neq j}\left[\log P(Z,X)\right]\}}{\sum\limits_{Z_j}\exp\{\mathbb{E}_{i\neq j}\left[\log P(Z,X)\right]\}}$