10.3 Scaled Dot-Product 与 Multi-Head Attention：先把 Mask 语义写对

RNN 把历史压进持续更新的状态。Attention 则在当前位置形成 query，与一组 key 比较，再按匹配权重汇总对应的 value。

Attention 就是一个可微检索/加权汇总。真正容易出错的不是 Q/K/V 比喻，而是 shape、缩放、softmax 维度和 mask 语义。

本课目标

推导 scaled dot-product attention 的 shape 与缩放；
区分 self、cross、causal 与 padding attention；
正确构造和测试 Boolean/additive masks；
理解 multi-head projection、拼接与输出映射；
识别 attention 的复杂度和解释限制。

1. Q、K、V 是 Learned Projections

输入表示 $X \in R^{B \times T \times d_{m o d e l}}$ 。Self-attention 常计算：

Q = X W_{Q}, K = X W_{K}, V = X W_{V} .

Query/key 决定匹配分数，value 提供被汇总内容。这些角色由训练目标学出，不保证等同数据库 key/value 或人类语义标签。

Cross-attention 中 Q 来自 target/decoder states，K/V 来自 source/encoder states，因此 query length $L$ 与 source length $S$ 可不同。

2. Scaled Dot-Product Attention

单头：

A = softmax (\frac{Q K^{T}}{\sqrt{d_{k}}} + M),

O = A V .

若：

text

Q: [B, L, d_k]
K: [B, S, d_k]
V: [B, S, d_v]
scores/A: [B, L, S]
O: [B, L, d_v]

Softmax 沿 key/source 维 $S$ 。每个 query row 的 allowed weights 和为 1。

若 Q/K components 近似零均值、单位方差且独立，dot-product variance 随 $d_{k}$ 增长；除以 $\sqrt{d_{k}}$ 把典型尺度拉回，减少 softmax 过早饱和。这是近似动机，不保证实际 trained representations 满足独立假设。

3. Mask 必须在 Softmax 前应用

Additive mask 在禁止位置使用 $- \infty$ ，使 softmax 权重为 0；有限负常数只会相对合法分数压低权重，不能无条件保证为零。Softmax 后再把权重乘 0 会让剩余权重不再归一，除非重新归一化。

常见 masks：

key padding mask：所有 queries 忽略 source padding；
causal mask：位置 $t$ 不看未来 $> t$ ；
local/block/sparse mask：只允许指定连接；
cross-attention source mask：忽略 encoder padding。

Query padding 也要在 loss/output 处 mask；只屏蔽 padded keys 不会自动删除 padded query outputs。

4. Boolean Mask 语义并不统一

这是工程高危点。当前 PyTorch API 中：

nn.MultiheadAttention 的 Boolean key_padding_mask/attn_mask：True 表示不允许/忽略；
torch.nn.functional.scaled_dot_product_attention 的 Boolean attn_mask：True 表示允许参与。

从一个 API 迁移到另一个可能需要取反。不要凭变量名猜，写一个 $3 \times 3$ 人工例子验证 forbidden weights 为 0。

若某个 query 的所有 keys 都被 mask，softmax 没有合法分布，可能得到 NaN 或未定义行为。Batch 构造必须保证至少一个合法 key，或显式定义 empty-row 结果。

5. 从零写一个单头版本

以下 Python 3 / PyTorch 2.14 CPU 例子可直接运行，支持下述三维有限浮点张量，batch 对齐，key 与特征维非空。Boolean mask 须在同一 device，且能广播到 [B,L,S]。

python

import math
import torch

def scaled_dot_product_attention(query, key, value, allowed=None):
    # query: [B, L, D], key: [B, S, D], value: [B, S, Dv]
    if query.ndim != 3 or key.ndim != 3 or value.ndim != 3:
        raise ValueError("query, key and value must be three-dimensional")
    if not query.shape[0] == key.shape[0] == value.shape[0]:
        raise ValueError("batch dimensions must match")
    if query.shape[-1] != key.shape[-1] or key.shape[1] != value.shape[1]:
        raise ValueError("key dimensions and source lengths must match")
    if key.shape[1] == 0 or query.shape[-1] == 0:
        raise ValueError("source length and key dimension must be positive")
    scores = query @ key.transpose(-2, -1) / math.sqrt(query.shape[-1])
    if not torch.isfinite(scores).all() or not torch.isfinite(value).all():
        raise ValueError("attention scores and values must be finite")

    if allowed is not None:
        if allowed.dtype != torch.bool:
            raise TypeError("allowed must be a Boolean tensor")
        # allowed 应可 broadcast 到 [B, L, S]；True 表示允许。
        allowed = torch.broadcast_to(allowed, scores.shape)
        if not torch.all(allowed.any(dim=-1)):
            raise ValueError("Every query must have at least one allowed key")
        scores = scores.masked_fill(~allowed, float("-inf"))

    weights = torch.softmax(scores, dim=-1)
    output = weights @ value
    return output, weights

q = torch.randn(2, 4, 8)
k = torch.randn(2, 6, 8)
v = torch.randn(2, 6, 5)
out, weight = scaled_dot_product_attention(q, k, v)
assert out.shape == (2, 4, 5)
assert torch.allclose(weight.sum(dim=-1), torch.ones(2, 4))

q = k = torch.zeros(1, 3, 2)
v = torch.tensor([[[2.], [4.], [8.]]])
allowed = torch.ones(3, 3, dtype=torch.bool).tril()
out, weight = scaled_dot_product_attention(q, k, v, allowed)
assert torch.all(weight[0][~allowed] == 0)
assert torch.allclose(out[0, :, 0], torch.tensor([2., 3., 14. / 3.]))

Q/K 分数全为零时，每个因果行只平均已允许的值，依次为 2、3、 $14 / 3$ 。Mask 一旦取反，结果立刻不同；比起只看输出形状，这个检查更容易抓住“答案从未来寄来”的错误。

生产代码优先用 framework optimized kernels；这个版本用于建立 shape/mask tests，不处理 fused precision、dropout、GQA 和高性能布局。

6. Causal Mask

长度 $T$ 的 decoder self-attention，allowed matrix：

text

q0: k0
q1: k0 k1
q2: k0 k1 k2
...

即下三角（含 diagonal）。Causal mask 防止训练时 token $t$ 读取未来 targets。仅把 labels shift 一位却忘记 mask，模型会通过 hidden representation 偷看答案，训练 loss 异常低。

Prefix-LM、bidirectional encoder、sequence-to-sequence decoder 的 mask 不同。不要把“Transformer mask”当单一模板。

7. Multi-head Attention

将 $d_{m o d e l}$ 投影成 $H$ 组 head subspaces：

h e a d_{h} = A t t e n t i o n (X W_{Q}^{(h)}, X W_{K}^{(h)}, X W_{V}^{(h)}),

M H A (X) = C o n c a t (h e a d_{1}, \dots, h e a d_{H}) W_{O} .

常见实现要求 $d_{m o d e l}$ 可被 heads 整除， $d_{h} = d_{m o d e l} / H$ 。增加 heads 不一定增加总 projection 参数，但会改变每头维度、kernel efficiency 和表示分解。

“每头分别学习语法、指代、位置”只是某些分析观察，不是训练约束。Heads 可冗余、混合或不稳定，不能给每头强行命名。

8. Framework API 的性能与 Dropout

nn.MultiheadAttention(batch_first=True) 输入可用 [B,T,D]。只需输出、不需要 weights 时设 need_weights=False，框架更可能使用优化的 scaled-dot-product path。

直接调用当前 F.scaled_dot_product_attention 时，dropout_p>0 会按参数应用 dropout，不会自动读取 module 的 train/eval；module 需要显式选择概率。

在 module 方法内，假设 self.dropout 存放配置的概率，对应片段为：

python

dropout_p = self.dropout if self.training else 0.0

Fused Flash/memory-efficient/math kernels 可能因 dtype、device、mask shape 而切换，并有浮点差异。性能测试要记录 backend，不要只从公式估 latency。

9. Attention 没有位置信息

若没有位置表示，也没有固定的依赖位置的 mask，置换 self-attention 输入 token 会让输出作相同置换。Causal 或 local mask 本身已加入位置结构；有此类 mask 时，置换等变需要同时置换 mask 的两个 token 轴。

位置可通过 absolute embedding、sinusoidal、relative bias、rotary 等机制注入。不同方法影响长度外推和 cache 实现，第 10.4 课展开。

10. 复杂度不只有 $O (T^{2})$ 一行

Self-attention score matrix 大小约 [B,H,T,T]，标准 dense attention 的 score 计算/存储随 $T^{2}$ 增长；projection 和 FFN 还依赖 $T d_{m o d e l}^{2}$ 。

在短序列/大 hidden 下 projection/FFN 可能主导；长序列时 attention matrix 常成为瓶颈。FlashAttention 类算法可减少 materialized memory/I/O，并保持 exact attention 语义（允许浮点实现差异），但不会把所有理论工作都自动变线性。

Sparse/local/linear attention 改变连接或 kernel approximation，要验证任务质量和实际加速。

11. Attention Weight 不是因果解释

高 weight 表示当前 head/layer 中 value mixture 的系数，不等于 token 对最终输出的唯一贡献：

value vectors 和 output projection 会改变效果；
residual/FFN/后续 layers 有其他路径；
多组 weights 可产生相近输出；
干预 token 会同时改变 Q/K/V。

可把 weights 用作诊断，但应结合 ablation、gradient/perturbation 与 counterfactual analysis，仍不能直接宣称因果。

常见误区

Q/K/V 是固定语义字段：它们是 learned projections。
Mask 后乘 0 就够：应在 softmax 前屏蔽并测试归一化。
PyTorch 所有 Boolean mask 都是 True=保留：不同 API 语义可相反。
每个 head 会自动分工成人类概念：没有这个约束。
Attention weight 就是解释：最终输出还有 value、projection 和其他路径。

练习

为 cross-attention 写出 Q/K/V/scores/output shapes。
构造 $3 \times 3$ causal mask，验证未来 weights 为 0。
对比 MultiheadAttention 与 SDPA Boolean mask 语义。
令某 query 全部被 mask，观察并修复失败。
估算不同 $T, D, H$ 下 attention score 与 FFN 的内存/运算。

小结

Scaled dot-product attention 用 Q/K 相似度产生每个 query 对 values 的归一化混合。缩放控制 softmax 输入尺度，mask 定义信息流，多头在不同 learned subspaces 并行计算。它擅长直接连接位置，也引入二次矩阵、API 语义和解释风险。

下一课把 attention 放进 residual、normalization 和 FFN，加入位置与 causal objective，组成 encoder、decoder 和完整 Transformer。

第1章逻辑

第2章集合论

第3章证明方法

第4章递归、归纳与递推

第5章关系与函数

第6章组合数学

第7章图论

第8章概率与统计

第9章数论

第10章线性代数

第11章微积分与优化

第12章信息论与数值计算