Use KL divergence as loss between two multivariate Gaussians

username: rasoolfa
cooked: <p>Hi,</p>
<p>I want to use KL divergence as loss function between two  multivariate Gaussians. Is the following right way to do it?</p>
<blockquote>
<p>mu1 = torch.rand((B, D), requires_grad=True)<br>
std1 = torch.rand((B, D), requires_grad=True)<br>
p  = torch.distributions.Normal(mu1, std1)<br>
mu2 = torch.rand((B, D))<br>
std2 = torch.rand((B, D))<br>
q  = torch.distributions.Normal(mu2, std2)</p>
<p>loss = torch.distributions.kl_divergence(p, q).mean()<br>
loss.backward()<br>
…</p>
</blockquote>
<p>My understanding is that <a href="https://github.com/pytorch/pytorch/blob/master/torch/distributions/kl.py#L377" rel="noopener nofollow ugc">torch.distributions.kl_divergence</a> computes kl(p,q) like derivations in section 9 of <a href="http://stanford.edu/~jduchi/projects/general_notes.pdf" rel="noopener nofollow ugc">this</a> document.</p>

username: Rojin
cooked: <p>any update on this question?</p>

username: Nikronic
cooked: <p>Hi,</p>
<p>You are right. When you are using distributions from <code>torch.distribution</code> package, you are doing fine by using <code>torch.distribution.kl_divergence</code>. But if you want to get <code>kl</code> by passing two tensors obtain elsewhere, you can do following approach:</p>
<aside class="quote quote-modified" data-post="2" 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/nikronic/48/71873_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/kl-divergence-between-two-layers/53024/2">KL-divergence between two multivariate gaussian</a> 
  </div>
  <blockquote>
    Hi, 
Yes, this is the correct approach. 
Just be aware that the input a must should contain log-probabilities and the target b should contain probability. 


<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> 
By the way, PyTorch use this approach: 
 <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">[image]</a> 
<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> 
Good luck 
Nik
  </blockquote>
</aside>

<p><a class="mention" href="/u/rojin">@Rojin</a> I have posted this on your thread actually.</p>
<p>This is the <code>kl</code> between two arbitrary layers.<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: LiUzHiAn
cooked: <p>Hi，<br>
I test the <code>kl_divergence</code> as follows, the tensors’ shapes are <code>[batch_size,n]</code>, where <code>n</code> is my gaussian distribution dimension (i.e. 3 here),</p>
<pre><code class="lang-auto">        mu1 = torch.Tensor([[1., 2., 3.],
						[2., 3., 4.]])
        var_1 = torch.Tensor([[1., 1., 1.],
						  [4., 9., 16.]])

	mu2 = torch.Tensor([[1., 3., 4.],
						[2., 3., 4.]])
	var_2 = torch.Tensor([[1., 4., 9.],
						  [4., 9., 16.]])

	p = torch.distributions.Normal(mu1, var_1)
	q = torch.distributions.Normal(mu2, var_2)
	kl_loss = torch.distributions.kl_divergence(p, q)

	print(kl_loss)
</code></pre>
<p>The outpur is :</p>
<pre><code class="lang-auto">tensor([[0.6542, 0.9488, 1.7096],
        [0.0000, 0.0000, 0.0000]])
</code></pre>
<p>Why 3 numbers there for each instance? In my thought, the case above contains two pairs <code>(p,q)</code> , each pair should have a scalar KL loss.</p>
<p>Any ideas? Thank you</p>

username: LiUzHiAn
cooked: <p>Ooh, I got it !</p>
<p>In my above case, I assume the distribution to be multivariate gaussian, I should use <code>torch.distributions.MultivariateNormal(mu1, var_1)</code> instead, which meets my hope. <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=9" title=":sweat_smile:" class="emoji" alt=":sweat_smile:"></p>

username: Michael-J98
cooked: <p>Hi,<br>
can I ask a question that why you set requires_grad = True in mu1&amp;std1 while mu2&amp;std2 did not?<br>
Is there something we should watch out?</p>

username: thomas_liu
cooked: <p>Hi,<br>
I run into a problem where MultivariateNormal does not allow [batch ,n] shape because batch and n has to be equal, Do you manage to solve this?</p>

username: drmeerkatli
cooked: <p>Does the KL calculation correct?<br>
log(a/b) = log(a) - b???</p>