Answers.Body: <p>Your three dimensional latent representation consists of two images of <em>mean pixels</em> and <em>covariance pixels</em> as shown in Fig. 3. Which represents a Gaussian distribution with the mean and covariance for each pixel in the latent representation. Each pixel value is a random variable.</p>
  <p>Now, have a close look at KL-loss Eq. 3 and it's corresponding description in the paper:</p>
  <p><span class="math-container">$$\mathcal{L}_{KL} = \frac{1}{2 \times (\frac{W}{16} \times \frac{H}{16}) } \sum^M_{m = 1}[\mu^2_m + \sigma^2_m - \log(\sigma^2_m) - 1]$$</span></p>
  <blockquote>
  <p>Finally, <span class="math-container">$M$</span> is  the  dimensionality  of  the  latent  features <span class="math-container">$\theta \in \mathbb{R}^M$</span> with  mean <span class="math-container">$\mu  = [\mu_1,...,\mu_M]$</span> and covariance matrix <span class="math-container">$\Sigma = \text{diag}(\sigma_1^2,...,\sigma_M^2)$</span>, [...].</p>
  </blockquote>
  <p>The covariance matrix is diagonal, thus <strong>all pixel values are independent of each other</strong>. That is the reason why we have this nice analytical form for the KL-divergence given by Eq. 3. Therefore you can treat your 2D random matrix simply as a random vector of size <span class="math-container">$M = \frac{W}{16} \times \frac{H}{16}$</span> (<span class="math-container">$\times 3$</span> if you like to include color dimension). The third dimension (RGB channel) can be considered independent as well, therefore it can be also flattened to a vector and appended. Indeed this is what is done in the paper indicated by the second half of the sentence from above:</p>
  <blockquote>
  <p>that are reparameterized by via sampling from a standard multivariate Gaussian <span class="math-container">$\epsilon \sim \mathcal{N}(0,I_M)$</span>, i.e. <span class="math-container">$\theta = \mu + \Sigma^{\frac{1}{2}}\epsilon$</span>.</p>
  </blockquote>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2021-02-15T13:23:49.813Z
Answers.Text: To clarify, will flattening the W/16 x H/16 matrix into a vector of size M maintain the spatial correspondence _because_ the pixel values are independent of each other? When I say spatial correspondence, I refer to this part of the paper: "we use three-dimensional latent feature maps, i.e., channel, height and width dimensions, rather than one-dimensional latent vectors with only channel dimension, for improving the reconstruction quality and preserve more spatial information."
Answers.UserDisplayName: magmacollaris
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2021-02-16T13:16:41.390Z
Answers.Text: As an additional point of clarification, does this mean that each mean pixel will get its own covariance? Since the covariance is diagonal, would this mean that the covariance matrix is (H/16 x W/16) x (H/16 x W/16) in size?
Answers.UserDisplayName: magmacollaris
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2021-02-09T21:00:12.117Z
Answers.LastActivityDate: 2021-02-10T12:29:57.027Z
Answers.LastEditDate: 2021-02-10T12:29:57.027Z
Answers.UserDisplayName: Tinu
Body: <p>I'm trying to implement a variational auto-encoder (as seen in Section 3.1 here: <a href="https://arxiv.org/pdf/2004.06271.pdf" rel="nofollow noreferrer">https://arxiv.org/pdf/2004.06271.pdf</a>).</p>
  <p>It differs from a traditional VAE because it encodes its input images to <strong>three-dimensional latent feature maps</strong>. In other words, the latent feature maps have a width, height and channel dimension rather than just a channel dimension like a traditional VAE.</p>
  <p>When calculating the Kullback-Liebler divergence as part of the loss function, I need the <strong>mean</strong> and <strong>covariance</strong> that is the output of the encoder. However, if the latent feature maps are <strong>three-dimensional</strong>, this means that the output of the encoder is three-dimensional, and therefore each latent feature is a 2D matrix.</p>
  <p><strong>How can I derive a mean and covariance from a 2D matrix to calculate the KL divergence?</strong></p>
  
Tags: <convolutional-neural-networks><computer-vision><math><variational-autoencoder><kl-divergence>
Title: How do you calculate KL divergence on a three-dimensional space for a Variational Autoencoder?
UserDisplayName: magmacollaris
ThreadText: # How do you calculate KL divergence on a three-dimensional space for a Variational Autoencoder?
  
  (Asked by: username_0 on 2021-02-09)
  
  I'm trying to implement a variational auto-encoder (as seen in Section 3.1 here: <https://arxiv.org/pdf/2004.06271.pdf>).
  
  It differs from a traditional VAE because it encodes its input images to **three-dimensional latent feature maps**. In other words, the latent feature maps have a width, height and channel dimension rather than just a channel dimension like a traditional VAE.
  
  When calculating the Kullback-Liebler divergence as part of the loss function, I need the **mean** and **covariance** that is the output of the encoder. However, if the latent feature maps are **three-dimensional**, this means that the output of the encoder is three-dimensional, and therefore each latent feature is a 2D matrix.
  
  **How can I derive a mean and covariance from a 2D matrix to calculate the KL divergence?**
  
  **Tags:** convolutional-neural-networks, computer-vision, math, variational-autoencoder, kl-divergence
  
  ---
  
  ## 1 answers
  
  ### Answer by username_1 on 2021-02-09. Score: 1 | Accepted answer
  
  Your three dimensional latent representation consists of two images of *mean pixels* and *covariance pixels* as shown in Fig. 3. Which represents a Gaussian distribution with the mean and covariance for each pixel in the latent representation. Each pixel value is a random variable.
  
  Now, have a close look at KL-loss Eq. 3 and it's corresponding description in the paper:
  
  $$\mathcal{L}\_{KL} = \frac{1}{2 \times (\frac{W}{16} \times \frac{H}{16}) } \sum^M\_{m = 1}[\mu^2\_m + \sigma^2\_m - \log(\sigma^2\_m) - 1]$$
  
  > Finally, $M$ is the dimensionality of the latent features $\theta \in \mathbb{R}^M$ with mean $\mu = [\mu\_1,...,\mu\_M]$ and covariance matrix $\Sigma = \text{diag}(\sigma\_1^2,...,\sigma\_M^2)$, [...].
  
  The covariance matrix is diagonal, thus **all pixel values are independent of each other**. That is the reason why we have this nice analytical form for the KL-divergence given by Eq. 3. Therefore you can treat your 2D random matrix simply as a random vector of size $M = \frac{W}{16} \times \frac{H}{16}$ ($\times 3$ if you like to include color dimension). The third dimension (RGB channel) can be considered independent as well, therefore it can be also flattened to a vector and appended. Indeed this is what is done in the paper indicated by the second half of the sentence from above:
  
  > that are reparameterized by via sampling from a standard multivariate Gaussian $\epsilon \sim \mathcal{N}(0,I\_M)$, i.e. $\theta = \mu + \Sigma^{\frac{1}{2}}\epsilon$.
  
  > **Comments:**
  > **username_0**: To clarify, will flattening the W/16 x H/16 matrix into a vector of size M maintain the spatial correspondence _because_ the pixel values are independent of each other? When I say spatial correspondence, I refer to this part of the paper: "we use three-dimensional latent feature maps, i.e., channel, height and width dimensions, rather than one-dimensional latent vectors with only channel dimension, for improving the reconstruction quality and preserve more spatial information."  
  > **username_0**: As an additional point of clarification, does this mean that each mean pixel will get its own covariance? Since the covariance is diagonal, would this mean that the covariance matrix is (H/16 x W/16) x (H/16 x W/16) in size?  
  
  ---
  
  