Transformer解码器Attention mask形状错误及掩码应用时机疑问

一、解决RuntimeError：注意力掩码维度不匹配 #

你的报错核心是输入张量维度与nn.MultiheadAttention的期望格式不匹配，导致模型误判序列长度和batch大小，进而对掩码形状产生错误要求。具体解决方式分两种：

方式1：开启batch_first=True（推荐） #

如果你希望保持输入x的形状(batch, seq_len, d_model)（即(1,100,10)），初始化MultiheadAttention时必须显式设置batch_first=True：

self.self_attn = nn.MultiheadAttention(embed_dim=10, num_heads=2, batch_first=True)

此时你的(100,100)形状的lookahead mask会被自动广播到batch维度，无需额外修改掩码形状，直接传入即可。

方式2：调整输入张量维度并适配掩码 #

如果不想开启batch_first=True，需要先将输入x转置为nn.MultiheadAttention默认期望的(seq_len, batch, d_model)格式：

x_transposed = x.transpose(0, 1)  # 形状变为(100,1,10)
hidden_states, _ = self.self_attn(x_transposed, x_transposed, x_transposed, attn_mask=lookahead_mask)
# 之后再将输出转置回(batch, seq_len, d_model)格式
hidden_states = hidden_states.transpose(0, 1)

这种情况下，(100,100)的掩码形状是合法的，可以直接使用。

额外验证：手动扩展掩码维度 #

如果上述方式仍有问题，可以手动给掩码增加batch维度，将(100,100)变为(1,100,100)：

lookahead_mask = lookahead_mask.unsqueeze(0)  # 扩展后形状(1,100,100)
hidden_states, _ = self.self_attn(x, x, x, attn_mask=lookahead_mask)

这种方式适配所有batch_first设置，尤其适合batch内样本掩码一致的场景。

二、词嵌入+位置嵌入后应用注意力掩码的可行性 #

完全可行，这是Transformer标准实现的核心流程之一：

• 词嵌入负责将token映射到语义空间，位置嵌入补充序列的顺序信息，两者相加后得到的张量包含了token的语义和位置信息。
• 在Decoder的自注意力层应用lookahead mask，正是为了让模型在预测当前token时，只能看到当前及之前的token，完全符合Transformer的设计逻辑，没有任何问题。

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