← 课程库
语言模型与架构模块三 · 注意力机制代码12 分钟

多头注意力:并行看多个子空间

为什么一个「头」不够用

学完这节你能:
  • 理解多头的动机与做法
  • 算清 d_model、头数 h、每头维度 dₖ 的关系
  • 读懂多头注意力的形状变换

一个注意力头只能学一种「关注模式」。多头注意力把表示切成若干子空间,各自并行做注意力,再拼接——不同头能分别捕捉语法、指代、位置等不同关系。

MultiHead=Concat(head1,,headh)WO\text{MultiHead} = \text{Concat}(\text{head}_1,\dots,\text{head}_h)\,W_O
每个 headi=Attention(QWiQ,  KWiK,  VWiV)\text{head}_i = \text{Attention}(QW_i^Q,\; KW_i^K,\; VW_i^V)

🔑 维度关系:dₖ = d_model / h。总计算量与单头相近——把一个大空间切成 h 个小空间并行算,而不是额外增加成本。

形状怎么变
d_model=512, 头数 h=8 → 每头 dₖ=64
输入 (batch=2, seq=10, 512)
拆头 → (2, 8, 10, 64):8 个头各自算 10×10 的注意力
拼回 (2, 10, 512) 再过 W_O 输出
多头的形状变换(核心就是 reshape 成多头再算再拼回)
# x: (batch, seq, d_model)
q = (x @ Wq).view(batch, seq, h, d_model // h).transpose(1, 2)  # (b, h, seq, dk)
k = (x @ Wk).view(batch, seq, h, d_model // h).transpose(1, 2)
v = (x @ Wv).view(batch, seq, h, d_model // h).transpose(1, 2)
scores = (q @ k.transpose(-2, -1)) / (dk ** 0.5)   # (b, h, seq, seq)
out = softmax(scores, dim=-1) @ v                   # (b, h, seq, dk)
out = out.transpose(1, 2).reshape(batch, seq, d_model) @ Wo
自测d_model=512、8 个头,每个头的维度 dₖ 是:
// 带走一句话

多头=把 d_model 切成 h 个 dₖ=d_model/h 的子空间并行做注意力再拼接,用相近成本捕捉多种关系。

独立练习

做到这些,才算完成本节

  1. 在本地运行代码,打印每个关键张量的 shape。
  2. 关闭示例后从空文件重写核心部分,并补一个异常或边界测试。
学习记录

记录本节阅读进度

建议先独立回答自测。读完记录不等于已经掌握。

继续下一课
← 上一课
自注意力 vs 交叉注意力
下一课 →
位置编码与 RoPE 旋转位置编码

学完不要急着赶进度:去练习自测,或完成对应的阶段实验