nn.MultiheadAttention with triangular mask

username: TeDataPro
cooked: <p>Hi PyTorch family <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=12" title=":grinning:" class="emoji" alt=":grinning:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/heart.png?v=12" title=":heart:" class="emoji" alt=":heart:" loading="lazy" width="20" height="20"></p>
<p>Currently trying to implement a traditional transformer decoder with triangular mask, to train on next token prediction task, such that the triangular mask both encodes padding mask AND triangular mask at once.</p>
<p>Reading the documentation for nn.MultiheadAttention says that the parameter attn_mask should be a boolean tensor which indicates True if the attention is forbidden between 2 tokens.</p>
<p>Hence, for an embedded tokens batch with 2 samples, 6 tokens and padding for the second one, I got the following mask :</p>
<pre><code class="lang-auto">mask = torch.tensor([[[1, 0, 0, 0, 0, 0],
         [1, 1, 0, 0, 0, 0],
         [1, 1, 1, 0, 0, 0],
         [1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 0],
         [1, 1, 1, 1, 1, 1]],

        [[1, 0, 0, 0, 0, 0],
         [1, 1, 0, 0, 0, 0],
         [1, 1, 1, 0, 0, 0],
         [1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 0],
         [0, 0, 0, 0, 0, 0]]])
</code></pre>
<p>But it gives me NaN embedded token for the last one of the second sample (padded token), with the following code  :</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

embedded_tokens = torch.rand(2, 6, 512).transpose(0, 1)
decoder = nn.MultiheadAttention(512, 1)

attn_mask = ~mask.bool()
updated_tokens = decoder(embedded_tokens, embedded_tokens, embedded_tokens, attn_mask=attn_mask)
</code></pre>
<p>However, by allowing padded tokens to interact with itself, such that :</p>
<pre><code class="lang-auto">mask = torch.tensor([[[1, 0, 0, 0, 0, 0],
         [1, 1, 0, 0, 0, 0],
         [1, 1, 1, 0, 0, 0],
         [1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 0],
         [1, 1, 1, 1, 1, 1]],

        [[1, 0, 0, 0, 0, 0],
         [1, 1, 0, 0, 0, 0],
         [1, 1, 1, 0, 0, 0],
         [1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 0],
         [0, 0, 0, 0, 0, 1]]])
</code></pre>
<p>It works correctly and avoid the introduction of NaN.</p>
<p>Please, can you tell me if my implementation if correct ? If not, what should be the optimal way to implement it ?</p>
<p>Thank you guys so much for your help, cannot wait to have your answer and discuss about it with passionate PyTorch developers ! <img src="https://discuss.pytorch.org/images/emoji/apple/heart.png?v=12" title=":heart:" class="emoji" alt=":heart:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=12" title=":grinning:" class="emoji" alt=":grinning:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/smiling_face_with_three_hearts.png?v=12" title=":smiling_face_with_three_hearts:" class="emoji" alt=":smiling_face_with_three_hearts:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/heart_eyes.png?v=12" title=":heart_eyes:" class="emoji" alt=":heart_eyes:" loading="lazy" width="20" height="20"></p>