题目重点：按题目“手搓 MHA”的技术方向构建可验证的预投影 Q/K/V 训练版，解释 shape、softmax 稳定性与因果 mask。
题意与假设
这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。
- Q/K/V 已完成线性投影，都是 [T,D]，T>0，D 可被 head 数 H 整除；不实现投影矩阵、batch 维或 output projection。
- 每头维 d=D/H，按最后维连续切分；可选 causal mask 禁止位置 i 看到 j>i。题目未限定这些规则，均为训练约定。
解题路线
- 每头取 Q、K、V 的对应维段，计算每对位置的点积除以 √d；按目标 mask 排除非法 key。
- 对每个 query 先减最大分数再取 exp，做数值稳定的 softmax；用权重对该头的 V 加权求和。
- 把 H 个头的输出按原始维段顺序拼接，得到 [T,D]。完整 MHA 通常还有 QKV 与输出的可训练线性投影，不能把训练版冒充完整生产层。
- 形状、非法 H、空序列和因果首 token 都要测；训练代码与独立 NumPy 矩阵计算做随机对照。
手动走一遍
T=2、D=4、H=2 时每头 d=2，每头分数矩阵 2×2，输出也是 2×2；拼接后为 2×4。causal 模式下位置 0 只能看自身，故输出第 0 行每头等于 V 第 0 行对应维段。
Python 参考实现
下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。
代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。
"""Reference multi-head attention for preprojected Q, K and V."""
from math import exp, sqrt
def multihead_attention(
query: list[list[float]],
key: list[list[float]],
value: list[list[float]],
heads: int,
causal: bool = False,
) -> list[list[float]]:
length = len(query)
if length == 0 or len(key) != length or len(value) != length:
raise ValueError("Q, K and V must have the same nonzero sequence length")
width = len(query[0])
if heads <= 0 or width == 0 or width % heads:
raise ValueError("width must be positive and divisible by heads")
if any(len(row) != width for matrix in (query, key, value) for row in matrix):
raise ValueError("Q, K and V must have the same width")
head_width = width // heads
output = [[0.0] * width for _ in range(length)]
for head in range(heads):
start = head * head_width
stop = start + head_width
for position in range(length):
allowed = range(position + 1) if causal else range(length)
scores = [
sum(query[position][dim] * key[other][dim] for dim in range(start, stop))
/ sqrt(head_width)
for other in allowed
]
maximum = max(scores)
weights = [exp(score - maximum) for score in scores]
total = sum(weights)
for offset in range(start, stop):
output[position][offset] = sum(
weights[other] / total * value[other][offset]
for other in allowed
)
return output
复杂度
标准全注意力计算每头 T² 个分数、每分数 d 次乘加，合计 O(T²D)；投影另约 O(TD²)。展示代码按 query 临时保留 O(T) 分数，输出 O(TD)，若批量显式存注意力矩阵通常 O(HT²) 空间。
易错点与追问
容易失分的边界
- 忘除以 √d，维度增长时 softmax 易饱和。
- 用 exp(rawScore) 导致数值溢出。
- 因果 mask 方向反了，让早期 token 读到未来。
- 把预投影参考代码声称为包含训练参数的完整 Transformer 层。
继续深挖
- 加入 batch、padding mask 与不同 Q/K 序列长度的 cross-attention 时 shape 如何变化？
- KV cache 如何减少自回归推理重复计算，存储随上下文怎样增长？
