nn.MultiheadAttention如何实现不同头使用不同掩码

在PyTorch中，nn.MultiheadAttention模块提供了mask参数来进行掩码操作。该参数可以是一个张量，用于指定要掩盖的位置，也可以是一个函数，用于根据序列长度动态生成掩码。

如果要为不同头使用不同的掩码，可以为每个头单独创建一个掩码张量，然后在调用MultiheadAttention模块时将这些张量传递给mask参数。

具体来说，可以将掩码张量组成一个列表，每个列表元素对应一个注意力头，然后将该列表传递给mask参数。这样，每个头就可以独立地应用自己的掩码，从而实现不同头使用不同掩码的效果。

共用一个mask

通常情况下我们是怎么做的

输入：

import torch
seq_len=5
attn_mask=torch.triu(torch.ones(seq_len,seq_len),diagonal=1).to(dtype=torch.bool)
print(attn_mask)

输出：

tensor([[False,  True,  True,  True,  True],
        [False, False,  True,  True,  True],
        [False, False, False,  True,  True],
        [False, False, False, False,  True],
        [False, False, False, False, False]])

输入：

self.multi_en = nn.MultiheadAttention(embed_dim = dim_model, num_heads = heads_en, dropout = dropout)
out, attn_wt = self.multi_en(out, out, out,attn_mask=attn_mask) 

这样所有的头就会共用一个mask

不共用mask

import torch.nn as nn

# 假设有4个注意力头
num_heads = 4

# 创建4个掩码张量，每个张量的大小都是(N, L)，其中N是batch size，L是序列长度
mask1 = torch.tensor([[1, 1, 0], [1, 1, 1], [0, 1, 1], [1, 0, 1]])  # 第1个头的掩码
mask2 = torch.tensor([[1, 0, 1], [1, 1, 0], [0, 1, 1], [1, 1, 1]])  # 第2个头的掩码
mask3 = torch.tensor([[0, 1, 0], [1, 1, 0], [1, 0, 1], [0, 1, 1]])  # 第3个头的掩码
mask4 = torch.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0], [1, 1, 1]])  # 第4个头的掩码

# 将4个掩码张量组成列表
mask_list = [mask1, mask2, mask3, mask4]

# 创建MultiheadAttention模块
multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)

# 调用模块时传递mask参数
output, _ = multihead_attn(query, key, value, mask=mask_list)