Answers.Body: <p>This is the analytical form of the KL divergence between two multivariate Gaussian densities with diagonal covariance matrices (i.e. we assume independence). More precisely, it's the KL divergence between the <strong>variational distribution</strong></p>
  <p><span class="math-container">$$
  q_{\boldsymbol{\phi}}(\mathbf{z}) 
  = 
  \mathcal{N}\left(\mathbf{z} ; \boldsymbol{\mu}, \mathbf{\Sigma} = \boldsymbol{\sigma}^{2}\mathbf{I}\right)
  =
  \frac{\exp \left(-\frac{1}{2}\left(\mathbf{z} - \boldsymbol{\mu}\right)^{\mathrm{T}} \mathbf{\Sigma}^{-1}\left(\mathbf{z}-\boldsymbol{\mu} \right)\right)}{\sqrt{(2 \pi)^{J}\left|\mathbf{\Sigma}\right|}}
  \tag{1}\label{1}
  $$</span></p>
  <p>and the prior (it's the same as above, but with mean and covariance equal to the zero vector and the identity matrix, respectively)</p>
  <p><span class="math-container">$$
  p(\mathbf{z})=\mathcal{N}(\mathbf{z} ;\boldsymbol{0}, \mathbf{I})
  =
  \frac{\exp \left(-\frac{1}{2}\mathbf{z}^{\mathrm{T}}\mathbf{z}\right)}{\sqrt{(2 \pi)^{J}}}
  \tag{2}\label{2}
  $$</span></p>
  <p>where</p>
  <ul>
  <li><span class="math-container">$\boldsymbol{\mu} \in \mathbb{R}^J$</span> is the mean vector (we assume column vectors, so <span class="math-container">$\boldsymbol{\mu}^T$</span> would be a row vector)</li>
  <li><span class="math-container">$\mathbf{\Sigma} = \boldsymbol{\sigma}^{2}\mathbf{I} \in \mathbb{R}^{J \times J}$</span> is a diagonal covariance matrix (with the vector <span class="math-container">$\boldsymbol{\sigma}^{2}$</span> on the diagonal of the identity)</li>
  <li><span class="math-container">$\mathbf{z} \in \mathbb{R}^J$</span> is a sample (latent vector) from these Gaussians with dimensionality <span class="math-container">$J$</span> (or, at the same time, the input variable of the density)</li>
  <li><span class="math-container">$\left|\mathbf{\Sigma}\right| = \operatorname{det} \mathbf{\Sigma}$</span> is the determinant (so a number) of the diagonal covariance matrix, which is just the product of the diagonal elements for a diagonal matrix (which is the case); so, in the case of the identity, the determinant is <span class="math-container">$1$</span></li>
  <li><span class="math-container">$\boldsymbol{0} \in \mathbb{R}^J$</span> is the zero vector</li>
  <li><span class="math-container">$\mathbf{I} \in \mathbb{R}^{J \times J}$</span> is an identity matrix</li>
  <li><span class="math-container">$\mathbf{z}^{\mathrm{T}}\mathbf{z} = \sum_{i=1}^J z_i^2 \in \mathbb{R}$</span> is the dot product (hence a number)</li>
  </ul>
  <p>Now, the (negative of the) KL divergence is defined as follows</p>
  <p><span class="math-container">\begin{align}
  -D_{K L}\left(q_{\boldsymbol{\phi}}(\mathbf{z}) \| p(\mathbf{z})\right)
  &amp;=
  \int q_{\boldsymbol{\phi}}(\mathbf{z})\left(\log p(\mathbf{z})-\log q_{\boldsymbol{\phi}}(\mathbf{z})\right) d \mathbf{z} \\
  &amp;=
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} \left[ \log p(\mathbf{z})-\log q_{\boldsymbol{\phi}}(\mathbf{z})\right] \label{3}\tag{3}
  \end{align}</span></p>
  <p>Given that we have logarithms here, let's compute the logarithm of equations \ref{1} and \ref{2}</p>
  <p><span class="math-container">\begin{align}
  \log 
  \left(
  \mathcal{N}\left(\mathbf{z} ; \boldsymbol{\mu}, \mathbf{\Sigma} \right) 
  \right)
  &amp;= 
  \dots \\
  &amp;=
  -\frac{1}{2}(\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}} \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})-\frac{J}{2} \log (2 \pi)-\frac{1}{2} \log |\mathbf{\Sigma} |
  \end{align}</span></p>
  <p>and</p>
  <p><span class="math-container">\begin{align}
  \log 
  \left( 
  \mathcal{N}(\mathbf{z} ;\boldsymbol{0}, \mathbf{I})
  \right)
  &amp;= 
  \dots \\
  &amp;=
  -\frac{1}{2}\mathbf{z}^{\mathrm{T}} \mathbf{z}-\frac{J}{2} \log (2 \pi)
  \end{align}</span></p>
  <p>We can now replace these in equation \ref{3} (below, I have already performed some simplifications, to remove verbosity, but you can check them!)</p>
  <p><span class="math-container">\begin{align}
  \frac{1}{2}
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} \left[ 
  -\mathbf{z}^{\mathrm{T}} \mathbf{z}
  + (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}} \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})
  +  \log |\mathbf{\Sigma} |
  \right] 
  \tag{4}\label{4}
  \end{align}</span>
  Now, given that <span class="math-container">$\mathbf{\Sigma}$</span> is diagonal and the log of a product is just a sum of the logarithms, we have <span class="math-container">$\log |\mathbf{\Sigma} | = \sum_{i=1}^J \log \sigma_{ii}$</span>, so we can continue</p>
  <p><span class="math-container">\begin{align}
  \frac{1}{2}
  \left(
  -
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} 
  \left[ 
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  + 
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}} \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})
  \right] 
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right) 
  &amp;= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} 
  \left[ 
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  + 
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  \operatorname{tr} \left(
  \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}}
  \right)
  \right] 
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right) 
  &amp;= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} 
  \left[ 
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  + 
  \operatorname{tr} \left(
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}}
  \right] 
  \right)
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right) 
  &amp;= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} 
  \left[ 
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  + 
  \operatorname{tr} \left(
  \mathbf{\Sigma}^{-1}
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  (\mathbf{z}-\boldsymbol{\mu})
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}}
  \right] 
  \right)
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right) 
  &amp;= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} 
  \left[ 
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  + 
  \operatorname{tr} \left(
  \mathbf{\Sigma}^{-1}
  \mathbf{\Sigma}
  \right)
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right)
  &amp;= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} 
  \left[ 
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  + 
  J
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right)
  &amp;= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} 
  \left[ 
  \operatorname{tr}
   \left(
  \mathbf{z} \mathbf{z}^{\mathrm{T}}
  \right)
  \right]
  + 
  \sum_{i=1}^J 1
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right)
  &amp;= \\
  \frac{1}{2}
  \left(
  -
  \operatorname{tr}
  \left(
  \mathbf{\Sigma}\right) 
  - 
  \operatorname{tr}
  \left(
  \boldsymbol{\mu} \boldsymbol{\mu}^T
  \right)
  + 
  \sum_{i=1}^J 1
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right)
  &amp;= \\
  \frac{1}{2}
  \left(
  -
  \sum_{i=1}^J \sigma_{ii}
  - 
  \sum_{i=1}^J \mu_{i}^2
  + 
  \sum_{i=1}^J 1
  + 
  \sum_{i=1}^J \log \sigma_{ii}
  \right)
  &amp;= \\
  \frac{1}{2}
  \sum_{i=1}^J
  \left(
  1
  + 
  \log \sigma_{ii}
  -
  \sigma_{ii}
  - 
  \mu_{i}^2
  \right)
  \end{align}</span></p>
  <p>In the above simplifications, I also applied the following rules.</p>
  <ul>
  <li>the linearity of the expectations,</li>
  <li><span class="math-container">$\operatorname{tr}(ba^T) = a^Tb$</span> (where <span class="math-container">$\operatorname{tr}$</span> is the trace of a matrix)</li>
  <li><a href="https://math.stackexchange.com/q/3268268/168764">you can swap the trace and the expectation</a></li>
  <li><a href="https://en.wikipedia.org/wiki/Covariance_matrix" rel="noreferrer"><span class="math-container">$\mathbf{\Sigma} = 
  \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  (\mathbf{z}-\boldsymbol{\mu})
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}}
  \right] $</span></a> (this is just the definition of the <em>covariance matrix</em>)</li>
  <li><span class="math-container">$\operatorname{tr}\left( \mathbf{\Sigma}^{-1}
  \mathbf{\Sigma} \right) = \operatorname{tr}\left(\mathbf{I} \right) = J = \sum_{i=1}^J 1$</span></li>
  <li><a href="https://www.probabilitycourse.com/chapter6/6_1_5_random_vectors.php" rel="noreferrer"><span class="math-container">$ \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z})} 
  \left[ 
  \mathbf{z} \mathbf{z}^{\mathrm{T}}
  \right] = \mathbf{\Sigma} + \boldsymbol{\mu} \boldsymbol{\mu}^T$</span></a> (definition of the <em>correlation matrix</em>)</li>
  </ul>
  <p>The official PyTorch implementation of the VAE, which can be found <a href="https://github.com/pytorch/examples/blob/master/vae/main.py" rel="noreferrer">here</a>, also uses this formula. This formula can also be found in <a href="https://arxiv.org/pdf/1312.6114.pdf#page=10" rel="noreferrer">Appendix B of the VAE paper</a>, but the long proof that I've just written above is not given. Note that, in my proof above, <span class="math-container">$\sigma$</span> is the variance and is denoted by <span class="math-container">$\sigma^2$</span> in the paper (as it is usually the case to denote the variance as the square of the standard deviation <span class="math-container">$\sigma$</span>, but, again, in my proof above <span class="math-container">$\sigma$</span> is the variance).</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2021-02-16T12:56:21.020Z
Answers.LastActivityDate: 2021-02-16T14:11:11.690Z
Answers.LastEditDate: 2021-02-16T14:11:11.690Z
Answers.UserDisplayName: nbro
Answers.Body: <p>The code is correct. Since OP asked for a proof, one follows.</p>
  <p>The usage in the code is straightforward if you observe that the authors are using the symbols unconventionally: <code>sigma</code> is the <em>natural logarithm of the variance,</em> where usually a normal distribution is characterized in terms of a mean <span class="math-container">$\mu$</span> and variance. Some of the functions in OP's link even have arguments named <code>log_var</code>.<span class="math-container">$^*$</span></p>
  <p>If you're not sure how to derive the standard expression for KL Divergence in this case, you can start from the definition of KL divergence and crank through the arithmetic. In this case, <span class="math-container">$p$</span> is the normal distribution given by the encoder and <span class="math-container">$q$</span> is the standard normal distribution.
  <span class="math-container">$$\begin{align}
  D_\text{KL}(P \| Q) &amp;= \int_{-\infty}^{\infty} p(x) \log\left(\frac{p(x)}{q(x)}\right) dx \\
  &amp;= \int_{-\infty}^{\infty} p(x) \log(p(x)) dx - \int_{-\infty}^{\infty} p(x) \log(q(x)) dx
  \end{align}$$</span>
  The first integral is recognizable as almost definition of entropy of a Gaussian (up to a change of sign).
  <span class="math-container">$$
  \int_{-\infty}^{\infty} p(x) \log(p(x)) dx = -\frac{1}{2}\left(1 + \log(2\pi\sigma_1^2) \right)
  $$</span>
  The second one is more involved.
  <span class="math-container">$$
  \begin{align}
  -\int_{-\infty}^{\infty} p(x) \log(q(x)) dx 
  &amp;= \frac{1}{2}\log(2\pi\sigma_2^2) - \int p(x) \left(-\frac{\left(x - \mu_2\right)^2}{2 \sigma_2^2}\right)dx \\
  &amp;= \frac{1}{2}\log(2\pi\sigma_2^2) + \frac{\mathbb{E}_{x\sim p}[x^2] - 2 \mathbb{E}_{x\sim p}[x]\mu_2 +\mu_2^2} {2\sigma_2^2} \\
  &amp;= \frac{1}{2}\log(2\pi\sigma_2^2) + \frac{\sigma_1^2 + \mu_1^2-2\mu_1\mu_2+\mu_2^2}{2\sigma_2^2} \\
  &amp;= \frac{1}{2}\log(2\pi\sigma_2^2) + \frac{\sigma_1^2 + (\mu_1 - \mu_2)^2}{2\sigma_2^2}
  \end{align}
  $$</span>
  The key is recognizing this gives us a sum of several integrals, and each can apply the law of the unconscious statistician. Then we use the fact that <span class="math-container">$\text{Var}(x)=\mathbb{E}[x^2]-\mathbb{E}[x]^2$</span>. The rest is just rearranging.</p>
  <p>Putting it all together:
  <span class="math-container">$$
  \begin{align}
  D_\text{KL}(P \| Q) &amp;= -\frac{1}{2}\left(1 + \log(2\pi\sigma_1^2) \right) + \frac{1}{2}\log(2\pi\sigma_2^2) + \frac{\sigma_1^2 + (\mu_1 - \mu_2)^2}{2\sigma_2^2} \\
  &amp;= \log (\sigma_2) - \log(\sigma_1) + \frac{\sigma_1^2 + (\mu_1 - \mu_2)^2}{2\sigma_2^2} - \frac{1}{2}
  \end{align}
  $$</span></p>
  <p>In this special case, we know that <span class="math-container">$q$</span> is a standard normal, so
  <span class="math-container">$$
  \begin{align}
  D_\text{KL}(P \| Q) &amp;= -\log \sigma_1 + \frac{1}{2}\left(\sigma_1^2 + 
  \mu_1^2 - 1 \right) \\
  &amp;= - \frac{1}{2}\left(1 + 2\log \sigma_1- \mu_1^2 -\sigma_1^2   \right)
  \end{align}
  $$</span>
  In the case that we have a <span class="math-container">$k$</span>-variate normal with diagonal covariance for <span class="math-container">$p$</span>, and a multivariate normal with covariance <span class="math-container">$I$</span>, this is the sum of <span class="math-container">$k$</span> univariate normal distributions because in this case the distributions are independent.</p>
  <p>The code is a correct implementation of this expression because <span class="math-container">$\log(\sigma_1^2) = 2 \log(\sigma_1)$</span> and in the code, <code>sigma</code> is the logarithm of the variance.</p>
  <hr />
  <p><span class="math-container">$^*$</span>The reason that it's convenient to work on the scale of the <em>log-variance</em> is that the log-variance can be any real number, but the <em>variance</em> is constrained to be non-negative by definition. It's easier to perform optimization on the unconstrained scale than it is to work on the constrained scale in <span class="math-container">$\eta^2$</span>. Also, we want to avoid &quot;round-tripping,&quot; where we compute <span class="math-container">$\exp(y)$</span> in one step and then <span class="math-container">$\log(\exp(y))$</span> in a later step, because this incurs a loss of precision. In any case, autograd takes care of all of the messy details with adjustments to gradients resulting from moving from one scale to another.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2021-02-17T05:12:39.420Z
Answers.Text: The part I'm a little confused by is that how do you know the mean of p(x)? Especially in neural network case? Cause it's intractable, or does reducing it down to a multivariate Gauss make it tractable? But if so is the mean and SD based on what the data produces?
Answers.UserDisplayName: user8714896
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2021-02-17T05:15:18.800Z
Answers.Text: The latent representation of a VAE emits a mean and (a transformation of) the variance for each input: the $i$th sample has an associated distribution with parameters $(\mu_i, \sigma_i^2)$. In other words, there's a distribution $p_i(x)$ for each of the $i$ inputs.
Answers.UserDisplayName: Sycorax
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2021-02-16T04:17:46.020Z
Answers.LastActivityDate: 2021-02-18T03:37:57.480Z
Answers.LastEditDate: 2021-02-18T03:37:57.480Z
Answers.UserDisplayName: Sycorax
Body: <p>I found the following PyTorch code (from <a href="https://debuggercafe.com/getting-started-with-variational-autoencoder-using-pytorch/" rel="nofollow noreferrer">this link</a>)</p>
  <pre><code>-0.5 * torch.sum(1 + sigma - mu.pow(2) - sigma.exp())
  </code></pre>
  <p>where <code>mu</code> is the mean parameter that comes out of the model and <code>sigma</code> is the sigma parameter out of the encoder. This expression is apparently equivalent to the KL divergence. But I don't see how this calculates the KL divergence for the latent.</p>
  
Tags: <pytorch><proofs><implementation><variational-autoencoder><kl-divergence>
Title: How is this Pytorch expression equivalent to the KL divergence?
UserDisplayName: user8714896
ThreadText: # How is this Pytorch expression equivalent to the KL divergence?
  
  (Asked by: username_0 on 2021-02-13)
  
  I found the following PyTorch code (from [this link](https://debuggercafe.com/getting-started-with-variational-autoencoder-using-pytorch/))
  
  ```
  -0.5 * torch.sum(1 + sigma - mu.pow(2) - sigma.exp())
  
  ```
  
  where `mu` is the mean parameter that comes out of the model and `sigma` is the sigma parameter out of the encoder. This expression is apparently equivalent to the KL divergence. But I don't see how this calculates the KL divergence for the latent.
  
  **Tags:** pytorch, proofs, implementation, variational-autoencoder, kl-divergence
  
  
  ---
  
  ## 2 answers
  
  ### Answer by username_2 on 2021-02-16. Score: 6 | Accepted answer
  
  The code is correct. Since OP asked for a proof, one follows.
  
  The usage in the code is straightforward if you observe that the authors are using the symbols unconventionally: `sigma` is the *natural logarithm of the variance,* where usually a normal distribution is characterized in terms of a mean $\mu$ and variance. Some of the functions in OP's link even have arguments named `log_var`.$^\*$
  
  If you're not sure how to derive the standard expression for KL Divergence in this case, you can start from the definition of KL divergence and crank through the arithmetic. In this case, $p$ is the normal distribution given by the encoder and $q$ is the standard normal distribution.
  $$\begin{align}
  D\_\text{KL}(P \| Q) &= \int\_{-\infty}^{\infty} p(x) \log\left(\frac{p(x)}{q(x)}\right) dx \\
  &= \int\_{-\infty}^{\infty} p(x) \log(p(x)) dx - \int\_{-\infty}^{\infty} p(x) \log(q(x)) dx
  \end{align}$$
  The first integral is recognizable as almost definition of entropy of a Gaussian (up to a change of sign).
  $$
  \int\_{-\infty}^{\infty} p(x) \log(p(x)) dx = -\frac{1}{2}\left(1 + \log(2\pi\sigma\_1^2) \right)
  $$
  The second one is more involved.
  $$
  \begin{align}
  -\int\_{-\infty}^{\infty} p(x) \log(q(x)) dx
  &= \frac{1}{2}\log(2\pi\sigma\_2^2) - \int p(x) \left(-\frac{\left(x - \mu\_2\right)^2}{2 \sigma\_2^2}\right)dx \\
  &= \frac{1}{2}\log(2\pi\sigma\_2^2) + \frac{\mathbb{E}\_{x\sim p}[x^2] - 2 \mathbb{E}\_{x\sim p}[x]\mu\_2 +\mu\_2^2} {2\sigma\_2^2} \\
  &= \frac{1}{2}\log(2\pi\sigma\_2^2) + \frac{\sigma\_1^2 + \mu\_1^2-2\mu\_1\mu\_2+\mu\_2^2}{2\sigma\_2^2} \\
  &= \frac{1}{2}\log(2\pi\sigma\_2^2) + \frac{\sigma\_1^2 + (\mu\_1 - \mu\_2)^2}{2\sigma\_2^2}
  \end{align}
  $$
  The key is recognizing this gives us a sum of several integrals, and each can apply the law of the unconscious statistician. Then we use the fact that $\text{Var}(x)=\mathbb{E}[x^2]-\mathbb{E}[x]^2$. The rest is just rearranging.
  
  Putting it all together:
  $$
  \begin{align}
  D\_\text{KL}(P \| Q) &= -\frac{1}{2}\left(1 + \log(2\pi\sigma\_1^2) \right) + \frac{1}{2}\log(2\pi\sigma\_2^2) + \frac{\sigma\_1^2 + (\mu\_1 - \mu\_2)^2}{2\sigma\_2^2} \\
  &= \log (\sigma\_2) - \log(\sigma\_1) + \frac{\sigma\_1^2 + (\mu\_1 - \mu\_2)^2}{2\sigma\_2^2} - \frac{1}{2}
  \end{align}
  $$
  
  In this special case, we know that $q$ is a standard normal, so
  $$
  \begin{align}
  D\_\text{KL}(P \| Q) &= -\log \sigma\_1 + \frac{1}{2}\left(\sigma\_1^2 +
  \mu\_1^2 - 1 \right) \\
  &= - \frac{1}{2}\left(1 + 2\log \sigma\_1- \mu\_1^2 -\sigma\_1^2 \right)
  \end{align}
  $$
  In the case that we have a $k$-variate normal with diagonal covariance for $p$, and a multivariate normal with covariance $I$, this is the sum of $k$ univariate normal distributions because in this case the distributions are independent.
  
  The code is a correct implementation of this expression because $\log(\sigma\_1^2) = 2 \log(\sigma\_1)$ and in the code, `sigma` is the logarithm of the variance.
  
  ---
  
  $^\*$The reason that it's convenient to work on the scale of the *log-variance* is that the log-variance can be any real number, but the *variance* is constrained to be non-negative by definition. It's easier to perform optimization on the unconstrained scale than it is to work on the constrained scale in $\eta^2$. Also, we want to avoid "round-tripping," where we compute $\exp(y)$ in one step and then $\log(\exp(y))$ in a later step, because this incurs a loss of precision. In any case, autograd takes care of all of the messy details with adjustments to gradients resulting from moving from one scale to another.
  
  > **Comments:**
  > **username_0**: The part I'm a little confused by is that how do you know the mean of p(x)? Especially in neural network case? Cause it's intractable, or does reducing it down to a multivariate Gauss make it tractable? But if so is the mean and SD based on what the data produces?  
  > **username_2**: The latent representation of a VAE emits a mean and (a transformation of) the variance for each input: the $i$th sample has an associated distribution with parameters $(\mu_i, \sigma_i^2)$. In other words, there's a distribution $p_i(x)$ for each of the $i$ inputs.  
  
  ---
  
  ### Answer by username_1 on 2021-02-16. Score: 8
  
  This is the analytical form of the KL divergence between two multivariate Gaussian densities with diagonal covariance matrices (i.e. we assume independence). More precisely, it's the KL divergence between the **variational distribution**
  
  $$
  q\_{\boldsymbol{\phi}}(\mathbf{z})
  =
  \mathcal{N}\left(\mathbf{z} ; \boldsymbol{\mu}, \mathbf{\Sigma} = \boldsymbol{\sigma}^{2}\mathbf{I}\right)
  =
  \frac{\exp \left(-\frac{1}{2}\left(\mathbf{z} - \boldsymbol{\mu}\right)^{\mathrm{T}} \mathbf{\Sigma}^{-1}\left(\mathbf{z}-\boldsymbol{\mu} \right)\right)}{\sqrt{(2 \pi)^{J}\left|\mathbf{\Sigma}\right|}}
  \tag{1}\label{1}
  $$
  
  and the prior (it's the same as above, but with mean and covariance equal to the zero vector and the identity matrix, respectively)
  
  $$
  p(\mathbf{z})=\mathcal{N}(\mathbf{z} ;\boldsymbol{0}, \mathbf{I})
  =
  \frac{\exp \left(-\frac{1}{2}\mathbf{z}^{\mathrm{T}}\mathbf{z}\right)}{\sqrt{(2 \pi)^{J}}}
  \tag{2}\label{2}
  $$
  
  where
  
  * $\boldsymbol{\mu} \in \mathbb{R}^J$ is the mean vector (we assume column vectors, so $\boldsymbol{\mu}^T$ would be a row vector)
  * $\mathbf{\Sigma} = \boldsymbol{\sigma}^{2}\mathbf{I} \in \mathbb{R}^{J \times J}$ is a diagonal covariance matrix (with the vector $\boldsymbol{\sigma}^{2}$ on the diagonal of the identity)
  * $\mathbf{z} \in \mathbb{R}^J$ is a sample (latent vector) from these Gaussians with dimensionality $J$ (or, at the same time, the input variable of the density)
  * $\left|\mathbf{\Sigma}\right| = \operatorname{det} \mathbf{\Sigma}$ is the determinant (so a number) of the diagonal covariance matrix, which is just the product of the diagonal elements for a diagonal matrix (which is the case); so, in the case of the identity, the determinant is $1$
  * $\boldsymbol{0} \in \mathbb{R}^J$ is the zero vector
  * $\mathbf{I} \in \mathbb{R}^{J \times J}$ is an identity matrix
  * $\mathbf{z}^{\mathrm{T}}\mathbf{z} = \sum\_{i=1}^J z\_i^2 \in \mathbb{R}$ is the dot product (hence a number)
  
  Now, the (negative of the) KL divergence is defined as follows
  
  \begin{align}
  -D\_{K L}\left(q\_{\boldsymbol{\phi}}(\mathbf{z}) \| p(\mathbf{z})\right)
  &=
  \int q\_{\boldsymbol{\phi}}(\mathbf{z})\left(\log p(\mathbf{z})-\log q\_{\boldsymbol{\phi}}(\mathbf{z})\right) d \mathbf{z} \\
  &=
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})} \left[ \log p(\mathbf{z})-\log q\_{\boldsymbol{\phi}}(\mathbf{z})\right] \label{3}\tag{3}
  \end{align}
  
  Given that we have logarithms here, let's compute the logarithm of equations \ref{1} and \ref{2}
  
  \begin{align}
  \log
  \left(
  \mathcal{N}\left(\mathbf{z} ; \boldsymbol{\mu}, \mathbf{\Sigma} \right)
  \right)
  &=
  \dots \\
  &=
  -\frac{1}{2}(\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}} \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})-\frac{J}{2} \log (2 \pi)-\frac{1}{2} \log |\mathbf{\Sigma} |
  \end{align}
  
  and
  
  \begin{align}
  \log
  \left(
  \mathcal{N}(\mathbf{z} ;\boldsymbol{0}, \mathbf{I})
  \right)
  &=
  \dots \\
  &=
  -\frac{1}{2}\mathbf{z}^{\mathrm{T}} \mathbf{z}-\frac{J}{2} \log (2 \pi)
  \end{align}
  
  We can now replace these in equation \ref{3} (below, I have already performed some simplifications, to remove verbosity, but you can check them!)
  
  \begin{align}
  \frac{1}{2}
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  -\mathbf{z}^{\mathrm{T}} \mathbf{z}
  + (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}} \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})
  + \log |\mathbf{\Sigma} |
  \right]
  \tag{4}\label{4}
  \end{align}
  Now, given that $\mathbf{\Sigma}$ is diagonal and the log of a product is just a sum of the logarithms, we have $\log |\mathbf{\Sigma} | = \sum\_{i=1}^J \log \sigma\_{ii}$, so we can continue
  
  \begin{align}
  \frac{1}{2}
  \left(
  -
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})}
  \left[
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  +
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}} \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})
  \right]
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})}
  \left[
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  +
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  \operatorname{tr} \left(
  \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}}
  \right)
  \right]
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})}
  \left[
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  +
  \operatorname{tr} \left(
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  \mathbf{\Sigma}^{-1}(\mathbf{z}-\boldsymbol{\mu})
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}}
  \right]
  \right)
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})}
  \left[
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  +
  \operatorname{tr} \left(
  \mathbf{\Sigma}^{-1}
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})} \left[
  (\mathbf{z}-\boldsymbol{\mu})
  (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}}
  \right]
  \right)
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})}
  \left[
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  +
  \operatorname{tr} \left(
  \mathbf{\Sigma}^{-1}
  \mathbf{\Sigma}
  \right)
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})}
  \left[
  \mathbf{z}^{\mathrm{T}} \mathbf{z}
  \right]
  +
  J
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \left(
  -
  \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})}
  \left[
  \operatorname{tr}
  \left(
  \mathbf{z} \mathbf{z}^{\mathrm{T}}
  \right)
  \right]
  +
  \sum\_{i=1}^J 1
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \left(
  -
  \operatorname{tr}
  \left(
  \mathbf{\Sigma}\right)
  -
  \operatorname{tr}
  \left(
  \boldsymbol{\mu} \boldsymbol{\mu}^T
  \right)
  +
  \sum\_{i=1}^J 1
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \left(
  -
  \sum\_{i=1}^J \sigma\_{ii}
  -
  \sum\_{i=1}^J \mu\_{i}^2
  +
  \sum\_{i=1}^J 1
  +
  \sum\_{i=1}^J \log \sigma\_{ii}
  \right)
  &= \\
  \frac{1}{2}
  \sum\_{i=1}^J
  \left(
  1
  +
  \log \sigma\_{ii}
  -
  \sigma\_{ii}
  -
  \mu\_{i}^2
  \right)
  \end{align}
  
  In the above simplifications, I also applied the following rules.
  
  * the linearity of the expectations,
  * $\operatorname{tr}(ba^T) = a^Tb$ (where $\operatorname{tr}$ is the trace of a matrix)
  * [you can swap the trace and the expectation](https://math.stackexchange.com/q/3268268/168764)
  * [$\mathbf{\Sigma} =
    \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})} \left[
    (\mathbf{z}-\boldsymbol{\mu})
    (\mathbf{z}-\boldsymbol{\mu})^{\mathrm{T}}
    \right] $](https://en.wikipedia.org/wiki/Covariance_matrix) (this is just the definition of the *covariance matrix*)
  * $\operatorname{tr}\left( \mathbf{\Sigma}^{-1}
    \mathbf{\Sigma} \right) = \operatorname{tr}\left(\mathbf{I} \right) = J = \sum\_{i=1}^J 1$
  * [$ \mathbb{E}\_{q\_{\boldsymbol{\phi}}(\mathbf{z})}
    \left[
    \mathbf{z} \mathbf{z}^{\mathrm{T}}
    \right] = \mathbf{\Sigma} + \boldsymbol{\mu} \boldsymbol{\mu}^T$](https://www.probabilitycourse.com/chapter6/6_1_5_random_vectors.php) (definition of the *correlation matrix*)
  
  The official PyTorch implementation of the VAE, which can be found [here](https://github.com/pytorch/examples/blob/master/vae/main.py), also uses this formula. This formula can also be found in [Appendix B of the VAE paper](https://arxiv.org/pdf/1312.6114.pdf#page=10), but the long proof that I've just written above is not given. Note that, in my proof above, $\sigma$ is the variance and is denoted by $\sigma^2$ in the paper (as it is usually the case to denote the variance as the square of the standard deviation $\sigma$, but, again, in my proof above $\sigma$ is the variance).
  
  
  ---
  
  