语言模型与架构模块三 · 注意力机制代码约 12 分钟
多头注意力:并行看多个子空间
为什么一个「头」不够用
学完这节你能:
- 理解多头的动机与做法
- 算清 d_model、头数 h、每头维度 dₖ 的关系
- 读懂多头注意力的形状变换
一个注意力头只能学一种「关注模式」。多头注意力把表示切成若干子空间,各自并行做注意力,再拼接——不同头能分别捕捉语法、指代、位置等不同关系。
每个
🔑 维度关系: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 的子空间并行做注意力再拼接,用相近成本捕捉多种关系。
独立练习
做到这些,才算完成本节
- 在本地运行代码,打印每个关键张量的 shape。
- 关闭示例后从空文件重写核心部分,并补一个异常或边界测试。
学习记录
记录本节阅读进度
建议先独立回答自测。读完记录不等于已经掌握。