KL-Divergence of $Q(z|X)$ and $P(z)$ in Variational Autoencoder (VAE)

I aim to understand how $$D_{KL}[Q(z | X) || P(z)]$$ can be converted to $$\frac{1}{2} \sum_{k} (\Sigma(X) + \mu^{2}(X) - 1 - \log \Sigma(X))$$, where $$k$$ is the dimension of the Gaussian distribution.

We have known that $$Q(z | X)$$ distributes as $$N(\mu(X), \Sigma(X))$$ and $$P(z)$$ satisfies $$N(0, 1)$$. Therefore, I can understand

$$D_{KL}[Q(z | X) || P(z)] = D_{KL}[N(\mu(X), \Sigma(X)) || N(0, 1) = \frac{1}{2} (\operatorname{tr}(\Sigma(X)) + \mu(X)^{T}\mu(X) - k - \log \det(\Sigma(X))).$$

We implement $$\Sigma(X)$$ as a vector as it's a diagonal matrix.

However, I don't understand the following, $$D_{KL}[N(\mu(X), \Sigma(X)) || N(0, 1) = \frac{1}{2} \big(\Sigma_{k} \Sigma(X) + \Sigma_{k} \mu^{2}(X) - \Sigma_{k}1 - \log \prod_{k} \Sigma(X) \big).$$

Can anyone please enlighten me how is the last step deduced?

$$D_{KL} \big[ Q(Z \mid X) \parallel P(z) \big] = \frac{1}{2} \Big[ \text{Tr} \big( \Sigma(X) \big) + \mu(X)^T \mu(X) -k - \log \det \big( \Sigma(X) \big) \Big]$$

Then, convert everything to scalar notation. Let $$\Sigma_j(X)$$ be the $$j$$th entry along the diagonal of the covariance matrix, and let $$\mu_j(X)$$ be the $$j$$th entry of the mean vector. Note that: 1) the trace of a matrix is the sum of entries along the diagonal, 2) the inner product of a vector with itself is the sum of the squared entries, and 3) the determinant of a diagonal matrix is the product of entries along the diagonal. This gives:

$$\frac{1}{2} \Big[ \sum_{j=1}^k \Sigma_j(X) + \sum_{j=1}^k \big( \mu_j(X) \big)^2 -k - \log \prod_{j=1}^k \Sigma_j(X) \Big]$$

Now let's make some substitutions. For the third term, note that $$k = \sum_{j=1}^k 1$$. For the fourth term, the log of a product is equal to a sum of logs. This gives:

$$\frac{1}{2} \Big[ \sum_{j=1}^k \Sigma_j(X) + \sum_{j=1}^k \big( \mu_j(X) \big)^2 - \sum_{j=1}^k 1 - \sum_{j=1}^k \log \Sigma_j(X) \Big]$$

Finally, combine the sums to obtain your original expression:

$$\frac{1}{2} \sum_{j=1}^k \Big[ \Sigma_j(X) + \big( \mu_j(X) \big)^2 - 1 - \log \Sigma_j(X) \Big]$$

• So $$\frac{1}{2} \big(\Sigma_{k} \Sigma(X) + \Sigma_{k} \mu^{2}(X) - \Sigma_{k}1 - \log \prod_{k} \Sigma(X) \big)$$ is an abbrevation of $$\frac{1}{2} \big(\Sigma_{k} \Sigma_k(X)+ \Sigma_{k} \mu_k^{2}(X) - \Sigma_{k}1 - \log \prod_{k} \Sigma_k(X) \big)$$? Nov 22 '18 at 3:24