Calculate p pair-wise KL divergence

Hoang_Phan: <p>Hi everyone, is there any way to efficiently calculate pair-wise KL divergence between 2 sets of samples A (batch_size x dimension) and B (batch_size x dimension), which returns a tensor of (batch_size x batch_size) in Pytorch. It’s similar to torch.cdist but with KL divergence rather than p-norm distance? Thank you so much.</p>

KFrank: <p>Hi Hoang!</p>
<aside class="quote no-group" data-username="Hoang_Phan" data-post="1" data-topic="131424" data-full="true">
<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/hoang_phan/48/31567_2.png" class="avatar"> Hoang_Phan:</div>
<blockquote>
<p>is there any way to efficiently calculate pair-wise KL divergence between 2 sets of samples A (batch_size x dimension) and B (batch_size x dimension)</p>
</blockquote>
</aside>
<p>Yes, it can be done using pytorch tensor operations (and without loops),<br>
so it should be efficient.  (However, I don’t see anyway to do this using<br>
pytorch’s <code>kl_div()</code>, itself.)</p>
<p>The idea is that you need get a batch-wise “outer product” and <code>einsum()</code><br>
seems to me to be the most straightforward way to achieve this.</p>
<p>Here we check a loopless implementation against using <code>kl_div()</code><br>
with loops:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2021)
&gt;&gt;&gt;
&gt;&gt;&gt; A = torch.randn (3, 5)
&gt;&gt;&gt; B = torch.rand (3, 5)
&gt;&gt;&gt;
&gt;&gt;&gt; A
tensor([[ 2.2871,  0.6413, -0.8615, -0.3649, -0.6931],
        [ 0.9023, -2.7183, -1.4478,  0.6238,  0.4822],
        [-2.3055,  0.9176,  1.5674, -0.1284, -1.0042]])
&gt;&gt;&gt; B
tensor([[0.8303, 0.5216, 0.7438, 0.8290, 0.0219],
        [0.0813, 0.0172, 0.1464, 0.7492, 0.9450],
        [0.6737, 0.1135, 0.7421, 0.7810, 0.9446]])
&gt;&gt;&gt;
&gt;&gt;&gt; batch_klA = torch.zeros (3, 3)
&gt;&gt;&gt;
&gt;&gt;&gt; for  i in range (3):
...     for  j in range (3):
...         batch_klA[i, j] = torch.nn.functional.kl_div (A[i], B[j], reduction = 'sum')
...
&gt;&gt;&gt; batch_klB = (B * B.log()).sum (dim = 1) - torch.einsum ('ik, jk -&gt; ij', A, B)
&gt;&gt;&gt;
&gt;&gt;&gt; batch_klA
tensor([[-2.2283,  0.0326, -1.0158],
        [ 0.2646, -1.5627, -1.1489],
        [-0.5549,  0.1626,  0.3534]])
&gt;&gt;&gt; batch_klB
tensor([[-2.2283,  0.0326, -1.0158],
        [ 0.2646, -1.5627, -1.1489],
        [-0.5549,  0.1626,  0.3534]])
&gt;&gt;&gt;
&gt;&gt;&gt; torch.allclose (batch_klA, batch_klB)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Hoang_Phan: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>, when I tried the above implementation with A=B, it returned a tensor containing all nans while we need a tensor with a zero-diagonal tensor, I’m still trying to figure out the problem</p>

KFrank: <p>Hi Hoang!</p>
<aside class="quote no-group" data-username="Hoang_Phan" data-post="4" data-topic="131424" data-full="true">
<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/hoang_phan/48/31567_2.png" class="avatar"> Hoang_Phan:</div>
<blockquote>
<p>when I tried the above implementation with A=B, it returned a tensor containing all nans</p>
</blockquote>
</aside>
<p>Please note that both the implementation I posted above, as well as,<br>
by default, pytorch’s <code>kl_div()</code>, expect the <code>input</code> (<code>A</code> in the example<br>
above) to be log-probabilities and the <code>target</code> (<code>B</code>) to be probabilities.<br>
If you pass in a probability that is zero or negative, you will get a <code>nan</code>.</p>
<p>If this doesn’t clear up your issue, please post a short, complete, runnable<br>
example that reproduces the issue you see.</p>
<p>Best.</p>
<p>K. Frank</p>