Variational inference: how to rewrite ELBO?

Question

I am reading this paper on variational inference and this website.

One thing I am confused about is how they get to decompose ELBO, where $ELBO(q) = E_q[log~p(z,x)] - E_q[log~q(z)]$, when focusing on one latent variable's variational distribution $q_j$ like this:

$$ ELBO(q_j) = E_j[E_{-j}[log~p(z_j, z_{-j}, x)]] - E_j[log~q_j(z_j)] + C $$

They say that they use iterated expectation but I had a hard time decomposing $ELBO(q_j)$ using that ($E[X] = E[E[X|Y]]$).

Can anyone elaborate on this? Thanks!

UPDATE: $q(z)=\prod_{i} q(z_i)$ is an assumption so I understand the decomposition of the 2nd term.

Answer 1

Your update has stated that you are using the mean-field variational family, or in other words that $q(z)=\prod_{i} q(z_i)$ which means that $$ \log q(z) = \sum_i \log q(z_i) \tag{*}. $$

So \begin{align*} \text{ELBO}(q) &= E_q[\log p(z,x)] - E_q[\log q(z)] \\ &= E_q[\log p(z_j, z_{-j},x)] - E_q[\log q(z_j, z_{-j})] \\ &= E_q[\log p(z_j, z_{-j},x)] - E_q[\log q(z_j) + \sum_{i\neq j}\log q(z_i)] \tag{*}\\ &= E_q[E\left(\log p(z_j, z_{-j},x) \mid z_{j} \right)] - E_q[\log q(z_j)] - E\left[\sum_{i\neq j}\log q(z_i)\right].\\ \end{align*} This is equivalent to equation (19) in your first linked document.

Answer 2

I read those two materials again and think that the decomposition may be done via the following rearrangement:

\begin{align} ELBO & = E_q[log~p(x,z)] - E_q[log~q(z)] \\ & = E_q[log~p(x, z_j, z_{-j})] - \sum_{q_l}E_{q_l}[q_l(z_l)] \\ & [\text{Here we use the fact that } q(z) \text{ can be factorized}]\\ & = E_j\Big[E_{-j}\big[ log~p(x, z_j, z_{-j}) \vert z_{j} \big] \Big] - E_{q_j}[q_j] + const \\ \end{align}

Now, according to the definition of expectation, we have: \begin{align} E_{-j}\big[ log~p(x, z_j, z_{-j}) \vert z_{j} \big] &= \int_{-j} log~p(x, z_j, z_{-j})~q(z_{-j}|z_j) dq_{-j} \\ & = \int_{-j} log~p(x, z_j, z_{-j})~q(z_{-j}) dq_{-j} \\ & = E_{-j}\big[ log~p(x, z_j, z_{-j}) \big] \end{align} [We assume independence between latent variables' variational distributions $q(z)$]

Therefore we have:

$$ELBO = E_{-j}\big[ log~p(x, z_j, z_{-j}) \big] - E_{q_j}[q_j] + const $$