29
$\begingroup$

In a VAE, the encoder learns to output two vectors:

$$\mathbf{\mu} \in\ \mathbb{R}^{z}$$ $$\mathbf{\sigma} \in\ \mathbb{R}^{z}$$

which are the mean and variances for the latent vector $\mathbf{z}$, the latent vector $\mathbf{z}$ is then calculated by:

$$\mathbf{z} = \mu + \sigma \epsilon$$

where: $\epsilon = N(0, \mathbf{I}_{z \times z})$

The KL divergence loss for a VAE for a single sample is defined as (referenced from this implementation and this explanation):

$$\frac{1}{2} \left[ \left(\sum_{i=1}^{z}\mu_{i}^{2} + \sum_{i=1}^{z}\sigma_{i}^{2} \right) - \sum_{i=1}^{z} \left(log(\sigma_{i}^{2}) + 1 \right) \right]$$

Though, I'm not sure how they got their results, would anyone care to explain or point me to the right resources?

$\endgroup$

1 Answer 1

50
$\begingroup$

The encoder distribution is $q(z|x)=\mathcal{N}(z|\mu(x),\Sigma(x))$ where $\Sigma=\text{diag}(\sigma_1^2,\ldots,\sigma^2_n)$.

The latent prior is given by $p(z)=\mathcal{N}(0,I)$.

Both are multivariate Gaussians of dimension $n$, for which in general the KL divergence is: $$ \mathfrak{D}_\text{KL}[p_1\mid\mid p_2] = \frac{1}{2}\left[\log\frac{|\Sigma_2|}{|\Sigma_1|} - n + \text{tr} \{ \Sigma_2^{-1}\Sigma_1 \} + (\mu_2 - \mu_1)^T \Sigma_2^{-1}(\mu_2 - \mu_1)\right] $$ where $p_1 = \mathcal{N}(\mu_1,\Sigma_1)$ and $p_2 = \mathcal{N}(\mu_2,\Sigma_2)$.

In the VAE case, $p_1 = q(z|x)$ and $p_2=p(z)$, so $\mu_1=\mu$, $\Sigma_1 = \Sigma$, $\mu_2=\vec{0}$, $\Sigma_2=I$. Thus: \begin{align} \mathfrak{D}_\text{KL}[q(z|x)\mid\mid p(z)] &= \frac{1}{2}\left[\log\frac{|\Sigma_2|}{|\Sigma_1|} - n + \text{tr} \{ \Sigma_2^{-1}\Sigma_1 \} + (\mu_2 - \mu_1)^T \Sigma_2^{-1}(\mu_2 - \mu_1)\right]\\ &= \frac{1}{2}\left[\log\frac{|I|}{|\Sigma|} - n + \text{tr} \{ I^{-1}\Sigma \} + (\vec{0} - \mu)^T I^{-1}(\vec{0} - \mu)\right]\\ &= \frac{1}{2}\left[-\log{|\Sigma|} - n + \text{tr} \{ \Sigma \} + \mu^T \mu\right]\\ &= \frac{1}{2}\left[-\log\prod_i\sigma_i^2 - n + \sum_i\sigma_i^2 + \sum_i\mu^2_i\right]\\ &= \frac{1}{2}\left[-\sum_i\log\sigma_i^2 - n + \sum_i\sigma_i^2 + \sum_i\mu^2_i\right]\\ &= \frac{1}{2}\left[-\sum_i\left(\log\sigma_i^2 + 1\right) + \sum_i\sigma_i^2 + \sum_i\mu^2_i\right]\\ \end{align}

$\endgroup$
7
  • 1
    $\begingroup$ Can you comment on why this looks so different from the univariate case? (stats.stackexchange.com/q/7440/141343 bottom of question) $\endgroup$
    – BlackBear
    Sep 5, 2019 at 16:48
  • $\begingroup$ @BlackBear your univariate example doesn't assume $\sigma_2 = 1$, which is the only difference $\endgroup$
    – Firebug
    Sep 5, 2019 at 16:50
  • 1
    $\begingroup$ @BlackBear as Firebug notes, $p$ is assumed to be multivariate standard normal (most of the time) in VAEs. If you set $n=1$ here and $\sigma_2=1,\,\mu_2=0$ there, they should match. Also, their answer starts from the definition of the KL divergence as an expectation of the log difference, whereas I started from the general expression for two Gaussians for simplicity. $\endgroup$ Sep 5, 2019 at 16:56
  • $\begingroup$ What is $n$ in the equation? Why is it replaced with 1? $\endgroup$ Mar 15, 2021 at 10:22
  • $\begingroup$ @GergőHorváth, $n$ is the dimension of the vector z. It is not replaced with 1, it is replaced with $\sum_{i=1}^n 1$. $\endgroup$
    – toliveira
    Apr 16, 2021 at 23:42

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.