Discrepancy Between key_padding_mask and attn_mask in MultiheadAttention Layer

evaggelos: <p>Hello,</p>
<p>I am working with the <code>MultiheadAttention</code> layer in PyTorch and encountered a discrepancy between using <code>key_padding_mask</code> and <code>attn_mask</code> for handling variable length sequences with padding. My goal is to ensure that padded positions do not influence the attention scores. However, the results differ when using <code>key_padding_mask</code> versus an equivalent <code>attn_mask</code>.</p>
<p>Here’s a minimal reproducible example:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

# Function to create key padding mask from padded sequences
def create_key_padding_mask(embeddings, padding_value=0):
    return (embeddings == padding_value).all(dim=-1)

# Function to create an attention mask from key padding mask
def create_attn_mask_from_key_padding_mask(key_padding_mask, num_heads):
    batch_size, seq_length = key_padding_mask.shape
    attn_mask = torch.ones((batch_size, seq_length, seq_length), dtype=torch.bool)
    for i in range(batch_size):
        valid_len = seq_length - key_padding_mask[i].sum().item()
        attn_mask[i, :valid_len, :valid_len] = False
    
    attn_mask = attn_mask.unsqueeze(1).expand(batch_size, num_heads, seq_length, seq_length)
    attn_mask = attn_mask.reshape(batch_size * num_heads, seq_length, seq_length)
    return attn_mask

# Create a batch of embeddings with variable lengths (padded with 0)
batch_size = 5
seq_length = 6
embed_dim = 100

# Simulate embeddings with padding (random example)
embeddings = torch.rand(batch_size, seq_length, embed_dim)
embeddings[0, 4:] = 0  # First sequence is of length 4
embeddings[1, 3:] = 0  # Second sequence is of length 3
embeddings[2, 5:] = 0  # Third sequence is of length 5
embeddings[3, 2:] = 0  # Fourth sequence is of length 2
embeddings[4, 1:] = 0  # Fifth sequence is of length 1

# Create the key padding mask
key_padding_mask = create_key_padding_mask(embeddings)

# Create the attention mask from key padding mask
num_heads = 10
attn_mask = create_attn_mask_from_key_padding_mask(key_padding_mask, num_heads)

# Define the MultiheadAttention layer
multihead_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)

# Perform multi-head attention using the key padding mask
attn_output_key_mask, attn_output_weights_key_mask = multihead_attn(
    embeddings, embeddings, embeddings, key_padding_mask=key_padding_mask)

# Perform multi-head attention using the attention mask
attn_output_attn_mask, attn_output_weights_attn_mask = multihead_attn(
    embeddings, embeddings, embeddings, attn_mask=attn_mask)

# Compare the outputs
print("Attention Output using Key Padding Mask:")
print(attn_output_key_mask)

print("Attention Output using Attention Mask:")
print(attn_output_attn_mask)

# Check if the outputs are similar
print("Are the outputs similar?")
print(torch.allclose(attn_output_key_mask, attn_output_attn_mask, atol=1e-6))
</code></pre>
<p>I suspect that when using key_padding_mask the last padded_columns (p) are correctly -inf in the final mask but the last p rows are not -inf, but zeros. So when the softmax is applied on dim=1 after, the last columns are 0 (which is correct) but the last rows are 1/(non_padded_items) because we have (non_padded_items zeros followed by p -inf’s)<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/9/1/91d0c88dd639cffa051e00c4ae262a46d9da5a7c.png" alt="disc" data-base62-sha1="kNWy4J9OUeajClGTjjsIMluJGJm" width="552" height="394"></p>
<p>Any insights or guidance on resolving this discrepancy would be greatly appreciated.</p>
<p>Thank you!</p>

vdw: <p>I think(!) that it actually does not matter what the values for the padding entries are since any follow-up layer should ignore those values anyway. At least this was my take-away message when I asked about <a href="https://stats.stackexchange.com/questions/598239/how-is-padding-masking-considered-in-the-attention-head-of-a-transformer" rel="noopener nofollow ugc">this</a> as well a while ago.</p>
<p>But I would be curious if anyone can confirm or clarify this.</p>

evaggelos: <p>Hello Chris,</p>
<p>Thank you for your answer. I see in the <code>multi_head_attention_forward</code> implementation that the output directly goes into the linear layer without using padding. So, the final 𝑝 rows of the result (which is the first p rows of the input, averaged, due to the uniform attention) pass through the network.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/9/d98bb048c319549c159b21ed7ec1a08b51455e52.png" data-download-href="https://discuss.pytorch.org/uploads/default/d98bb048c319549c159b21ed7ec1a08b51455e52" title="multi_head_att"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d98bb048c319549c159b21ed7ec1a08b51455e52_2_690x172.png" alt="multi_head_att" data-base62-sha1="v2uTkYLx9YXm9xU3qus1eZLWYbU" width="690" height="172" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d98bb048c319549c159b21ed7ec1a08b51455e52_2_690x172.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d98bb048c319549c159b21ed7ec1a08b51455e52_2_1035x258.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/9/d98bb048c319549c159b21ed7ec1a08b51455e52.png 2x" data-dominant-color="F7F6F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">multi_head_att</span><span class="informations">1066×266 42.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Can someone confirm if something happens under the hood that I might be missing and that solves this? Also, in the TransformerEncoderLayer, the output of <code>_sa_block</code> passes to <code>_ff_block</code> without padding (I had to create a hook to pad it manually). Is that right?</p>

evaggelos: <p>Hello Mel,</p>
<p>Thanks for your response. The correct mask is the one I was creating, which also has the last two rows set to -inf instead of 0. If you look closely at your output, you’ll see that the last 2 (=number of padded items) rows are the same and equal to the mean of the first (non-padded) rows. I was asking if the last two rows are handled internally, because they shouldn’t pass through the linear layer without being zeroed first. But if you see the <code>multi_head_attention_forward</code> code they pass without first zeroed</p>

konan009: <p>ahh yeah nice catch. For me, It seems to be a bug, the padded tokens should be set again to zeros after the forward so that it wont affect the further parameters especially if you have multiple self-attention layer or a feed-forward layer. The current workaround is just to set the padding tokens to zeroes after the forward.Good Observation <a class="mention" href="/u/evaggelos">@evaggelos</a> <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"> Regards</p>

konan009: <p>Looking at the <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L5518" rel="noopener nofollow ugc">code</a> it was then projected to different set of parameters.</p>

evaggelos: <p>Exactly! This shouldn’t be happening without first zeroing out, right?</p>

evaggelos: <p>Hello, do we have any updates on this topic? <a class="mention" href="/u/ptrblck">@ptrblck</a>, I apologize for the random tag, but I’ve noticed that you are very active and helpful. I was wondering if you might have any information on this subject.</p>

gabrielff: <p>Hi <a class="mention" href="/u/evaggelos">@evaggelos</a>,<br>
The function create_attn_mask_from_key_padding_mask is not doing the expected attn_mask.<br>
You can compare with how they do at <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L6179" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L6179</a><br>
where they first create the dimensions, expand, reshape, and later broadcast towards attn_mask:</p>
<pre><code class="lang-auto">    # merge key padding and attention masks
    if key_padding_mask is not None:
        assert key_padding_mask.shape == (
            bsz,
            src_len,
        ), f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
        key_padding_mask = (
            key_padding_mask.view(bsz, 1, 1, src_len)
            .expand(-1, num_heads, -1, -1)
            .reshape(bsz * num_heads, 1, src_len)
        )
        if attn_mask is None:
            attn_mask = key_padding_mask
        else:
            attn_mask = attn_mask + key_padding_mask
</code></pre>
<p>So, you can come up with a function like:</p>
<pre><code class="lang-auto">def pad_to_attn_mask(pad_mask, num_heads):
    B, L, H = *pad_mask.shape, num_heads
    return pad_mask.view(B, 1, 1, L).expand(-1, H, L, -1).reshape(B * H, L, L)
</code></pre>
<p>And compare attention masks with a small example:</p>
<pre><code class="lang-auto">num_heads = 2
key_padding_mask_small = torch.tensor([[0, 0, 0, 1, 1], [0, 0, 0, 0, 1]], dtype=bool)

attn_mask = pad_to_attn_mask(key_padding_mask_small, num_heads)
attn_mask_ = create_attn_mask_from_key_padding_mask(key_padding_mask_small, num_heads)
</code></pre>
<p>which gives:</p>
<pre><code class="lang-auto">attn_mask
tensor([[[False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [False, False, False,  True,  True]],

        [[False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [False, False, False,  True,  True]],

        [[False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True]],

        [[False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True]]])
</code></pre>
<p>and</p>
<pre><code class="lang-auto">attn_mask_
tensor([[[False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [ True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True]],

        [[False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [False, False, False,  True,  True],
         [ True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True]],

        [[False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True],
         [ True,  True,  True,  True,  True]],

        [[False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True],
         [False, False, False, False,  True],
         [ True,  True,  True,  True,  True]]])
</code></pre>
<p>which are different.<br>
The former is what we expect as explained here: <a href="https://gmongaras.medium.com/how-do-self-attention-masks-work-72ed9382510f" rel="noopener nofollow ugc">https://gmongaras.medium.com/how-do-self-attention-masks-work-72ed9382510f</a></p>
<p>I just came up on this issue this morning, while trying to use the SDPA on a custom transformer and digged.<br>
Thanks for posting!<br>
Sincerely,<br>
Gabriel</p>

evaggelos: <p>Hello Gabriel,</p>
<p>Thank you for your response, and I apologize for the delayed reply. The article was great, and I finally understood why it’s called <code>key_padding_mask</code>. The purpose is to prevent the non-<code>&lt;PAD&gt;</code> tokens from communicating with the <code>&lt;PAD&gt;</code> tokens, while still allowing the <code>&lt;PAD&gt;</code> tokens to have a representation. The <code>&lt;PAD&gt;</code> token communicates with the previous tokens only in its own row/representation.</p>