Derivation of Variational Inference

Question

I'm reading Blei et al. (2017) "Variational Inference: A Review for Statisticians" to understand Variational Inference (VI). I follow the paper's notations: $\mathbf{x}_{1:n}$ (observations), $\mathbf{z}_{1:m}$ (latent variables), and $q(\mathbf{z})$ (a variational factors) $$ q(\mathbf{z}) = \prod_{j=1}^{m} q_j(z_j).$$

On page 10, the paper explains how to get optimal $q_j(z_j)$. The log-likelihood of the observations is $$ \log p(\mathbf{x}) = {\rm KL}( q(\mathbf{z}) || p(\mathbf{z} | \mathbf{x} ) ) + {\rm ELBO}(q) . $$ We can decompose ELBO as $$ {\rm ELBO}(q) = \mathbb{E}[ \log p(\mathbf{z}, \mathbf{x})] - \mathbb{E}[\log q(\mathbf{z})]. $$

Using a factorization assumption, we can focus on the $j$th variational factor $q_j(z_j)$ $$ {\rm ELBO}(q_j) = \mathbb{E}_j \left[ \mathbb{E}_{-j} [\log p(z_j, \mathbf{z}_{-j}, \mathbf{x})] \right] - \mathbb{E}_j \left[ \log q_j (z_j) \right] + const \ \ ({\rm Eq.A}) $$

From here, the paper claims we can get Eq.(18) (p.9) $$ q_j^{*}(z_j) \propto \exp \left\{ \mathbb{E}_{-j} \left[ \log p(z_j, \mathbf{z}_{-j}, \mathbf{x}) \right] \right\}. $$

Since we want to maximize the ELBO with respect to $q_j$, we rewrite Eq.A as $$ \left[ \int q(z_j) \left( \int q(\mathbf{z}_{-j}) \cdot \log p(z_j, \mathbf{z}_{-j}, \mathbf{x}) d\mathbf{z}_{-j} \right) d z_j \right] - \left[ \int q(z_j) \log q(z_j) d z_j \right] + const \ \ \ ({\rm Eq.B}) $$ and get $$ \frac{\rm Eq.B}{\partial q(z_j)} = \left[ \int \int q(\mathbf{z}_{-j}) \log p(z_j, \mathbf{z}_{-j}, \mathbf{x}) d \mathbf{z}_{-j} dz_j \right] - \int \log q(z_j) dz_j + const \ \ ({\rm Eq.C}) $$

If we can say $$ {\rm Eq.C} = \mathbb{E}_{\mathbf{z}_{-j}} \left[ \log p (z_j, \mathbf{z}_{-j}, \mathbf{x}) \right] - \log q(z_j) + const, \ \ ({\rm Eq.D)} $$ we can set ${\rm Eq.C}= 0$ and get Eq.(18) in the original paper.

However, I cannot derive it. For the first term, $p(z_j, \mathbf{z}_{-j}, \mathbf{x})$ has $\mathbf{z}_{-j}$ and $z_j$, and for the second term $q(z_j)$ has $z_j$. Don't we have to consider them when we calculate integral (expectation)?

Note: I think my question is similar to this post, but I have a trouble with how to get Eq.(18) in the original paper while the linked post focuses on how to get Eq.A.

Answer 1

We have $$\textsf{ELBO}(q_j)=\mathbb{E}_{q_j}[\mathbb{E}_{q_{-j}}[\log(p(z_j,z_{-j},x))]]-\mathbb{E}_{q_j}[\log(q_j(z_j))]-constant,$$ which we can rewrite as $$\textsf{ELBO}(q_j)=-\mathbb{E}_{q_j}\left[\log\left(\frac{q_j(z_j)}{\exp\left[\mathbb{E}_{q_{-j}}[\log(p(z_j,z_{-j},x))]\right]}\right)\right]-constant,$$ which we recognize as a KL divergence (up to a constant) $$\textsf{ELBO}(q_j)=-D_{KL}\left(q_j(z_j)||\exp\left[\mathbb{E}_{q_{-j}}[\log(p(z_j,z_{-j},x))]\right]\right)-constant.$$

Since we'd like to maximize the $\textsf{ELBO}$, we'd like to minimize the KL divergence. This happens when we let $q_j(z_j)\propto\exp\left[\mathbb{E}_{q_{-j}}[\log(p(z_j,z_{-j},x))]\right]$.

The reason we only specify this up to a constant of proportionality is because we were being a little sloppy before! That technically wasn't a KL divergence since $\exp\left[\mathbb{E}_{q_{-j}}[\log(p(z_j,z_{-j},x))]\right]$ wasn't normalized (we can just add in the normalizing constant since it won't depend on $z_j$ and thus we'll still have the same optimization problem, so this isn't too important).