在 Transformer 架构中自注意力层最早解决的核心问题是序列里的每个位置如何通过计算与所有其他位置的关联程度来更新自己的表示。可是单头自注意力只能把全部信息压缩到一组 Query、Key、Value 投影里等价于只保留一种关系视角。多头注意力机制正是对自注意力层的重要完善它把 Q、K、V 拆分到多个子空间每个头独立计算注意力再把结果拼接并通过输出投影融合回原始维度。下面会从单头注意力的局限讲起解释多头注意力的数学原理给出一个可直接运行的 PyTorch 实现并把它放到残差连接、层归一化和多层感知机构成的完整 Transformer 块里验证行为最后整理实现中容易踩的坑和工程落地建议。1. 单头自注意力能做什么又缺什么1.1 自注意力的核心Q、K、V 和缩放点积自注意力层用一个输入矩阵 $X$ 同时生成三组向量Query查询、Key键、Value值。对序列中某个位置来说Query 表示“我想找什么信息”Key 表示“我能提供什么信息”Value 表示“我实际携带的内容”。两两位置之间的关联程度通过 Query 与 Key 的点积计算点积越大说明这个位置越值得关注。点积结果要除以 $\sqrt{d_k}$ 进行缩放再经过 softmax 变成权重最后用权重对 Value 做加权求和。单头缩放点积注意力的公式可以写成$$ \text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$在常见的 Transformer 自注意力层中输入 $X$ 的维度是 $(B, L, D)$其中 $B$ 是 batch size$L$ 是序列长度$D$ 是模型维度。三个线性投影矩阵将 $X$ 映射为 $Q$、$K$、$V$每个向量的维度是 $d_k$。$d_k$ 通常取 64$D$ 通常取 512 或 768。这段描述覆盖了自注意力的最小公式结构后面实现多头时仍然会复用这套计算。1.2 单头只能提供一种“关系视角”单头注意力的问题不在计算速度而在表达能力。一个头只有一组 $Q$、$K$、$V$ 投影softmax 之后只能得到一组注意力权重分布。这组权重必须同时承担多种关系判断比如句法依赖、指代消解、局部共现、远距离语义关联。把这些关系全部揉进一个分布里模型很容易顾此失彼。举个例子在“小张因为下雨没有去公园他留在了家里”这句话里“他”需要同时关联“小张”“下雨”“公园”“家里”这些位置。单头注意力虽然也能把部分权重分给这些词但缺少“按类型拆分”的机制。当一个位置的表示需要叠加多种关系时单头通常只能学到一种模糊的平均而不是能够区分不同语义关系的多个子表示。1.3 “多头”要解决的三个实际问题引入多头注意力主要是为了解决三个具体问题。第一表达能力不足。把 $Q$、$K$、$V$ 切割成 $h$ 份每份在低维子空间里独立计算可以让模型同时维持 $h$ 套注意力分布而不是强迫一个分布做所有事情。第二监督信号稀释。反向传播时单头的梯度把所有关系混合在一起难以对某一类关系做精细调整。多头让每个头更容易专注某一类模式训练时梯度也按子空间隔离。第三位置信息的编码能力。不同头可以分别关注“相邻词”和“远距离词”后续叠加位置编码后这种多视角特性会明显提升序列建模能力。这三个问题并不等于说“头数越多越好”但至少解释了为什么单头结构在复杂语言任务上经常不够用。2. 多头注意力机制的拆解、拼接与数学解释2.1 多头注意力的定义多头注意力的标准定义是先将 $Q$、$K$、$V$ 通过线性投影拆到 $h$ 个子空间每个子空间独立执行缩放点积注意力然后把 $h$ 个头的输出拼接起来最后再过一次输出线性投影。公式$$ \begin{aligned} \text{head}_i \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) \ \text{MultiHead}(Q,K,V) \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O \end{aligned} $$其中 $W_i^Q \in \mathbb{R}^{D \times d_k}$$W_i^K \in \mathbb{R}^{D \times d_k}$$W_i^V \in \mathbb{R}^{D \times d_v}$$W^O \in \mathbb{R}^{h d_v \times D}$。通常取 $d_k d_v D/h$。从实现角度看标准做法不是把输入复制 $h$ 份而是先用一个大的线性层把输入从 $D$ 维映射到 $D$ 维然后按 head 数量做 reshape。这样做参数更紧凑GPU 算子也更高效。后面写代码时会严格按这个路径来实现。2.2 为什么点积要除以 sqrt(d_k)点积 $Q \cdot K$ 的结果是一个标量当 $d_k$ 增大时点积的方差也会增大。假设 $Q$ 和 $K$ 中每个元素是均值为 0、方差为 1 的随机变量那么 $Q$ 与 $K$ 的点积期望是 0但方差接近 $d_k$。如果维度是 64点积结果的标准差就接近 8。softmax 的输入如果方差过大输出分布会非常尖锐大部分概率集中在最大值附近其他位置的梯度接近 0。这会直接造成梯度消失训练早期几乎学不动。除以 $\sqrt{d_k}$ 是为了把方差拉回 1 附近让 softmax 工作在更平稳的区域。这是它在数学上的关键作用不是随便凑出来的缩放。2.3 张量形状变化与参数关系下面用一组具体数字说明形状变化。设 batch size $B2$序列长度 $L10$模型维度 $D64$head 数量 $h8$则每个 head 的维度 $d_k D/h 8$。张量形状说明输入 $X$$(2, 10, 64)$默认序列维度在中间拆分前 Q$(2, 10, 64)$线性层输出未改变维度拆分后 Q$(2, 8, 10, 8)$第 2 维是 head第 3 维是序列注意力分数$(2, 8, 10, 10)$每个 head 一个完整注意力矩阵注意力输出$(2, 8, 10, 8)$对 V 加权求和后仍按 head 排列拼接后$(2, 10, 64)$合并 head恢复原始维度输出投影后$(2, 10, 64)$与输入保持相同形状这个形状变化是实现时最容易出问题的地方。很多维度报错都来自transpose和view混用后面排错章节会专门展开。2.4 一个容易混淆的问题多头是否增加了参数量如果定义是“把 $Q$、$K$、$V$ 投影到 $D$ 维再按 head 切分”那么多头注意力并没有比单头多出一整组参数。以 $D512$、$h8$ 为例$Q$ 的大线性层形状是 $(512, 512)$$K$ 和 $V$ 同理。输出投影也是 $(512, 512)$。加起来是 $4 \times 512 \times 512$。单头的写法是 $W_Q, W_K, W_V \in \mathbb{R}^{D \times D}$也是同样的参数量。所以“多头”带来的不是参数数量膨胀而是参数重组方式的变化原来的 512 维投影被切成 8 个 64 维子空间每个子空间独立计算注意力并独立回传梯度。这也是多头机制能够在不过度增加开销的前提下提升表达力的原因。3. 用 PyTorch 从零实现多头自注意力层3.1 环境准备实现用 PyTorch推荐环境如下Python 3.8 或更高版本PyTorch 1.13 或 2.x可选 CUDA但下面示例在 CPU 上也能运行安装依赖时直接使用 pip 即可示例代码只依赖 torch避免