Skip to main content
3 of 3
Add the expectation of KL divergence between 2 normal distributions
Cuong
  • 451
  • 4
  • 14

You are getting closed. To ease the derivation, let's redefine some notations:

\begin{cases} p(\pmb{\mu}, \pmb{\Lambda}) = p(\pmb{\mu} \vert \pmb{\Lambda}) p(\pmb{\Lambda}) & = \mathcal{N}\left( \pmb{\mu} \vert \pmb{\mu}_{p}, (\lambda_{p} \pmb{\Lambda})^{-1} \right) \mathcal{W} \left( \pmb{\Lambda} \vert \mathbf{W}_{p}, \nu_{p} \right)\\ q(\pmb{\mu}, \pmb{\Lambda}) = q(\pmb{\mu} \vert \pmb{\Lambda}) q(\pmb{\Lambda}) & = \mathcal{N}\left( \pmb{\mu} \vert \pmb{\mu}_{q}, (\lambda_{q} \pmb{\Lambda})^{-1} \right) \mathcal{W} \left( \pmb{\Lambda} \vert \mathbf{W}_{q}, \nu_{q} \right) \end{cases}

The KL divergence of interest is:

\begin{align} & D_{\mathrm{KL}} \left[ p(\pmb{\mu}, \pmb{\Lambda}) \Vert q(\pmb{\mu}, \pmb{\Lambda}) \right] \\ & = \int_{\mu} \int_{\Lambda} p(\pmb{\mu}, \pmb{\Lambda}) \ln \frac{p(\pmb{\mu}, \pmb{\Lambda})}{q(\pmb{\mu}, \pmb{\Lambda})} d\pmb{\mu} d\pmb{\Lambda} \\ & = \int_{\mu} \int_{\Lambda} p(\pmb{\mu} \vert \pmb{\Lambda}) p(\pmb{\Lambda}) \ln \frac{p(\pmb{\mu} \vert \pmb{\Lambda}) p(\pmb{\Lambda})}{q(\pmb{\mu} \vert \pmb{\Lambda}) q(\pmb{\Lambda})} d\pmb{\mu} d\pmb{\Lambda} \\ & = \int_{\mu} \int_{\Lambda} p(\pmb{\mu} \vert \pmb{\Lambda}) p(\pmb{\Lambda}) \ln \frac{p(\pmb{\mu} \vert \pmb{\Lambda})}{q(\pmb{\mu} \vert \pmb{\Lambda})} d\pmb{\mu} d\pmb{\Lambda} + \int_{\mu} \int_{\Lambda} p(\pmb{\mu} \vert \pmb{\Lambda}) p(\pmb{\Lambda}) \ln \frac{p(\pmb{\Lambda})}{q(\pmb{\Lambda})} d\pmb{\mu} d\pmb{\Lambda}\\ & = \int_{\Lambda} p(\pmb{\Lambda}) \left[\int_{\mu} p(\pmb{\mu} \vert \pmb{\Lambda}) \ln \frac{p(\pmb{\mu} \vert \pmb{\Lambda})}{q(\pmb{\mu} \vert \pmb{\Lambda})} d\pmb{\mu} \right] d\pmb{\Lambda} + \int_{\Lambda} p(\pmb{\Lambda}) \ln \frac{p(\pmb{\Lambda})}{q(\pmb{\Lambda})} d\pmb{\Lambda}\\ & = \mathbb{E}_{p(\pmb{\Lambda})} \left[ D_{\mathrm{KL}} \left[ p(\pmb{\mu} \vert \pmb{\Lambda}) \Vert q(\pmb{\mu} \vert \pmb{\Lambda}) \right] \right] + D_{\mathrm{KL}} \left[ p(\pmb{\Lambda}) \Vert q(\pmb{\Lambda}) \right].\\ & \tag{eq:KL_normal_wishart} \label{eq:KL_normal_wishart} \end{align}

To what I am aware of, there is a closed-form for the KL divergence between two Wishart distributions, corresponding to the second term. However, the first term is complicated, and I believe that further assumptions (eg. diagonal normal distributions) should be made to have a closed-form solution.

The first term is an expectation of the KL divergence between two normal distributions w.r.t. $p(\pmb{\Lambda})$, and also has a closed-form solution. To be specific, the KL divergence between 2 normal distributions can be written as:

\begin{aligned}[b] & D_{\mathrm{KL}} \left[ p(\pmb{\mu} \vert \pmb{\Lambda}) \Vert q(\pmb{\mu} \vert \pmb{\Lambda}) \right] \\ & = \frac{1}{2} \left[ \mathrm{tr}\left( \lambda_{q} \pmb{\Lambda} \lambda_{p}^{-1} \pmb{\Lambda}^{-1} \right) + \left( \pmb{\mu}_{q} - \pmb{\mu}_{p} \right)^{\top} \lambda_{q} \pmb{\Lambda} \left( \pmb{\mu}_{q} - \pmb{\mu}_{p} \right) - D + \ln \frac{\mathrm{det}(\lambda_{p} \pmb{\Lambda})}{\mathrm{det}(\lambda_{q} \pmb{\Lambda})}\right] \\ & = \frac{1}{2} \left[ D \frac{\lambda_{q}}{\lambda_{p}} + \left( \pmb{\mu}_{q} - \pmb{\mu}_{p} \right)^{\top} \lambda_{q} \pmb{\Lambda} \left( \pmb{\mu}_{q} - \pmb{\mu}_{p} \right) - D + D \ln \frac{\lambda_{p}}{\lambda_{q}} \right]. \end{aligned}

Note that for a Wishart distribution: $\mathbb{E}_{p(\pmb{\Lambda})} \left[ \pmb{\Lambda} \right] = \nu_{p} \mathbf{W}_{p}$. Hence, the first term of $\eqref{eq:KL_normal_wishart}$ can be obtained as:

\begin{aligned}[b] &\mathbb{E}_{p(\pmb{\Lambda})} \left[ D_{\mathrm{KL}} \left[ p(\pmb{\mu} \vert \pmb{\Lambda}) \Vert q(\pmb{\mu} \vert \pmb{\Lambda}) \right] \right] \\ & = \frac{\lambda_{q}}{2} \left( \pmb{\mu}_{q} - \pmb{\mu}_{p} \right)^{\top} \nu_{p} \mathbf{W}_{p} \left( \pmb{\mu}_{q} - \pmb{\mu}_{p} \right) + \frac{D}{2} \left( \frac{\lambda_{q}}{\lambda_{p}} - \ln \frac{\lambda_{q}}{\lambda_{p}} - 1 \right). \end{aligned}

Cuong
  • 451
  • 4
  • 14