KL divergence loss

Asad_Ullah: <p>Trying to implement KL divergence loss but got nan always.</p>
<p>p = torch.randn((100,100))<br>
q = torch.randn((100,100))<br>
kl_loss  = torch.nn.KLDivLoss(size_average= False)(p.log(), q)<br>
output = nan<br>
p_soft = F.softmax( p )<br>
q_soft = F.softmax( q )<br>
kl_loss  = torch.nn.KLDivLoss(size_average= False)(p_soft.log(), q_soft)<br>
output = 96.7017<br>
Do we have to pass the distributions (p, q) through softmax function always?</p>

ptrblck: <p>According to the docs:</p>
<blockquote>
<p>As with <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.NLLLoss"> <code>NLLLoss</code> </a>, the input given is expected to contain <em>log-probabilities</em> and is not restricted to a 2D Tensor. The targets are given as <em>probabilities</em> (i.e. without taking the logarithm).</p>
</blockquote>
<p>your code snippet looks alright. I would recommend to use <code>log_softmax</code> instead of <code>softmax().log()</code>, as the former approach is numerically more stable.</p>

Hmrishav_Bandyopadhy: <p>Hi,<br>
As i guess, KL divergence is supposed to return a numerical value representing the  distance between 2 probability distributions in feature space. However, this is what I got using log softmax…</p>
<pre><code class="lang-auto">&gt;&gt;&gt; k=torch.rand(256)
&gt;&gt;&gt; k1=k.clone()
&gt;&gt;&gt; F.kl_div(F.log_softmax(k),k1,reduction="none").mean()
tensor(2.4483)
</code></pre>
<p>On the other hand when I use simple log, the answer is zero, which is expected. Can you let me know what i should use when comparing 2 layers in pytorch with KL divergence?</p>

ptrblck: <p>The target should be given as probabilities:</p>
<pre><code class="lang-python">k = torch.rand(256)
k1 = k.clone()
F.kl_div(F.log_softmax(k, 0), F.softmax(k1, 0), reduction="none").mean()
&gt; tensor(6.2333e-10)
</code></pre>

Hmrishav_Bandyopadhy: <p>Thank you ! I was missing that out <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=9" title=":sweat_smile:" class="emoji" alt=":sweat_smile:"></p>

pascal_notsawo: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="65393">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p><code>k = torch.rand(256)</code></p>
</blockquote>
</aside>
<p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
There is something I don’t understand.<br>
Let’s assume that: <code>shape(target) = shape(input) = (batch_size, N)</code></p>
<p>The log_softmax / softmax must be in dimension 0?<br>
<code>F.kl_div(F.log_softmax(logits, dim = 0), F.softmax(target, dim = 0), reduction="none").mean()</code></p>
<p>Or in dimension 1 ?<br>
<code>F.kl_div(F.log_softmax(logits, dim = 1), F.softmax(target, dim = 1), reduction="none").mean()</code></p>
<p>Personally I think in dimension 1 (N here is a bit like the number of classes for a classification)</p>

ptrblck: <p>You would apply the <code>log_softmax</code> in the class dimension, so usually in <code>dim1</code>.<br>
Note that my example is not a really representative one and I was just reusing the posted code.</p>

Lezcano: <p>The following issue is relevant for people using <code>kl_div</code> or its <code>nn</code> module, as its current behaviour is wrong</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/57459">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/57459" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/57459" target="_blank" rel="noopener nofollow ugc">KLDivLoss and F.kl_div compute KL(Q || P) rather than KL(P || Q)</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2021-05-03" data-time="08:24:55" data-timezone="UTC">08:24AM - 03 May 21 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/BenoitDalFerro" target="_blank" rel="noopener nofollow ugc">
          <img alt="BenoitDalFerro" src="https://avatars.githubusercontent.com/u/69694610?v=4" class="onebox-avatar-inline" width="20" height="20">
          BenoitDalFerro
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          high priority
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: docs
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: nn
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: numpy
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: correctness (silent)
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: special
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🐛 Bug
### Executive summary:
The inputs of `KLDivLoss` and `F.kl_div` are i<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">nverted. `target` should be `input` and `input` should be `target`. This is bad, as this function is not symmetric.

### Further points
torch.functional.kl_div() is inverting positional argument of source (P) and target (Q) distribution, incorrect development of logarithm fractions for the visible part, resulting in various problems up to and including negative Kl divergence (as reported in here (https://github.com/pytorch/pytorch/issues/32520) and Shannon Jensen Divergence not bounded at ]0,1] (bits) or ]0,log2(e)=1.4426950408889634...] (nats) and eventually for computation of the Shannon Jensen metric (square root of the JSD) to NaN issues (as a result of Kl underflow and JSM computing square root of negative numbers...)

The [documentation seems](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html) wrong too (and explains partly the errors experienced)

l(x,y)=L={l1​,…,lN​},ln​=yn​⋅(logyn​−xn​) 

Which actually should be : 
l(x,y)=L={l1​,…,lN​},ln​=xn​⋅(log(xn)​−log(yn​))=xn​⋅log(xn)​−xn​⋅log(yn​)
Because
l(x,y)= 1/NΣ xn⋅log(xn/ yn),ln​=xn​⋅(log(xn/yn​)=xn​⋅(log(xn)-logyn​))=xn​⋅log(xn)​−xn​⋅log(yn​)

Now taking into account that xn and yn are inverted and that xn is a log_softmax
l(x,y)=L={l1​,…,lN​},ln​=yn​⋅(logyn​−xn​), xn=log(wn) 
(=)
l(x,y)=L={l1​,…,lN​},ln​=yn​⋅(logyn​−xn​)=yn​⋅(logyn​−logwn)​ and here we fall back on the equation but with inversion of yn and xn terms or more exactly inversion of one and the antilog of the other, any reader will understand it as Kl(x||y) but in fact is Kl(y||w) with x=log(w).

Actually naming is misleading, it is not a Kl divergence since F.log_softmax has to be computed beforehand, I understand numercial stability issues in computing log of F.softmax() but that's something explicit in the documentation, the reader should be forewarned as this not so trivial to grasp issue, made very clear.

## To Reproduce

Steps to reproduce the behavior:

```
M =0.5*(p+q)
JSD = 0
JSD += 0.5*F.kl_div(F.log_softmax(M,dim=dim).log(), F.softmax(p, dim=dim), None, None, 'none')
JSD += 0.5*F.kl_div(F.log_softmax(M,dim=dim).log(), F.softmax(q, dim=dim), None, None, 'none')

return torch.sqrt(JSD)
``` 

## Expected behavior

First two positional arguments in wrong order as per mathematical error in the documentation, for the general Kullback-Leibner it should be something like

F.kl_div(F.softmax(p, dim=dim),F.log_softmax(q,dim=dim), None, None, 'none')

and in the context of the JSD's specifics : 

F.kl_div(F.softmax(p, dim=dim),F.log_softmax(M,dim=dim), None, None, 'none')

Second question is why do we need to input a F.log_softmax for the first (which should be in fact the second) positional argument in the first place ? I understand that F.softmax(M, dim=dim),log() has numerical stability compared to computing directly F.log_softmax yet could not be part of a routine treatment of the second term ? Otherwise it is not KL_divergence and the name is misleading !

## Proposed fix for Kl underflow

The Kl numerical underflow issue is actually documented in Cover &amp; al (2006), 2.3 Relative Entropy and Mutual Information p.45 : "we use the convention that 0 log 0/0 =0 and the convention (based on continuity arguments) that 0 log 0/q = 0 and p log p /0 = inf. Therefore postprocessing fix is : 

D_pM = (p*(torch.log(p)-F.log_softmax(M,dim=dim)))
D_qM = (q*(torch.log(q)-F.log_softmax(M,dim=dim)))
D_pM[torch.isnan(D_pM)] = 0
D_qM[torch.isnan(D_qM)] = 0

## Environment

Collecting environment information...
PyTorch version: 1.7.0
Is debug build: True
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Famille
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: GeForce GTX 1050
Nvidia driver version: 457.63
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\bin\cudnn64_7.dll
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] numpydoc==1.1.0
[pip3] pytorch-metric-learning==0.9.98
[pip3] torch==1.7.0 (UPDATED TO 1.8.1 problem remains)
[pip3] torchaudio==0.7.0
[pip3] torchvision==0.8.1
[conda] blas                      2.16                        mkl    conda-forge
[conda] cudatoolkit               10.2.89              hb195166_8    conda-forge
[conda] libblas                   3.8.0                    16_mkl    conda-forge
[conda] libcblas                  3.8.0                    16_mkl    conda-forge
[conda] liblapack                 3.8.0                    16_mkl    conda-forge
[conda] liblapacke                3.8.0                    16_mkl    conda-forge
[conda] mkl                       2020.1                      216
[conda] numpy                     1.19.1           py37hae9e721_0    conda-forge
[conda] numpydoc                  1.1.0                      py_1    conda-forge
[conda] pytorch                   1.7.0           py3.7_cuda102_cudnn7_0    pytorch
[conda] pytorch-metric-learning   0.9.98             pyh39e3cac_0    metric-learning
[conda] torchaudio                0.7.0                      py37    pytorch
[conda] torchvision               0.8.1                py37_cu102    pytorch

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @brianjo @mruberry @albanD @walterddr @rgommers @heitorschueroff</span></p>
  </div>

  </article>

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

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