PyTorch MultiheadAttention为何需为每个注意力头单独指定attn_mask？

问题1：为何需要为不同注意力头设置不同的掩码？ #

• 不同注意力头的核心作用就是捕捉序列里不一样的依赖关系：有的头盯着相邻的局部token，有的头能跨长距离找语义关联，甚至针对特定任务有专属的聚焦模式。
• 给不同头配不同掩码，能精准约束每个头的注意力范围：比如多任务场景里，让几个头专门处理任务A的token关联，剩下的头专注任务B的规则；或者做定制化建模时，给不同分工的头加不同限制，让模型更适配任务需求。
• 这也给研究和调参留了灵活度，比如做消融实验时，能单独控制某类头的注意力范围，精准验证它的作用。

问题2：3D掩码的排列顺序是怎样的？ #

3D掩码(N⋅num_heads, L, S)的排列顺序是**（样本1，注意力头1）、（样本1，注意力头2）……（样本1，注意力头num_heads）、（样本2，注意力头1）、（样本2，注意力头2）……**

• 这是因为PyTorch处理MultiheadAttention时，会先把输入张量从(N, L, E)（E是输入特征维度）转置重塑成(N⋅num_heads, L, E//num_heads)，这个过程是先把同一个样本的所有注意力头数据拼在一起，再处理下一个样本。掩码必须和这个张量维度对齐，才能准确对应到每个样本的每个头的注意力计算上。

内容的提问来源于stack exchange，提问作者Bastiaan