Efficient implementation of Jacobian of Softmax

Samue1: <p>Given a one-dimensional input tensor <code>S</code>, I want to evaluate the following expression:</p>
<p><code>J_ij = S_i(delta_ij - S_j)</code></p>
<p>where <code>delta_ij</code> represents the Kronecker delta. The result <code>J</code> of this expression is a square matrix. I would like to know how I can efficiently evaluate this expression using PyTorch? My current implementation is very slow and looks as follows:</p>
<pre><code>S = torch.arange(start=1, end=10)
I = torch.eye(S.size()[0])
S = S.repeat(S.size()[0], 1)
J = S.T * (I - S)
</code></pre>
<p>I would be very happy if someone could give me a hint.</p>

KFrank: <p>Hi Samuel!</p>
<aside class="quote no-group" data-username="Samue1" data-post="1" data-topic="125369" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/f04885/48.png" class="avatar"> Samue1:</div>
<blockquote>
<p><code>J_ij = S_i(delta_ij - S_j)</code></p>
</blockquote>
</aside>
<p>The following expression is more compact and likely faster:</p>
<pre data-code-wrap="python"><code class="lang-python">J = torch.diag (S) - torch.outer (S, S)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Samue1: <p>Does this also work for batched versions of <code>S</code>? For example, if <code>S</code> has dimensions <code>(batch_size, num_outputs)</code>? Is the result correct if I use</p>
<p><code>J = torch.diag_embed(S) - torch.outer(S, S)</code></p>

KFrank: <p>Hi Samuel!</p>
<aside class="quote no-group" data-username="Samue1" data-post="3" data-topic="125369" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/f04885/48.png" class="avatar"> Samue1:</div>
<blockquote>
<p>Does this also work for batched versions of <code>S</code>?</p>
</blockquote>
</aside>
<p>No.  If you had tried it, you would have discovered that <code>torch.outer()</code><br>
does not accept multidimensional tensors.</p>
<aside class="quote no-group">
<blockquote>
<p>Is the result correct if I use</p>
<p><code>J = torch.diag_embed(S) - torch.outer(S, S)</code></p>
</blockquote>
</aside>
<p>No, this will throw an error (because you pass a multidimensional<br>
tensor to <code>torch.outer()</code>).</p>
<p>You can, however, use pytorch’s swiss army knife of tensor<br>
multiplication functions to construct a batch version of outer:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt; S = torch.arange (6).reshape (2, 3).float()
&gt;&gt;&gt; S
tensor([[0., 1., 2.],
        [3., 4., 5.]])
&gt;&gt;&gt; torch.diag_embed (S) - torch.einsum ('ij, ik -&gt; ijk', S, S)
tensor([[[  0.,   0.,   0.],
         [  0.,   0.,  -2.],
         [  0.,  -2.,  -2.]],

        [[ -6., -12., -15.],
         [-12., -12., -20.],
         [-15., -20., -20.]]])
</code></pre>
<p>(As an aside, none of this has anything to do with the title you gave<br>
this thread, namely “Jacobian of Softmax.”)</p>
<p>Best.</p>
<p>K. Frank</p>