My code did not return the same result as pytorch's softmax

StuChen: <p>I need to extract Transformer’s attention weights on GPU, but my results sometimes differ slightly from PyTorch.</p>
<p>Here is my code snippet:</p>
<pre data-code-wrap="python"><code class="lang-python"># queries.shape: torch.Size([16, 3, 4, 5])
# keys.shape: torch.Size([16, 3, 4, 5])
# n_fea_hid: 15

# Compute scaled dot-product attention scores
attention_scores = torch.matmul(queries, keys.transpose(-2, -1))
attention_scores.div_(n_fea_hid**0.5)

# pytorch implementation
attention_weights1 = attention_scores.softmax(dim=3)

# My implementation
_attention_scores = attention_scores - attention_scores.max(3, keepdim=True).values
attention_scores_exp = _attention_scores.exp()
attention_weights2 = attention_scores_exp / attention_scores_exp.sum(
dim=3, keepdim=True
)
</code></pre>
<p>For the most times, <code>(attention_weights1 == attention_weights2.all() </code> returns <code>True</code>, but in rare cases, it returns <code>False</code>.</p>
<p>When return <code>False</code>, not all values are different. For example,</p>
<pre data-code-wrap="python"><code class="lang-python">attention_weights1[0,0] == attention_weights2[0,0]
</code></pre>
<p>returns</p>
<pre data-code-wrap="python"><code class="lang-python">tensor([[0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [1.4901e-08, 2.9802e-08, 2.9802e-08, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00]], device='cuda:1',
       grad_fn=&lt;SubBackward0&gt;)
</code></pre>
<p>but</p>
<pre data-code-wrap="python"><code class="lang-python">(attention_weights1[0,1] == attention_weights2[0,1]).all()
</code></pre>
<p>and</p>
<pre data-code-wrap="python"><code class="lang-python">(attention_weights1[0,2] == attention_weights2[0,2]).all()
</code></pre>
<p>return <code>True</code></p>
<p>Where do these minor deviations come from? I googled a lot but still didn’t find the reason.</p>

KFrank: <p>Hi Stu!</p>
<aside class="quote no-group quote-modified" data-username="StuChen" data-post="1" data-topic="195801" 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/s/90ced4/48.png" class="avatar"> StuChen:</div>
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">attention_weights1 = attention_scores.softmax(dim=3)
...
attention_weights2 = attention_scores_exp / attention_scores_exp.sum(
dim=3, keepdim=True
)
</code></pre>
<p>…</p>
<pre data-code-wrap="python"><code class="lang-python">tensor([[0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [1.4901e-08, 2.9802e-08, 2.9802e-08, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00]], device='cuda:1',
       grad_fn=&lt;SubBackward0&gt;)
</code></pre>
</blockquote>
</aside>
<p>These are numerical <a href="https://en.wikipedia.org/wiki/Round-off_error" rel="noopener nofollow ugc">round-off errors</a> (and are to be expected).</p>
<p>While your two computations of <code>attention_weights</code> are <em>mathematically</em><br>
equivalent, they differ numerically.  You can verify and explore this by<br>
repeating the computations using <code>torch.float64</code> (“double precision”)<br>
and you will see the differences reduced by several orders of magnitude.</p>
<p>Best.</p>
<p>K. Frank</p>

StuChen: <p>As you mentioned, changing the precision to double() will reduce this difference by several orders of magnitude.</p>
<p>Another question: if I compute these <strong>exponential numbers</strong> first using torch on GPU, based upon them, then I use numpy matrix multiplication to calculate softmax and Attention logits and then make predictions (load torch’s Linear weights under the same float32 precision), will these minor differences influence the probability distribution?</p>

KFrank: <p>Hi Stu!</p>
<aside class="quote no-group" data-username="StuChen" data-post="3" data-topic="195801" 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/s/90ced4/48.png" class="avatar"> StuChen:</div>
<blockquote>
<p>then I use numpy matrix multiplication to calculate softmax and Attention logits and then make predictions</p>
</blockquote>
</aside>
<p>Note, if you perform part of your computation with numpy instead of<br>
pytorch, you won’t be able to use autograd to backpropagate through<br>
that part of the computation (unless you write your own <code>backward()</code><br>
function to support that part of the backward pass).</p>
<p>(I don’t really understand what you’re asking here.  Why would you do<br>
things this way?  Why not use pytorch for all of the forward pass?)</p>
<aside class="quote no-group">
<blockquote>
<p>will these minor differences influence the probability distribution?</p>
</blockquote>
</aside>
<p>You would expect “these minor differences” to have a minor effect on<br>
“the probability distribution,” but you probably don’t care about such<br>
minor differences.  If they do turn out to be causing a problem (unlikely),<br>
I would recommend performing that part of the computation (or perhaps<br>
the whole computation) with pytorch in double precision (<code>torch.float64</code>).</p>
<p>Note that often when two mathematically-equivalent computations (that<br>
are numerically different) produce results that differ by round-off error<br>
near the level of machine precision, it is <em>not</em> the case that one result is<br>
more correct than the other.  They are both likely to be equally good<br>
floating-point approximations to the exact mathematical result.</p>
<p>Best.</p>
<p>K. Frank</p>