为何多头注意力中Query不影响输出且输出沿首维度重复？

问题背景 #

我正在进行多头注意力实验，遇到两个困惑：

• 为什么代码中Query嵌入的值对注意力输出没有影响？
• 为什么输出会沿序列维度重复？

实验代码 #

import torch
import torch.nn as nn

multihead_attn = nn.MultiheadAttention(embed_dim=2, num_heads=1, batch_first=True)

kv = torch.randn(1, 1, 2)
print("Key / Val")
print(kv)
for _ in range(2):
    q = torch.randn(1, 2, 2)
    print("Query")
    print(q)
    attn_output, _ = multihead_attn(q, kv, kv)
    print("Output")
    print(attn_output)

实验输出 #

Key / Val
tensor([[[ 0.1782, -1.3460]]])
Query
tensor([[[-0.7521,  0.6856],
         [-0.8761, -1.6864]]])
Output
tensor([[[ 0.0517, -0.3687],
         [ 0.0517, -0.3687]]], grad_fn=<TransposeBackward0>)
Query
tensor([[[-0.9609, -1.0166],
         [-1.1555, -1.3593]]])
Output
tensor([[[ 0.0517, -0.3687],
         [ 0.0517, -0.3687]]], grad_fn=<TransposeBackward0>)

疑问解答 #

1. Query值不影响输出的原因

核心原因是你的Key序列长度只有1（kv的形状是(1,1,2)，对应batch_size=1, 序列长度=1, 嵌入维度=2）。多头注意力的计算逻辑是：先算每个Query和所有Key的相似度，再通过softmax得到注意力权重，最后用权重对Value加权求和。当Key只有1个元素时，不管Query是什么，softmax后的权重必然是1（只有一个候选，归一化后就是1），所以最终输出就是这个唯一的Value经过模型内部线性投影后的结果——完全和Query无关。

另外要注意，nn.MultiheadAttention会自动对Q、K、V做线性变换（哪怕是单头），所以你看到的输出不是原始的kv值，而是变换后的结果，但这个变换只和模型参数有关，和Query没有关系。

2. 输出沿序列维度重复的原因

你的Query序列长度是2（q的形状是(1,2,2)，序列长度=2），每个Query都会和唯一的Key计算注意力，得到的权重都是1，所以每个Query对应的输出都是同一个变换后的Value。因此输出的2个序列位置的值完全相同，看起来就像是沿序列维度重复了。

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