KL-divergence between two multivariate gaussian

username: Rojin
cooked: <p>I have two multivariate Gaussian distributions that I would like to calculate the kl divergence between them. each is defined with a vector of <strong>mu</strong> and a vector of <strong>variance</strong> (similar to <strong>VAE</strong>  <strong>mu</strong>  and  <strong>sigma layer</strong>). What is the best way to calculate the KL between the two? Is this even doable? because I do not have the covariance matrix. If that is not doable, what if I take samples from both distributions and calculate the KL between the samplings? is that the right way of doing this? is there any build-in function for any of these calculations?</p>

username: Nikronic
cooked: <p>Hi,</p>
<p>Yes, this is the correct approach.<br>
Just be aware that the input <code>a</code> must should contain <code>log-probabilities</code> and the target <code>b</code> should contain <code>probability</code>.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/35fed93b1ef05175143f883c6f89f06c6dd9429b/aten/src/ATen/native/Loss.cpp#L71">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/35fed93b1ef05175143f883c6f89f06c6dd9429b/aten/src/ATen/native/Loss.cpp#L71" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/35fed93b1ef05175143f883c6f89f06c6dd9429b/aten/src/ATen/native/Loss.cpp#L71" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/35fed93b1ef05175143f883c6f89f06c6dd9429b/aten/src/ATen/native/Loss.cpp#L71</a></h4>



    <pre class="onebox"><code class="lang-cpp">
      <ol class="start lines" start="61" style="counter-reset: li-counter 60 ;">
          <li>  return apply_loss_reduction(output, reduction);</li>
          <li>}</li>
          <li></li>
          <li>Tensor margin_ranking_loss(const Tensor&amp; input1, const Tensor&amp; input2, const Tensor&amp; target, double margin, int64_t reduction) {</li>
          <li>  auto output =  (-target * (input1 - input2) + margin).clamp_min_(0);</li>
          <li>  return apply_loss_reduction(output, reduction);</li>
          <li>}</li>
          <li></li>
          <li>Tensor kl_div(const Tensor&amp; input, const Tensor&amp; target, int64_t reduction) {</li>
          <li>  auto zeros = at::zeros_like(target);</li>
          <li class="selected">  auto output_pos = target * (at::log(target) - input);</li>
          <li>  auto output = at::where(target &gt; 0, output_pos, zeros);</li>
          <li>  return apply_loss_reduction(output, reduction);</li>
          <li>}</li>
          <li></li>
          <li>Tensor kl_div_backward_cpu(const Tensor&amp; grad, const Tensor&amp; input, const Tensor&amp; target, int64_t reduction) {</li>
          <li>  auto grad_input = at::zeros_like(input);</li>
          <li>  auto grad_expand = grad.expand_as(input);</li>
          <li>  AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "kl_div_backward_cpu", [&amp;]() {</li>
          <li>    at::CPU_tensor_apply3&lt;scalar_t, scalar_t, scalar_t&gt;(</li>
          <li>        grad_input,</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p><a href="https://pytorch.org/docs/stable/nn.functional.html?highlight=kl_div#kl-div" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/nn.functional.html?highlight=kl_div#kl-div</a></p>
<p>By the way, PyTorch use this approach:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/0/0a58d819377f67d7c4735e76be5db794ad56f05d.png" data-download-href="https://discuss.pytorch.org/uploads/default/0a58d819377f67d7c4735e76be5db794ad56f05d" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/2X/0/0a58d819377f67d7c4735e76be5db794ad56f05d.png" alt="image" data-base62-sha1="1tx7ArRswZihgE4LqGaoEKrRijP" width="690" height="124" data-dominant-color="212324"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">692×125 8.41 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<a href="https://pytorch.org/docs/stable/distributions.html?highlight=kl_div#torch.distributions.kl.kl_divergence" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/distributions.html?highlight=kl_div#torch.distributions.kl.kl_divergence</a></p>
<p>Good luck<br>
Nik</p>

username: Rojin
cooked: <aside class="quote no-group quote-modified" data-username="Rojin" data-post="1" data-topic="53024">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/rojin/48/11062_2.png" class="avatar"> Rojin:</div>
<blockquote>
<p>F.kl_div(a, b)</p>
</blockquote>
</aside>
<p>Thanks Nick for your input. I should restate my question. I have two multi-variate distributions each defined with “n” mu and sigma. Now I would like to do a KL between these two. Any idea how that can be done?</p>

username: zhl515
cooked: <p>You can sample x1 and x2  from 𝑝1(𝑥|𝜇1,Σ1) and 𝑝2(𝑥|𝜇2,Σ2) respectively, then compute KL divergence using torch.nn.functional.kl_div(x1, x2).</p>

username: Rojin
cooked: <p>As <a class="mention" href="/u/nikronic">@Nikronic</a> mentioned the kl_div requires <code>a</code> to be  <code>log-probabilities</code>  and the target  <code>b</code>  to be  <code>probability</code>. correct me if I’m wrong, but I think sampling from the two distributions is not going to give  log-probabilities and probabilities</p>

username: Rojin
cooked: <p>any suggestion? I don’t think that what you said can be used for sampling layer</p>

username: Nikronic
cooked: <p><a class="mention" href="/u/rojin">@Rojin</a><br>
Actually, I have never used distribution class but based on docs, in this situation you just need to pass them directly to <a href="https://pytorch.org/docs/stable/distributions.html#torch.distributions.kl.kl_divergence" rel="nofollow noopener">kl_div from distribution</a> if the aforementioned distributions are obtained from <a href="https://pytorch.org/docs/stable/distributions.html#torch.distributions.distribution.Distribution" rel="nofollow noopener">Distribution subclasses</a> or you have registered your own.</p>
<p>Please refer to this links for more information:</p>
<ol>
<li><a href="https://pytorch.org/docs/stable/distributions.html#torch.distributions.kl.kl_divergence" rel="nofollow noopener">https://pytorch.org/docs/stable/distributions.html#torch.distributions.kl.kl_divergence</a></li>
<li><a href="https://pytorch.org/docs/stable/distributions.html#torch.distributions.distribution.Distribution" rel="nofollow noopener">https://pytorch.org/docs/stable/distributions.html#torch.distributions.distribution.Distribution</a></li>
</ol>
<p>well, you know, my name is Nik, somehow abbreviated of Nikan (a Persian name) <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>
<p>Bests</p>

username: Rojin
cooked: <p>Hey Nikan!</p>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"> Yeah I realized that from your last name!</p>
<p>Sorry if I’m restating my question ( I modified it a couple of times at the top lol) and probably repeating your answer, I’m kinda new to the KL implementation in pytorch, so I really appreciate your help.</p>
<p>My question is :<br>
I have two VAEs, each with a <strong>mu</strong> and <strong>standard deviation</strong> layers. Now I would like to calculate the KL between these two VAEs using either the <strong>mu</strong> and <strong>sds</strong>, or the <strong>sampling</strong> layers. This sounds to me like a multivariate gaussian KL divergence problem, so I looked at the formula and I noticed that I actually need the covariance matrix of q (if we assume that <strong>KL(p||q)</strong> ). And as far as I know there is no way to calculate the covariance in this case, am I correct?</p>
<p><a class="mention" href="/u/zhl515">@zhl515</a> suggested to directly calculate the KL between the two sampling layers as you said before using the bellow function, but as you said this requires probabilities which I do not have, so I think the following function cannot be used on the sampling layers, is that right? :</p>
<pre><code class="lang-auto"> torch.nn.functional.kl_div(p, q).
</code></pre>
<p>And based on your last comment, you are suggesting to register the distribution, and then use</p>
<pre><code class="lang-auto">torch.distributions.kl.kl_divergence(p, q)
</code></pre>
<p>The only problem is that in order to register the distribution I need to have the covariance matrix, and I can’t obtain that because I only have <strong>mu</strong> and <strong>std</strong>. It makes me think whether that is a possible thing to do at all in neural network. In the VAE paper they don’t have that problem because the are assuming that q has <strong>mean</strong> <strong>0</strong> and <strong>covariance</strong> <strong>I</strong><br>
<aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stats.stackexchange.com/questions/318748/deriving-the-kl-divergence-loss-for-vaes" target="_blank" rel="nofollow noopener">stats.stackexchange.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stats.stackexchange.com/users/90815/yellowpillow" target="_blank" rel="nofollow noopener">
    <img alt="YellowPillow" src="https://www.gravatar.com/avatar/e5693aa1ca63ca1a49d3df043839e5ed?s=128&amp;d=identicon&amp;r=PG&amp;f=1" class="thumbnail onebox-avatar" width="128" height="128">
  </a>
<h4>
  <a href="https://stats.stackexchange.com/questions/318748/deriving-the-kl-divergence-loss-for-vaes" target="_blank" rel="nofollow noopener">Deriving the KL divergence loss for VAEs</a>
</h4>

<div class="tags">
  <strong>kullback-leibler, autoencoders, variational-bayes</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stats.stackexchange.com/users/90815/yellowpillow" target="_blank" rel="nofollow noopener">
    YellowPillow
  </a>
  on <a href="https://stats.stackexchange.com/questions/318748/deriving-the-kl-divergence-loss-for-vaes" target="_blank" rel="nofollow noopener">01:55AM - 14 Dec 17 UTC</a>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Thank you!</p>

username: zhl515
cooked: <p>You said you can’t obtain covariance matrix. In VAE paper, the author  assume the true (but intractable) posterior takes on a approximate Gaussian form with an approximately <strong>diagonal covariance</strong>. So just place  the <strong>std</strong>  on diagonal of convariance matrix, and other elements of matrix are zeros.</p>

username: zhl515
cooked: <p>Now you  can compute KL-divargence of two multivariate Gaussians directly from the below formula:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/5/56267d56ae80f1959d33174740d8a48f11d27d94.png" alt="1565939371(1)" data-base62-sha1="ci7wC9yzX6klNe0PVYUJqWVA9dq" width="622" height="74"></p>

username: Rojin
cooked: <p>Thank you that solves the problem <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: shivangi
cooked: <p>Hey <a class="mention" href="/u/rojin">@Rojin</a><br>
Could you share how you computed the KL-divergence? I am also trying to achieve the same</p>

username: LiUzHiAn
cooked: <p>Hi, <a class="mention" href="/u/zhl515">@zhl515</a>.</p>
<p>Thanks for your solution and I’ve coded as you post, like this:</p>
<pre><code class="lang-auto">def kl_divergence(mu1, mu2, sigma_1, sigma_2):
	sigma_diag_1 = np.eye(sigma_1.shape[0]) * sigma_1
	sigma_diag_2 = np.eye(sigma_2.shape[0]) * sigma_2

	sigma_diag_2_inv = np.linalg.inv(sigma_diag_2)

	kl = 0.5 * (np.log(np.linalg.det(sigma_diag_2) / np.linalg.det(sigma_diag_2))
				- mu1.shape[0] + np.trace(np.matmul(sigma_diag_2_inv, sigma_diag_1))
				+ np.matmul(np.matmul(np.transpose(mu2 - mu1), sigma_diag_2_inv), (mu2 - mu1))
				)

	return kl
</code></pre>
<p>But I have another problem that I cannot call it in a “vectorized” way, or say, I don’t know how to let my code support batch operation in Pytorch. For instance, I have a tensor <code>mu_1</code> with shape [batch_size,n] and another tensor <code>mu_2</code> with shape [batch_size,n]. The variance tensors are same with mu. The <code>n</code> here is the dimensions of Gaussian distribution. How can I calculate kl loss for each instance without using for-loop ?</p>
<p>Any ideas? thx</p>

username: NILOIDE
cooked: <p>Check <a href="https://discuss.pytorch.org/t/how-to-compute-the-diagonal-matrix-in-batch/35449">this thread</a>.<br>
torch.diag_embed() should do the trick.</p>

username: hf_Morii
cooked: <p>hi i don’t know much about the meaning of mu1 mu2 sigma_1 and sigma_2</p>

username: Dane_Mitrev
cooked: <p>I am not sure I understand fully your question, but IMO if you have mu1, mu2, sigma_1, sigma_2, why not create two instances of <code>MultivariateNormal</code> using diagonal covariance matrices?</p>
<pre><code class="lang-auto"># mu1 [batch_size, n]
# assuming that sigma1 is [batch_size, n] also
cov1 = torch.stack([torch.diag(sigma) for sigma in torch.exp(sigma_1)])
mvn1 = MultivariateNormal(mu1, cov1)

...
# similarly for mvn2
</code></pre>
<p><code>MultivariateNormal</code> will interpret the batch_size as the batch dimension automatically thus <code>mvn1</code> would have:</p>
<pre><code class="lang-auto">batch_shape = batch_size
event_shape = n
sample_shape = ()
</code></pre>
<p>when you sample it will take into consideration the batch_shape.</p>
<p>You can compute kl (mvn1, mvn2) using the Pytorch’s implementation.</p>
<p>I hope this can help.</p>

username: hf_Morii
cooked: <p>This is really helpful!<br>
Thanks a lot.</p>

username: muammar
cooked: <p>Is this formula applicable to any arbitrary distribution? e.g., non-Gaussian?</p>

username: banikr
cooked: <p>The cov-variance matrix must be positive definite.<br>
Any idea how to turn a diagonal matrix with zero non-diagonal elements into a positive definite matrix?</p>
<p>The output comes from an FC+BN+ReLU layer.<br>
then I can turn it into a diagonal by:</p>
<pre><code class="lang-auto">cov_var_diag = torch.diag_embed(output)
cov_var_spd = cov_var @ cov_var.mT + 1e-6 # this actually makes symmetric positive definite and didn't work
q = torch.distributions.multivariate_normal.MultivariateNormal(loc=mu, covariance_matrix=cov_var_spd)
</code></pre>
<p>Also looked into Cholesky decomposition into lower triangular matrix as <code>scale_tril</code> attribute in MVN function, but that also didn’t end well.</p>