There's something wrong with KL divergece

kirk86: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/4/14e6a37121dbca52e546aaeed41b959262fb5ecd.png" data-download-href="https://discuss.pytorch.org/uploads/default/14e6a37121dbca52e546aaeed41b959262fb5ecd" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/4/14e6a37121dbca52e546aaeed41b959262fb5ecd_2_517x318.png" alt="image" data-base62-sha1="2YTGMlXayrQAatgyW6wcMePaWrX" width="517" height="318" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/4/14e6a37121dbca52e546aaeed41b959262fb5ecd_2_517x318.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/4/14e6a37121dbca52e546aaeed41b959262fb5ecd_2_775x477.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/4/14e6a37121dbca52e546aaeed41b959262fb5ecd_2_1034x636.png 2x" data-dominant-color="EFEFEF"><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">1322×816 72.4 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>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e47dc20ea65fd62c619537dd96364f0ee77538d6.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/e47dc20ea65fd62c619537dd96364f0ee77538d6" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e47dc20ea65fd62c619537dd96364f0ee77538d6_2_345x212.jpeg" alt="image" data-base62-sha1="wBkivRATRDXceUXf1gbfAqkpWC2" width="345" height="212" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e47dc20ea65fd62c619537dd96364f0ee77538d6_2_345x212.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e47dc20ea65fd62c619537dd96364f0ee77538d6_2_517x318.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e47dc20ea65fd62c619537dd96364f0ee77538d6_2_690x424.jpeg 2x" data-dominant-color="FAFBFB"><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">1322×816 42.2 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>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/1/51c379737afca4b5cbe9c08ed9ac1afc814e554a.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/51c379737afca4b5cbe9c08ed9ac1afc814e554a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/1/51c379737afca4b5cbe9c08ed9ac1afc814e554a_2_345x228.jpeg" alt="image" data-base62-sha1="bFjtA28V0bImg5w73qTuy382ZZ0" width="345" height="228" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/1/51c379737afca4b5cbe9c08ed9ac1afc814e554a_2_345x228.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/5/1/51c379737afca4b5cbe9c08ed9ac1afc814e554a_2_517x342.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/1/51c379737afca4b5cbe9c08ed9ac1afc814e554a_2_690x456.jpeg 2x" data-dominant-color="F4F1F0"><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">1322×876 73.1 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>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/b/0b3f1a2b713e91d633923ed488a7bfec053a7591.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/0b3f1a2b713e91d633923ed488a7bfec053a7591" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b3f1a2b713e91d633923ed488a7bfec053a7591_2_345x228.jpeg" alt="image" data-base62-sha1="1BurMnsrYdbwzRy89U7LilDH5sZ" width="345" height="228" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b3f1a2b713e91d633923ed488a7bfec053a7591_2_345x228.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b3f1a2b713e91d633923ed488a7bfec053a7591_2_517x342.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b3f1a2b713e91d633923ed488a7bfec053a7591_2_690x456.jpeg 2x" data-dominant-color="F4F1EF"><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">1322×876 71.4 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></p>
<p>It seems like KL is not doing a good job at approximating this toy example or am I missing something?</p>

kirk86: <p>Just to answer this, I’ve come to the conclusion that there’s sth wrong with the KL, here’s why.</p>
<p>Verbatim from <a href="https://pytorch.org/docs/1.7.1/generated/torch.nn.KLDivLoss.html#torch.nn.KLDivLoss" rel="noopener nofollow ugc">KLDivLoss</a>:</p>
<pre><code class="lang-auto">As with NLLLoss, the input given is expected to contain log-probabilities
and is not restricted to a 2D Tensor. The targets are interpreted as probabilities by default,
but could be considered as log-probabilities with log_target set to True.
</code></pre>
<p>The above plots where generated when <code>P = F.log_softmax(*)</code> and <code>Q = F.softmax(*)</code>. To this point KL is formed as <code>F.kl_div(P, Q, reduction='batchmean')</code> according to the exact definition of the docs.</p>
<p>The problem is the <code>P = F.log_softmax(*)</code> which basically destroys the approximation of Q and KL cannot recover.</p>
<p>The same problem is persistent even with the flag <code>log_target=True</code> where both P and Q are transformed via <code>F.log_softmax(*)</code>.</p>
<p>And here’s the correct result when both P and Q are just outputs from <code>F.softmax(*)</code>.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/3/f313a9240495d5a41d35717694e4bfee190e9f10.png" data-download-href="https://discuss.pytorch.org/uploads/default/f313a9240495d5a41d35717694e4bfee190e9f10" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/3/f313a9240495d5a41d35717694e4bfee190e9f10_2_345x238.png" alt="image" data-base62-sha1="yGm9bD8zm0DfyYS42JnY6ADOwMg" width="345" height="238" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/3/f313a9240495d5a41d35717694e4bfee190e9f10_2_345x238.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/3/f313a9240495d5a41d35717694e4bfee190e9f10_2_517x357.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/3/f313a9240495d5a41d35717694e4bfee190e9f10_2_690x476.png 2x" data-dominant-color="F4F3F1"><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">748×518 51.4 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>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/c/5c209c47fbccae73c4ebda352b21f9e92e863ff9.png" data-download-href="https://discuss.pytorch.org/uploads/default/5c209c47fbccae73c4ebda352b21f9e92e863ff9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/c/5c209c47fbccae73c4ebda352b21f9e92e863ff9_2_345x238.png" alt="image" data-base62-sha1="d8ZNhAwuKWeOQ1pGoWa87sc9xjz" width="345" height="238" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/c/5c209c47fbccae73c4ebda352b21f9e92e863ff9_2_345x238.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/c/5c209c47fbccae73c4ebda352b21f9e92e863ff9_2_517x357.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/c/5c209c47fbccae73c4ebda352b21f9e92e863ff9_2_690x476.png 2x" data-dominant-color="F2EFEC"><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">748×518 64.6 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></p>
<p>Can someone from the pytorch devs verify and correct this if that’s the case?</p>