Why denominator in multi-head attention in PyTorch's implementation different from most proposed structure?

annisat: <p>As the title describes. I just compared the code of transformers from hugging face’s Bert and torch.nn.Transformer, directly from the code I found in my site-packages. Specifically, I’m looking for the part of dividing attention score before softmax. Most paper on transformer mentioned this:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/8/b/8bb13b69d63f73d53c1a91d0d250be2d65ad2d43.png" alt="image" data-base62-sha1="jVM5m2T904Bw6ilgJyHySMkkFuX" width="597" height="127"></p>
<p>In hugging face’s Bert, I believe it is this line of code, which implement this quite right.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/huggingface/transformers/blob/e4b234834a79541f31be227aadce13f5aafda85a/src/transformers/models/bert/modeling_bert.py#L324">
  <header class="source">

      <a href="https://github.com/huggingface/transformers/blob/e4b234834a79541f31be227aadce13f5aafda85a/src/transformers/models/bert/modeling_bert.py#L324" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/huggingface/transformers/blob/e4b234834a79541f31be227aadce13f5aafda85a/src/transformers/models/bert/modeling_bert.py#L324" target="_blank" rel="noopener nofollow ugc">huggingface/transformers/blob/e4b234834a79541f31be227aadce13f5aafda85a/src/transformers/models/bert/modeling_bert.py#L324</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="314" style="counter-reset: li-counter 313 ;">
          <li>    positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility</li>
          <li></li>
          <li>    if self.position_embedding_type == "relative_key":</li>
          <li>        relative_position_scores = torch.einsum("bhld,lrd-&gt;bhlr", query_layer, positional_embedding)</li>
          <li>        attention_scores = attention_scores + relative_position_scores</li>
          <li>    elif self.position_embedding_type == "relative_key_query":</li>
          <li>        relative_position_scores_query = torch.einsum("bhld,lrd-&gt;bhlr", query_layer, positional_embedding)</li>
          <li>        relative_position_scores_key = torch.einsum("bhrd,lrd-&gt;bhlr", key_layer, positional_embedding)</li>
          <li>        attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key</li>
          <li></li>
          <li class="selected">attention_scores = attention_scores / math.sqrt(self.attention_head_size)</li>
          <li>if attention_mask is not None:</li>
          <li>    # Apply the attention mask is (precomputed for all layers in BertModel forward() function)</li>
          <li>    attention_scores = attention_scores + attention_mask</li>
          <li></li>
          <li># Normalize the attention scores to probabilities.</li>
          <li>attention_probs = nn.functional.softmax(attention_scores, dim=-1)</li>
          <li></li>
          <li># This is actually dropping out entire tokens to attend to, which might</li>
          <li># seem a bit unusual, but is taken from the original Transformer paper.</li>
          <li>attention_probs = self.dropout(attention_probs)</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>I believe the closest thing in PyTorch is this line, which is actually dividing query, not the result of dot product.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/6b0b088c6caef54d35ad04defe07a4a94831dd27/torch/nn/functional.py#L4999">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/6b0b088c6caef54d35ad04defe07a4a94831dd27/torch/nn/functional.py#L4999" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/6b0b088c6caef54d35ad04defe07a4a94831dd27/torch/nn/functional.py#L4999" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/6b0b088c6caef54d35ad04defe07a4a94831dd27/torch/nn/functional.py#L4999</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="4989" style="counter-reset: li-counter 4988 ;">
          <li>        and E is embedding dimension.</li>
          <li>    - value: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,</li>
          <li>        and E is embedding dimension.</li>
          <li>    - attn_mask: either a 3D tensor of shape :math:`(B, Nt, Ns)` or a 2D tensor of</li>
          <li>        shape :math:`(Nt, Ns)`.</li>
          <li></li>
          <li>    - Output: attention values have shape :math:`(B, Nt, E)`; attention weights</li>
          <li>        have shape :math:`(B, Nt, Ns)`</li>
          <li>"""</li>
          <li>B, Nt, E = q.shape</li>
          <li class="selected">q = q / math.sqrt(E)</li>
          <li># (B, Nt, E) x (B, E, Ns) -&gt; (B, Nt, Ns)</li>
          <li>attn = torch.bmm(q, k.transpose(-2, -1))</li>
          <li>if attn_mask is not None:</li>
          <li>    attn += attn_mask</li>
          <li>attn = softmax(attn, dim=-1)</li>
          <li>if dropout_p &gt; 0.0:</li>
          <li>    attn = dropout(attn, p=dropout_p)</li>
          <li># (B, Nt, Ns) x (B, Ns, E) -&gt; (B, Nt, E)</li>
          <li>output = torch.bmm(attn, v)</li>
          <li>return output, attn</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>I’d like to know if they’re really equivalent? Wouldn’t that change the amount of backpropagation into Wq and Qk?<br>
I know training transformer from scratch is hard. I’m having trouble training torch.nn.Transformer from scratch in a recent project. But could this be a reason why the loss won’t go down?</p>

Matias_Vasquez: <p>Since 𝑑_ℎ represents a scalar, it should not matter if you do it before or after the matrix multiplication.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/8/487eabaf13b3fe5e84ea9cbbfea52626de6deb9d.png" alt="image" data-base62-sha1="aljLDZhIXsw7XdMg5D970Ividcx" width="307" height="137"></p>
<p>Of course due to resolution you will get slightly different results, but I do not think that this should be enough reason for your loss to not go down. Here is a small example to show that there are differences, but they are very small.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import math

q = torch.rand(1, 5, 30)
k = torch.rand(1, 5, 30)
v = torch.rand(1, 5, 30)


B, Nt, E = q.shape
q_ = q / math.sqrt(E)
attn1 = torch.bmm(q_, k.transpose(-2, -1))
attn2 = torch.bmm(q, k.transpose(-2, -1)) / math.sqrt(E)

print(f"Attn1: {attn1}")
print(f"Attn2: {attn2}")
print(f"Compa: {attn1==attn2}")
</code></pre>
<pre data-code-wrap="python"><code class="lang-python"># Output:
Attn1: tensor([[[1.4470, 1.2312, 1.4192, 1.0483, 1.0938],
         [1.5067, 1.2858, 1.3447, 1.3185, 1.2335],
         [1.3196, 1.1793, 1.3776, 1.1048, 1.0500],
         [1.2411, 1.1090, 1.3351, 0.9647, 0.9519],
         [1.6056, 1.5329, 1.6754, 1.3193, 1.4312]]])
Attn2: tensor([[[1.4470, 1.2312, 1.4192, 1.0483, 1.0938],
         [1.5067, 1.2858, 1.3447, 1.3185, 1.2335],
         [1.3196, 1.1793, 1.3776, 1.1048, 1.0500],
         [1.2411, 1.1090, 1.3351, 0.9647, 0.9519],
         [1.6056, 1.5329, 1.6754, 1.3193, 1.4312]]])
Compa: tensor([[[ True, False, False, False,  True],
         [False,  True,  True,  True, False],
         [ True, False, False, False,  True],
         [False,  True,  True, False, False],
         [False, False, False,  True, False]]])
</code></pre>
<p>Transformers do need a lot of training data and a lot of time to train. Would it be possible for you to take a pretrained model and fine tune it to your specific task?</p>