<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>

CC BY-SA 4.0
2021-02-15T13:23:49.813Z
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."
magmacollaris
CC BY-SA 4.0
2021-02-16T13:16:41.390Z
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?
magmacollaris
CC BY-SA 4.0
2021-02-09T21:00:12.117Z
2021-02-10T12:29:57.027Z
2021-02-10T12:29:57.027Z
Tinu
<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>

<convolutional-neural-networks><computer-vision><math><variational-autoencoder><kl-divergence>
How do you calculate KL divergence on a three-dimensional space for a Variational Autoencoder?
magmacollaris
# 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?  

---

