题目重点：以单头 scaled dot-product attention 训练版推导形状、掩码和数值稳定性，并用 NumPy 独立矩阵核验。
题意与假设
这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。
- 训练版是单头注意力：Q 为 n×d、K 为 m×d、V 为 m×dv，布尔 mask 的 True 表示可见；每个 query 至少一个 key 可见。
- 输出是 `softmax(QKᵀ/√d + mask) V`，不含可学习投影、dropout、位置编码或反向传播；题目的具体实现要求未知。
解题路线
- 先检查形状：Q 和 K 最后一维相同，K 和 V 序列长度相同；得分矩阵 n×m。除以 √d 控制点积随维度增大的尺度。
- 逐行应用 mask，在允许位置上减该行最大得分再做 exp 与归一化，避免大正得分导致溢出；所有位置被 mask 应报错或按明确定义处理。
- 每行权重和为 1，再与对应 V 向量加权求和，得到 n×dv。因果自注意力的 mask 常遮住未来位置，交叉注意力的 n 与 m 可以不同。
- 训练代码与 1,001 组独立 NumPy 矩阵计算对照，包含极大分数稳定性；题目未明确 GPU/矩阵库限制。
手动走一遍
一个 query 看两个 key，得分分别 2 和 1，权重为 exp(2)/(exp(2)+exp(1)) 与 exp(1)/(...)，输出是两条 value 的加权平均。若 mask 第二个 key，第一条权重变 1，直接返回第一条 value。
Python 参考实现
下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。
代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。
from math import exp, isfinite, sqrt
def scaled_dot_attention(q, k, v, mask=None):
"""Single-head Q[n,d], K[m,d], V[m,dv]; mask[n,m] True=allowed."""
if not q or not k or not v or len(k) != len(v):
raise ValueError('nonempty Q/K/V and same key/value length required')
d, dv = len(q[0]), len(v[0])
if d == 0 or dv == 0 or any(len(row) != d for row in q + k) or any(len(row) != dv for row in v):
raise ValueError('inconsistent dimensions')
if mask is not None and (len(mask) != len(q) or any(len(row) != len(k) for row in mask)):
raise ValueError('invalid mask shape')
if any(not isfinite(x) for row in q + k + v for x in row):
raise ValueError('nonfinite input')
result = []
scale = sqrt(d)
for i, query in enumerate(q):
allowed = [j for j in range(len(k)) if mask is None or mask[i][j]]
if not allowed:
raise ValueError('each query needs an allowed key')
scores = [sum(a * b for a, b in zip(query, k[j])) / scale for j in allowed]
offset = max(scores)
weights = [exp(score - offset) for score in scores]
denominator = sum(weights)
result.append([sum(weights[t] * v[j][axis] for t, j in enumerate(allowed)) / denominator
for axis in range(dv)])
return result
复杂度
朴素单头计算 QKᵀ 为 O(nmd)，乘 V 为 O(nm·dv)，存完整注意力矩阵 O(nm)；示例逐行实现只保留一行得分/权重，额外 O(m+ndv)，不含输入。长序列的内存瓶颈是后续 flash/paged/window 追问的出发点。
易错点与追问
容易失分的边界
- 忘记除以 √d。
- mask True/False 方向颠倒，泄露未来 token。
- 直接 exp 原始大分数导致上溢。
- Q/K 维度或 K/V 序列长度不匹配仍计算。
- 把这段单头前向代码说成完整 Transformer 多头训练。
继续深挖
- 为什么因果 mask 不改变每行 softmax 的归一化？
- 多头如何划分维度并拼回输出？
- FlashAttention 为什么不必写出完整 n×m 矩阵？
- cross attention 中 n≠m 时哪些形状要改？
