Gradual softmax?

username: ZimoNitrome
cooked: <p>Given a tensor of values in the range [0, 1], multiplying these values with a scalar <code>p</code> and applying a softmax gives scaled probabilities that sum to 1. Increasing <code>p</code> pushes the values to either 0 or 1. See example:</p>
<pre data-code-wrap="py"><code class="lang-py">values = torch.tensor([0.00, 0.25, 0.50, 0.75, 1.00])

softmax_values = torch.stack([torch.softmax(p*values, 0) for p in np.arange(0, 20)])

for input, outputs in zip(values, softmax_values.T):
    c = plt.gca()._get_lines.get_next_color()
    plt.plot([0], [input], 'o', c=c)
    plt.plot(outputs, c=c)

plt.xlabel("p")
    
plt.show()
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/e/deddd217f72151d0579220aa9b6456171aea7304.png" alt="download (35)" data-base62-sha1="vNzfjF8H8Ok2JLJ1NcAtJSGdimM" width="372" height="262"></p>
<p>Dots represent the input values and lines show softmax values of those values given a certain <code>p</code>.</p>
<p>Setting low values of <code>p</code> results in all softmax values converge to 1 / n_values = 1 / 5 = 0.2.</p>
<p>How do I instead create a function that gradually transform the input values to their softmax as p → inf?</p>
<p>That function would look like:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/8/880f0fd68cd1eb41064fdd7bc61a52a4b7783a2f.png" alt="download (36)" data-base62-sha1="jpDcLpMw2niaLJhFShGnBZIzokL" width="372" height="262"></p>

username: KFrank
cooked: <p>Hi Zimo!</p>
<aside class="quote no-group" data-username="ZimoNitrome" data-post="1" data-topic="121075" 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/zimonitrome/48/27591_2.png" class="avatar"> ZimoNitrome:</div>
<blockquote>
<p>How do I instead create a function that gradually transform the input values to their softmax as p → inf?</p>
</blockquote>
</aside>
<p>There are lots of possibilities.  (You could choose between them by<br>
imposing further conditions on the behavior of your function.)</p>
<p>Here is an illustration of one possible choice:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.7.1'
&gt;&gt;&gt; _ = torch.manual_seed (2021)
&gt;&gt;&gt; def gradual (values, p):
...     return torch.where (values == values.max(), 1.0 - (1.0 - values)**p, values**p)
...
&gt;&gt;&gt; values = torch.rand (5)
&gt;&gt;&gt; values
tensor([0.1304, 0.5134, 0.7426, 0.7159, 0.5705])
&gt;&gt;&gt; gradual (values, 1)
tensor([0.1304, 0.5134, 0.7426, 0.7159, 0.5705])
&gt;&gt;&gt; gradual (values, 2)
tensor([0.0170, 0.2636, 0.9337, 0.5125, 0.3254])
&gt;&gt;&gt; gradual (values, 3)
tensor([0.0022, 0.1353, 0.9829, 0.3669, 0.1857])
&gt;&gt;&gt; gradual (values, 5)
tensor([3.7761e-05, 3.5664e-02, 9.9887e-01, 1.8803e-01, 6.0419e-02])
&gt;&gt;&gt; gradual (values, 10)
tensor([1.4259e-09, 1.2719e-03, 1.0000e+00, 3.5355e-02, 3.6504e-03])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: ZimoNitrome
cooked: <p>Thanks for the answer K. Frank.</p>
<p>I will try to tinker with this solution in a few hours.</p>
<p>Do you happen to know if <code>torch.where</code> breaks the gradent tree? I am looking for a differentiable function even if I know that for large values of <code>p</code>, the gradient will be strangled.</p>

username: KFrank
cooked: <p>Hi Zimo!</p>
<aside class="quote no-group" data-username="ZimoNitrome" data-post="3" data-topic="121075" 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/zimonitrome/48/27591_2.png" class="avatar"> ZimoNitrome:</div>
<blockquote>
<p>Do you happen to know if <code>torch.where</code> breaks the gradent tree?</p>
</blockquote>
</aside>
<p><code>torch.where()</code> does not break the computation graph – that’s a good<br>
thing about it.</p>
<p>But, furthermore, in this specific case, each individual element of <code>values</code><br>
only participates in a single the branch of the <code>torch.where()</code> function.<br>
I’m just using it as a more compact way to single out the largest element<br>
in <code>values</code>, rather than indexing with <code>argmax()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

username: ZimoNitrome
cooked: <p>Yes it seems to work quite nicely <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<pre data-code-wrap="python"><code class="lang-python">values = torch.tensor([0.20, 0.40, 0.60, 0.80], requires_grad=True)

def gradual (values, p):
    return torch.where (values == values.max(), 1.0 - (1.0 - values)**p, values**p)


softmax_values = torch.stack([gradual(values, p) for p in np.arange(1, 20)])

for input, outputs in zip(values, softmax_values.T):
    c = plt.gca()._get_lines.get_next_color()
    plt.plot([0], [input], 'o', c=c)
    plt.plot(outputs.detach(), c=c)
    
plt.xlabel("p")
plt.show()

loss = softmax_values.mean()
loss.backward()
values.grad
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/4/045e4533ef30d0e3309e0db3ce78e4bebdefd771.png" alt="download (37)" data-base62-sha1="CDSKHWLpAYnPMxaXK8zLMRTbuF" width="372" height="262"></p>
<p><code>tensor([0.0206, 0.0365, 0.0822, 0.0206])</code></p>
<p>Do you know if that is the case generally, that using <code>torch.where(x == max(x), t1, t2)</code> is differentiable while <code>torch.argmax(x)</code> is not? I guess they have different use-cases but this feels like a very nice “trick” to maintain gradient.</p>