Self-attention gives incorrect weight

nlp
1

Hi, I was trying to use the pytorch implementation of multihead attention with masks (“key_padding_mask”) for the padded tokens. However, it seems the mask is doing weird things. Here is an example:
I have three tokens, each with embedding dimension of 4. Assuming the last token is padding, I added the mask [False, False, True]. I would expect the attention weight for the masked token to be zero

import torch
from torch import optim, nn
self_attn = nn.MultiheadAttention(4,2, dropout=0, batch_first=True)
x=torch.rand((3,4)) #sequence length x embed_d
res, attn=self_attn(x,x,x,key_padding_mask=torch.Tensor([False,False,True]),need_weights=True)

Then attention matrix gives a result like following:

attn= tensor([[0.2125, 0.2203, 0.5673],
        [0.2041, 0.2077, 0.5881],
        [0.2086, 0.2112, 0.5803]], grad_fn=<SqueezeBackward1>)

The attention weight for the masked token (last column) is not zero, but a big value.
What is going on?