速查：缩放点积注意力

Scaled dot-product attention —— 一页纸，面试前看这个。

来源：Attention Is All You Need §3.2 · nanoGPT model.py · PyTorch SDPA 文档

公式

Attention(Q, K, V) = softmax( QKT + M√dk ) V

M 是因果掩码：j ≤ i 处为 0，j > i 处为 −∞（严格说是先缩放再加掩码，−∞ 缩放后仍是 −∞，所以写在哪都一样）。不加 M 就是 encoder 的双向注意力。

五步 + 形状

#操作单头形状批量多头（nanoGPT）
0输入 X(T, d_model)(B, T, C)
1投影 Q=XW_Q, K=XW_K, V=XW_V(T, d_k) ×3拆头后 (B, nh, T, hs)
2分数 Q @ K.T / √d_k(T, T)(B, nh, T, T)
3因果掩码：上三角填 −∞(T, T)(B, nh, T, T)
4softmax（最后一维，即对 key 维）(T, T) 每行和=1(B, nh, T, T)
5加权求和 weights @ V(T, d_v)(B, nh, T, hs) → 合并头 (B, T, C)

口诀：(T, d) 进，(T, T) 中间，(T, d) 出。中间那个 T×T 矩阵就是计算量 O(T²d) 和显存瓶颈的来源（论文 Table 1），也是 FlashAttention 要优化掉的东西。

为什么除以 √dk（三行推导）

  1. 设 q、k 各分量独立，均值 0、方差 1。则每项 qiki 均值 0，方差 E[qi²]E[ki²] = 1。
  2. q·k = Σi=1..dk qiki 是 dk 个独立项之和 ⇒ 方差 = dk，标准差 √dk（论文脚注）。
  3. 除以 √dk ⇒ 方差回到 1。否则 dk 大时分数差距大，softmax 接近 one-hot，落在梯度极小的区域（"saturate"），训练初期每个 token 只盯住一个位置（Karpathy 视频）。

参考实现

NumPy（单头）

def attention(Q, K, V, causal=False):
    scores = Q @ K.swapaxes(-2, -1) / np.sqrt(Q.shape[-1])
    if causal:
        T = scores.shape[-1]
        scores = np.where(np.triu(np.ones((T, T), bool), k=1), -np.inf, scores)
    w = np.exp(scores - scores.max(-1, keepdims=True))
    w /= w.sum(-1, keepdims=True)
    return w @ V

PyTorch（与 nanoGPT 手写分支一致）

att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(tril[:T, :T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
y = att @ v
# 生产代码：y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

常见坑

坑后果正确做法
softmax 用错轴（dim=0）按列归一化，语义错，shape 却不报错dim=-1：每个 query 对所有 key 归一化
softmax 之后才掩码每行和不再为 1，且未来信息已混入分母掩码一定在 softmax 之前，填 −∞
用 K.T 处理批量张量转置了所有维度k.transpose(-2, -1)，只换最后两维
某一行被全部掩掉（如全是 padding）softmax(全 −∞) = NaN保证每行至少一个可见位置（因果掩码对角线总可见）
忘了缩放小 dk 看不出，大 dk 训练变慢/不稳除以 sqrt(d_k)，注意是 dk 不是 dmodel

面试一句话