Is taking the index of a maximum element and then using that element in further operations a differentiable process?

username: Lewis
cooked: <p>Hi. I have a model where I will need to use one element from a matrix of elements. I computed a softmax probability distribution over this matrix. Now, i would like to extract the element that corresponds to the maximum probability from the softmax, and then use that element further on:</p>
<pre><code class="lang-auto">_, idx = torch.max(softmax)
desired_element = matrix[idx]
</code></pre>
<p>and then the <code>desired_element</code> is used for further operations in the model.</p>
<p>Is that operation differentiable?<br>
Thanks!</p>

username: albanD
cooked: <p>Hi,</p>
<p>Gradients will flow back from <code>desired_element</code> to <code>matrix</code> as you only take one element out of a matrix which is differentiable.<br>
Gradients won’t flow back towards <code>softmax</code> though as indexing is not differentiable wrt the index and the argmax operation is not differentiable either.</p>

username: Lewis
cooked: <p>Thank you <a class="mention" href="/u/alband">@albanD</a> for your reply. So If i had some learnable weights (W) before the softmax operation as follows:</p>
<pre><code class="lang-auto">_, idx = torch.max(softmax(A.W.B))
desired_element = matrix[idx]
</code></pre>
<p>Then gradients won’t flow through W?</p>

username: albanD
cooked: <p>No they won’t.<br>
You can think about the gradients for idx as follow: how a small change in idx would change the output value? Well idx is an index, so it’s discrete. So gradients don’t make sense here <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>

username: Lewis
cooked: <p>Thank you <a class="mention" href="/u/alband">@albanD</a>  very much for your answer. What if i take a weighted sum of all the elements in matrix (softmax) but then multiplied the those weights with a mask generated based on the <code>idx</code>, as follows:</p>
<pre><code class="lang-auto">weights = softmax(A.W.B)
_, idx = torch.max(weights)
mask = torch.ones_like(weights)
mask = mask[weights&lt;idx] = 0
weighted_sum = torch.sum(weights * mask * matrix)
</code></pre>
<p>This would basically achieve the same purpose as before, because ignored elements are multiplied with 0. So Would that be differentiable?</p>

username: albanD
cooked: <p>Doing this you will get gradients back to your W, but it’s not exactly the same function as now you multiply your entry with <code>weights</code>. Not sure how stable this function will be during training though.<br>
Also in your example, I think you want the maximum value out of the max op, not the index of the maximum value.</p>

username: Lewis
cooked: <p>Hi <a class="mention" href="/u/alband">@albanD</a>  and thanks for for fast answers! I actually just want to take the maximum element from <code>matrix</code> and ignore everything else. But i still need the gradients to flow through W because I need it to learn how to choose the maximum.</p>
<pre><code class="lang-auto">outputs = softmax(A.W.B)
value, idx = torch.max(outputs)
mask = torch.ones_like(outputs)
mask[outputs&lt;value] = 0
mask[outputs&gt;=value] = 1
outputs = outputs * mask
weighted_sum = torch.sum(outputs * matrix)
</code></pre>
<p>This would basically keep only the maximum element. This would be fine I guess? We are just changing the output tensor of a specific layer.</p>

username: Brando_Miranda
cooked: <aside class="quote no-group" data-username="albanD" data-post="2" data-topic="54973">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>the argmax operation is not differentiable either</p>
</blockquote>
</aside>
<p>small caveat, argmax is not differentiable only wrt to the index.</p>

username: albanD
cooked: <p>But argmax only returns the index <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:">  Only max returns the value.</p>

username: Yongjie_Shi
cooked: <p>Meaningful explanation!!!</p>