Variational Inference

The posterior in Bayesian Models with latent variables is quite hard to compute if we do not choose a conjuagte prior for the likelihood function. However, we need to have some approximate version of the posterior in order to make inferences later on. One approach to approximate posterior is to find a distribution $q \in Q$ ( $Q$ is a family of distributions) that can minimize the KL divergence $\begin{array}{r} p (z | X) \approx min_{q \in Q} K L (q (z) ∥ p (z | X)) \end{array}$ where $z$ is the latent variable and $X$ is the data And we do not need to worry about the normalization constant $Z$ ( $p (X)$ ) $\begin{aligned} K L (q (z) ∥ p (z | X)) & = \int_{z} q (z) l o g \frac{q (z)}{p (X | z) p (z) / Z} \\ = \int_{z} q (z) l o g \frac{q (z)}{p (X | z) p (z)} + l o g (Z) \\ p (z | X) & \approx min_{q \in Q} K L (q (z) ∥ p (X | z) p (z)) (p (X | z) p (z) = p^{*} (z)) \end{aligned}$ since $\int_{z} q (z) d z = 1$ . We only need to work with likelihood and prior in this case.

E-step in EM

utilizes the variational inference technique. The E-step requires us to calculate the posterior of the latent variables, which can be hard to compute. VI can help us approximate that distribution assuming that we restrict ourselves to a family of distributions. $\begin{array}{r} q (z) = p (z | x, θ) \approx min_{q \in Q} K L q (z) p^{*} (z) \end{array}$ This is called Variational EM.

Mean Field Approximations

This is a Variational Inference method where we assume the distribution $q$ to be factorized over the latent variables across all dimensions $d$ , i.e., $\begin{array}{r} Q = {q | q (z) = \prod_{i = 1}^{d} q_{i} (z_{i})} \\ q (z) = {argmin}_{q \in Q} K L (\prod_{i = 1}^{d} q_{i} (z_{i}) ∥ p^{*} (z)) \end{array}$ and we minimize the KL divergence using coordinated gradient descent. First find the minima for $q_{1}$ keeping everything else fixed, then $q_{2}$ , and so on. We will repeat this loop until convergence.

Since we are minimizing over one component at a time, the functional form of any $q_{k} (z_{k})$ can be derived as follows $\begin{matrix} min_{q_{k}} K L ((\prod_{j = 1}^{d} q_{j}) (z_{j}) ∥ p^{*} (z)) = min_{q_{k}} \int (\prod_{j = 1}^{d} q_{j} (z_{j})) l o g \frac{\prod_{j = 1}^{d} q_{j} (z_{j})}{p^{*} (z)} d z \\ = min_{q_{k}} {\sum_{i = 1}^{d} \int (\prod_{j = 1}^{d} q_{i} (z_{i})) l o g (q_{i} (z_{i})) d z - \int (\prod_{j = 1}^{d} q_{j} (z_{j})) l o g (p^{*} (z)) d z} \\ = min_{q_{k}} {\int (q_{k} (z_{k}) l o g (q_{k} (z_{k})) (\int \prod_{j = 1, j \neq k}^{d} q_{j} (z_{j}) d z_{\neq j}) d z_{k}) \\ + \sum_{i = 1, i \neq k}^{d} \int (q_{i} (z_{i}) l o g (q_{i} (z_{i})) (\int \prod_{j = 1, j \neq i}^{d} q_{j} (z_{j}) d z_{\neq i}) d z_{i}) - \int (\prod_{j = 1}^{d} q_{j} (z_{j})) l o g (p^{*} (z)) d z} \\ Note that \int \prod_{j = 1}^{d} q_{j} (z_{j}) d z = \int (\prod_{j = 1}^{d}) q_{j} (z_{j}) d z_{1} d z_{2} \dots d z_{d} = \prod_{j = 1}^{d} (\int q_{j} (z_{j}) d z_{j}) = 1 \\ min_{q_{k}} {\int q_{k} (z_{k}) l o g (q_{k} (z_{k})) d z_{k} + \sum_{i = 1, i \neq k}^{d} q_{i} (z_{i}) l o g (q_{i} (z_{i})) d z_{i} \\ - \int q_{k} (z_{k}) ((\prod_{j = 1, j \neq k}^{d} q_{j} (z_{j})) l o g (p^{*} (z)) d z_{\neq j}) d z_{k}} \\ = min_{q_{k}} {\int q_{k} (z_{k}) [l o g (q_{k} (z_{k})) - ((\prod_{j = 1, j \neq k}^{d} q_{j} (z_{j})) l o g (p^{*} (z)) d z_{\neq j})] d z_{k}} \\ because \sum_{i = 1, i \neq k}^{d} q_{i} (z_{i}) l o g (q_{i} (z_{i})) d z_{i} is constant wrt q_{k} \end{matrix}$

We now make three observations

adding a constant to the minimization equation will still give the same result
$\int q_{k} (z_{k}) d z_{k} = 1$ because $q_{k}$ is a probability distribution
$\int (\prod_{j = 1, j \neq k}^{d} q_{j} (z_{j})) l o g (p^{*} (z)) d z_{\neq j} = E_{q_{- k}} [l o g (p^{*} (z))]$ since $\prod_{j = 1, j \neq k}^{d} q_{j} (z_{j})$ is a valid probability distribution. Notice that the expectation integrates over all $z$ except $z_{k}$ and thus is a function of just $z_{k}$ . $q_{- k}$ in the expectation just means that we are considering the product $\prod_{j = 1, j \neq k}^{d} q_{j} (z_{j})$ as the probability distribution. We convert this expectation to a positive value by taking the exponent, and then get a valid probability distribution as $\begin{array}{r} t (z_{k}) = \frac{e x p (E_{q_{- k}} [l o g (p^{*} (z))])}{\int e x p (E_{q_{- k}} [l o g (p^{*} (z))]) d z_{k}} \end{array}$ The denominator of this expression is a constant which we can introduce in the equation as is without any alteration

Rewriting the minimization equation so far $\begin{matrix} min_{q_{k}} {\int q_{k} (z_{k}) [l o g (q_{k} (z_{k})) - l o g (e x p (E_{q_{- k}} [l o g (p^{*} (z))]))] d z_{k} + (\int e x p (E_{q_{- k}} [l o g (p^{*} (z))]) d z_{k}) \int q_{k} (z_{k}) d z_{k}} \\ = min_{q_{k}} {\int q_{k} (z_{k}) l o g (q_{k} (z_{k})) - l o g \frac{e x p (E_{q_{- k}} [l o g (p^{*} (z))])}{\int e x p (E_{q_{- k}} [l o g (p^{*} (z))]) d z_{k}}} = min_{q_{k}} {\int q_{k} (z_{k}) l o g \frac{q_{k} (z_{k})}{t (z_{k})}} \\ min_{q_{k}} K L ((\prod_{j = 1}^{d} q_{j}) (z_{j}) ∥ p^{*} (z)) = min_{q_{k}} K L (q_{k} (z_{k}) ∥ t (z_{k})) \end{matrix}$ and we know that KL divergence is minimized when the two functions conincide, i.e. $\begin{array}{r} q_{k} (z_{k}) = t (z_{k}) = \frac{e x p (E_{q_{- k}} [l o g (p^{*} (z))])}{\int e x p (E_{q_{- k}} [l o g (p^{*} (z))]) d z_{k}} \\ ⟹ l o g (q_{k} (z_{k})) = E_{q_{- k}} [l o g (p^{*} (z))] - c o n s t a n t \end{array}$ which is the expectation of the posterior on $z$ over $q$ without the current component being minimized.