I'm currently trying to implement an Encoder-Decoder architecture for text summarization based on Transformers. Thus I need ti apply MultiHeadAttention on the Decoder site of the model. Since I want to ensure that the model doesn't attend to unseen tokens of the target sequence, I need to use the 3D attention mask (attn_mask) argument.
According to the documentation (https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html), the shape of the mask must be BATCH_SIZE * NUMBER_HEADS, SEQUENCE_LENGTH, SEQUENCE_LENGTH. Which is fine, as it provides the possibility to use different attentions between the Heads, which I don't need in my case ...
But the documentation doesn't state how the tensor needs to be filled regarding it's first dimension and I can't see/ find it in the implementation how it is actually used...
Is it:
[
[2D Attention for Batch 1 for Head 1]
[2D Attention for Batch 2 for Head 1]
...
[2D Attention for Batch 1 for Head 2]
[2D Attention for Batch 2 for Head 2]
...
[2D Attention for Batch n for Head n]
]
or
[
[2D Attention for Batch 1 for Head 1]
[2D Attention for Batch 1 for Head 2]
...
[2D Attention for Batch 2 for Head 1]
[2D Attention for Batch 2 for Head 2]
...
[2D Attention for Batch n for Head n]
]
Would be great if someone knows :)