Is clamp on torch.exp is a good alternative to softmax

Came: <pre><code class="lang-python">torch.softmax
</code></pre>
<p>is stable to work on some large data. For example,</p>
<pre><code class="lang-auto">x = torch.tensor([1., 2, 150])

F.softmax(x, dim = 0)
tensor([0., 0., 1.])
</code></pre>
<p>I actually have to manually calculated the softmax where I can not directly use softmax function.<br>
I am doing:</p>
<pre><code class="lang-auto">torch.exp(x)/torch.exp(x).sum()
tensor([0., 0., nan])
</code></pre>
<p>This would further leading to nan in loss.<br>
Thus, I choose to clip the input before the exponential function.</p>
<pre><code class="lang-auto">def clamp_exp(x):
    
    return torch.exp(torch.clamp(x, max = 20))

clamp_exp(x)/clamp_exp(x).sum()
tensor([5.6028e-09, 1.5230e-08, 1.0000e+00])
</code></pre>
<p>Since I am not quite familiar with the clamp function, I am wondering is there any good way to  avoiding exploiding, will the clamp affect the backpropagation?</p>

KFrank: <p>Hi Came!</p>
<aside class="quote no-group" data-username="Came" data-post="1" data-topic="162752" 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/came/48/9868_2.png" class="avatar"> Came:</div>
<blockquote>
<p>I actually have to manually calculated the softmax where I can not directly use softmax function.</p>
</blockquote>
</aside>
<p>First, you should modify the structure of your overall computation so that<br>
you can use <code>log_softmax()</code> for this step.  Whether you write your own<br>
version or use pytorch’s, <code>log_softmax()</code> will almost certainly lead to a<br>
numerically more stable computation as it is much less likely to “explode”<br>
or underflow to zero.</p>
<p>Also, you should figure out how to structure your computation to use<br>
pytorch’s version (whether it be <code>log_softmax()</code> or <code>softmax()</code>), rather<br>
than reinvent the wheel by writing your own.  Why do you need to write<br>
your own version?  If worst were to come to worst, couldn’t you just<br>
implement “your own version” by writing a wrapper around pytorch’s?</p>
<p>Best.</p>
<p>K. Frank</p>

Came: <p>Hi, Frank,<br>
I appreciate your kind replies.<br>
I was trying to implement some neighborhood attention (like Graph Attention Network) where softmax is not directly be able to applied. You have to use scatter_add  ( segment sum in tensor flow) and calculate the denominator manually.</p>
<p>I found folks try to subtract the maximum value to make exponential stable:<br>
<a href="https://github.com/gordicaleksa/pytorch-GAT/blob/39c8f0ee634477033e8b1a6e9a6da3c7ed71bbd1/models/definitions/GAT.py#L262" rel="noopener nofollow ugc">https://github.com/gordicaleksa/pytorch-GAT/blob/39c8f0ee634477033e8b1a6e9a6da3c7ed71bbd1/models/definitions/GAT.py#L262</a></p>
<aside class="onebox stackexchange" data-onebox-src="https://stats.stackexchange.com/questions/338285/how-does-the-subtraction-of-the-logit-maximum-improve-learning">
  <header class="source">

      <a href="https://stats.stackexchange.com/questions/338285/how-does-the-subtraction-of-the-logit-maximum-improve-learning" target="_blank" rel="noopener nofollow ugc">stats.stackexchange.com</a>
  </header>

  <article class="onebox-body">
      <a href="https://stats.stackexchange.com/users/176100/jav" target="_blank" rel="noopener nofollow ugc">
    <img alt="Jav" src="https://www.gravatar.com/avatar/bc4d4fbae28ecd225bd84e7e8cb6fc5e?s=256&amp;d=identicon&amp;r=PG" class="thumbnail onebox-avatar" width="256" height="256">
  </a>

<h4>
  <a href="https://stats.stackexchange.com/questions/338285/how-does-the-subtraction-of-the-logit-maximum-improve-learning" target="_blank" rel="noopener nofollow ugc">How does the subtraction of the logit maximum improve learning?</a>
</h4>

<div class="tags">
  <strong>deep-learning, logit, tensorflow, keras</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stats.stackexchange.com/users/176100/jav" target="_blank" rel="noopener nofollow ugc">
    Jav
  </a>
  on <a href="https://stats.stackexchange.com/questions/338285/how-does-the-subtraction-of-the-logit-maximum-improve-learning" target="_blank" rel="noopener nofollow ugc">12:45PM - 03 Apr 18 UTC</a>
</div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>