Jacobian of Cross Entropy Loss

jean_b: <p>I’m trying to find the Jacobian of nn.CrossEntropyLoss, where my inputs are some network output and target in the form of class indeces. I tired doing a simple test case as follows:</p>
<pre data-code-wrap="import"><code class="lang-nohighlight">import torch.nn as nn
import torch.autograd.functional

loss = nn.CrossEntropyLoss()

input = torch.randn(3, 5, requires_grad=True)
target = torch.empty(3, dtype=torch.long).random_(5)

jb = torch.autograd.functional.jacobian(loss, (input, target))
</code></pre>
<p>This throws RuntimeError: only Tensors of floating point dtype can require gradients</p>
<p>I tired converting target to a float tensor but that doesn’t work with the Cross Entropy, as it wants target to be either int or long. Is there any way around this? I’m new to pytorch so I apologize if I’m thinking about this completely wrong. Overall I want to be able to do forward mode AD on the loss so that I can do a directional derivative/jacobian vector product in the direction of some vector v, or in this case (since Cross Entropy outputs a scalar) the magnitude of the projection of the gradient along v.</p>

KFrank: <p>Hi Jean!</p>
<aside class="quote no-group" data-username="jean_b" data-post="1" data-topic="153858" 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/j/7c8e57/48.png" class="avatar"> jean_b:</div>
<blockquote>
<p>I’m trying to find the Jacobian of nn.CrossEntropyLoss</p>
</blockquote>
</aside>
<p>The Jacobian matrix is the matrix of partial derivatives of a number of<br>
results (the vector-valued result) of a function with respect to a number<br>
of inputs (a vector-valued input) to that function.</p>
<p>Could you specify which partial derivatives you want to compute, that is<br>
the derivatives of which specific results with respect to which specific<br>
inputs?</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="import"><code class="lang-import">target = torch.empty(3, dtype=torch.long).random_(5)

jb = torch.autograd.functional.jacobian(loss, (input, target))
</code></pre>
</blockquote>
</aside>
<p>Your code corresponds to the most common use case of <code>CrossEntropyLoss</code>,<br>
namely, where the <code>target</code> is a batch of integer class labels.  It doesn’t make<br>
sense to differentiate something with respect to an integer, because it doesn’t<br>
make sense to vary an integer value an infinitesimal amount away from being<br>
an integer.</p>
<p>In your Jacobian computation, are you intending to differentiate something<br>
with respect to your integer <code>target</code>?  If so, what is that supposed to mean,<br>
conceptually and mathematically?</p>
<p>Best.</p>
<p>K. Frank</p>

jean_b: <p>I was just looking to find the gradient of loss with respect to input, like would normally be done with backprop using loss.backward(). My hope was to be able to do it via forward mode (for various reasons) but since loss takes both input and target I wasn’t sure if there was an easy way to split it up or get around this error message.</p>

KFrank: <p>Hi Jean!</p>
<p>First a general answer, and then a work-around for your use case:</p>
<aside class="quote no-group quote-modified" data-username="jean_b" data-post="3" data-topic="153858" 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/j/7c8e57/48.png" class="avatar"> jean_b:</div>
<blockquote>
<p>I was just looking to find the gradient of loss with respect to input,<br>
…<br>
but since loss takes both input and target I wasn’t sure if there was an easy way to split it up</p>
</blockquote>
</aside>
<p>For context:</p>
<aside class="quote no-group" data-username="jean_b" data-post="1" data-topic="153858">
<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/j/7c8e57/48.png" class="avatar"> jean_b:</div>
<blockquote>
<pre><code class="lang-auto">input = torch.randn(3, 5, requires_grad=True)
target = torch.empty(3, dtype=torch.long).random_(5)

jb = torch.autograd.functional.jacobian(loss, (input, target))
</code></pre>
</blockquote>
</aside>
<p>The problem is that you want to apply <code>jacobian()</code> to a function one of<br>
whose arguments (<code>target</code>) is an integer tensor – not differentiable – hence<br>
<code>jacobian()</code> complains.  But you only want gradients of the floating-point,<br>
differentiable argument (<code>input</code>), which is, in principle, okay.</p>
<p>As a general approach, define a function-object class.  When you instantiate<br>
it, pass it <code>target</code>, and have it store <code>target</code> as a property.  Have its <code>__call__</code><br>
method only take <code>input</code> (and <code>self</code>) as an argument.  Now <code>jacobian()</code> will<br>
only try to calculate gradients of the floating-point <code>input</code>.</p>
<p>Thus:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2022)

class CrossEntropyWrapper:
    def __init__ (self, target):
        self.target = target
    def __call__ (self, input):
        return  torch.nn.functional.cross_entropy (input, self.target)

input = torch.randn (3, 5, requires_grad = True)
target = torch.empty (3, dtype = torch.long).random_ (5)

wrapper = CrossEntropyWrapper (target)
jbA = torch.autograd.functional.jacobian (wrapper, input)

print (jbA)
</code></pre>
<p>resulting in:</p>
<pre><code class="lang-plaintext">1.11.0
tensor([[-0.2766,  0.0651,  0.0589,  0.1144,  0.0382],
        [-0.3239,  0.0023,  0.1199,  0.1836,  0.0181],
        [ 0.0190,  0.0239,  0.0961,  0.1055, -0.2445]])
</code></pre>
<p>For your specific use case – because (as of a few versions ago)<br>
<code>CrossEntropyLoss</code> can also take floating-point “soft” targets (of<br>
a different shape) – you can let <code>jacobian()</code> calculate the gradients<br>
of both <code>input</code> and a floating-point version of <code>target</code>, and you can<br>
just ignore the gradients of <code>target</code>.</p>
<p>Continuing:</p>
<pre data-code-wrap="python"><code class="lang-python">loss = torch.nn.CrossEntropyLoss()

target_onehot_float = torch.nn.functional.one_hot (target, 5).float()

jbB = torch.autograd.functional.jacobian(loss, (input, target_onehot_float))

print (jbB)
</code></pre>
<p>with the result:</p>
<pre><code class="lang-plaintext">(tensor([[-0.2766,  0.0651,  0.0589,  0.1144,  0.0382],
        [-0.3239,  0.0023,  0.1199,  0.1836,  0.0181],
        [ 0.0190,  0.0239,  0.0961,  0.1055, -0.2445]]), tensor([[0.5905, 0.5442, 0.5775, 0.3565, 0.7225],
        [1.1887, 1.6564, 0.3409, 0.1988, 0.9707],
        [0.9557, 0.8783, 0.4147, 0.3834, 0.4406]]))
</code></pre>
<p>Best.</p>
<p>K. Frank</p>