How to not loose the gradient on this operation?

jayz: <p>Hey,</p>
<p>I have a model that outputs batches with size B of “A-dimensional” softmax probabilities for C classes. For instance, the output <code>w</code> could look like this, for B=2, A=4 and C=3:</p>
<pre><code class="lang-python"># B x A x C
w = torch.softmax(torch.rand([2,4,3]),dim=2) 
w

# standard A-dimensional class probabilities, shape: B x A x C
&gt;&gt; tensor([[[0.2901, 0.4441, 0.2658],
            [0.4681, 0.2511, 0.2808],
            [0.2814, 0.3119, 0.4067],
            [0.4007, 0.3358, 0.2635]]
            [[0.2750, 0.4774, 0.2476],
            [0.3758, 0.3680, 0.2562],
            [0.3018, 0.3007, 0.3976],
            [0.3610, 0.3699, 0.2691]]])
</code></pre>
<p>where the probs sum up along the class dimension.</p>
<p>In my loss function, I need to construct the labels / targets <code>y</code>in dependence of the input <code>w</code>, i.e. <code>y = y(w)</code>, before feeding them together into <code>nn.CrossEntropyLoss()(w.transpose(1,2),y)</code>. In order to do that, I need to convert <code>w</code> to a representation that shows the class numbers instead of the softmax probabilities. I do it like this:</p>
<pre><code class="lang-python"># first step
wa = (1*(w == torch.max(w,dim=(2))[0].unsqueeze(-1))) # B x A x C
wa 

# intermediate output 
# (indicates the location of the  maximum probabilities)
&gt;&gt; tensor([[[0, 1, 0],
            [1, 0, 0],
            [0, 0, 1],
            [1, 0, 0]],
           [[0, 1, 0],
            [1, 0, 0],
            [0, 0, 1],
            [0, 1, 0]]])

# second step 
# (converts the intermediate output vectors to the predicted class labels)
wp = torch.max(wa,dim=2)[1] 
wp

# class number representation, shape: B x A (with dimension C dropped)
&gt;&gt; tensor([[1, 0, 2, 0],
           [1, 0, 2, 1]])
</code></pre>
<p>I would like to achieve that my final labels that are constructed by further processing <code>wp</code> have a gradient with respect to <code>w</code>. This will be the case if <code>wp</code> has a gradient wrt <code>w</code>. However, the second operation drops the gradient, because it asks for the max_indices (list element [1]) of <code>wa</code>, which are not differentiable, I think. Or in other words, the second operation is an argmax operation, which does not have a gradient, afaic.</p>
<p>Is there any way to rewrite the operations (i.e some smart way to circumvent the argmax operation) such that the gradient of <code>wp</code> wrt <code>w</code> is somehow preserved? I did not find one, yet.</p>
<p>Thanks!</p>
<p>Best, JZ</p>

KFrank: <p>Hi Jay!</p>
<aside class="quote no-group" data-username="jayz" data-post="1" data-topic="158483" 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/jayz/48/46742_2.png" class="avatar"> jayz:</div>
<blockquote>
<p>I would like to achieve that my final labels that are constructed by further processing <code>wp</code> have a gradient with respect to <code>w</code>.</p>
</blockquote>
</aside>
<p>The short story is that the gradient you are asking for isn’t useful because<br>
it is zero almost everywhere.</p>
<p>Some more detail:</p>
<p>The problem is that your labels, <code>wp</code>, are integers (and would have integer<br>
values even if you were to convert them to floating-point numbers).  It<br>
doesn’t really make sense to differentiate integers.</p>
<p>Consider <code>wp[0, 0]</code>.  In the example you give it starts out as <code>1</code>.  As you<br>
vary the values in <code>w</code>, <code>wp[0, 0]</code> will stay <code>1</code> for a while, and over this<br>
range the gradient will be zero (as <code>wp[0, 0]</code> is not changing).</p>
<p>Then, at some point as you vary <code>w</code>, <code>wp[0, 0]</code> will “pop” to another<br>
value (say to <code>2</code>).  Right at this point the gradient will be infinity (or,<br>
if you prefer, undefined).</p>
<p>That is, for almost all values of <code>w</code>, the gradient will be zero and will<br>
otherwise be infinite (or undefined).</p>
<p>(Because gradients of things like <code>argmax()</code> aren’t useful, pytorch<br>
generally doesn’t compute them for you.)</p>
<p>As an aside, it sounds like you want to use as a loss the cross entropy<br>
of some floating-point numbers (<code>w</code>) with (integer) class labels (<code>wp</code>)<br>
derived from those floating-point numbers.  You will have useful gradients<br>
of your loss with respect to <code>w</code> (but not <code>wp</code>).</p>
<p>I don’t understand your use case, so I can’t say whether gradients<br>
directly with respect to <code>w</code> would make sense for you, but they would<br>
be conceptually valid, and pytorch will compute them for you.</p>
<p>Best.</p>
<p>K. Frank</p>

jayz: <p>Hello K. Frank,</p>
<p>thanks for your reply!<br>
Right, I totally forgot that the integer nature of the labels prevents a creation of the gradient, on top of using the argmax. I’ll try reformulating my loss as a continuous one and transform the problem to a regression instead of a classification setting. Then, it should be possible to get the derivative of the labels wrt the input. Also, yes, the gradient of the loss wrt w is there, but for my application its crucial to have one wrt the labels, as well.</p>
<p>Best, JZ</p>