第 1 周 · Day 1：自注意力与多头注意力

时长 8 小时：§1 自注意力（1.5h）· §2 缩放点积注意力拆开看，含反向推导（2h）· §3 掩码（1h）· §4 多头（1.5h）· §5 复杂度（0.5h）· §A API（1h）· §B–§C（0.5h）

对应正课 Lecture 1 后半段：注意力机制 → Transformer 架构。建议先读完本文 §1–§4 再看视频，视频会快很多。

配套代码 w1day1_checks.py（纯 numpy）：投影打破对称、置换等变、凸包性质、秩坍缩、注意力反向的手推 vs 数值微分、因果掩码的梯度隔离、左补齐的 NaN、单头 vs 多头的表达力、FLOPs 与显存、初始注意力近似均匀

§B 参考实现（code/ 目录）：w1d1_attention.py。先自己写，再对照。

今天的目标：能不看资料写出带掩码的多头注意力的前向，并推出它的反向；能说清每个设计（三个投影、 $\sqrt{d_k}$ 、掩码、多头）是为了解决什么问题。

§0 今天的主线

0.1 从 Week 0 接上

Week 0 Day 4 §5.6 留了一个三步跳跃的预告：

Bahdanau 注意力（解码器看编码器，RNN 的附件）
   ──跳跃一：序列看自己──────────► 自注意力              ← 今天 §1–§3
   ──跳跃二：扔掉 RNN────────────► 纯注意力 + 位置编码     ← Day 2
   ──跳跃三：一次查多种关系──────► 多头注意力            ← 今天 §4

今天完成跳跃一和跳跃三，Day 2 补上跳跃二以及让整个块能工作的其余零件（FFN、残差、LayerNorm、位置编码），Day 3 把它们组装成一个 mini-GPT，替换掉 Week 0 Day 5 的 LSTM。

0.2 今天要用到的旧结论

旧结论	出处	今天用在哪
注意力 = 可微的软查表， $\tau\to0$ 退化为硬查表	Week 0 Day 4 §5.3	§1.2
Q/K/V 的来历：key 和 value 为什么要拆开	Week 0 Day 4 §5.4	§1.4
$\mathrm{Var}[q^\top k]=d$ ，softmax 饱和时雅可比塌缩	Week 0 Day 4 §5.5	§2.1
softmax 的雅可比 $\mathrm{diag}(p)-pp^\top$	Week 0 Day 1 §5.2.1	§2.3
矩阵乘的反向： $\bar A=\bar CB^\top$ ， $\bar B=A^\top\bar C$	Week 0 Day 1 §5.8	§2.3
五步推导法	Week 0 Day 1 §5.0	§2.3
掩码规则：每行至少保留一个未被掩的位置	Week 0 Day 4 §A.4	§3.4
teacher forcing：一次前向算 $T$ 个位置的损失	Week 0 Day 5 §2.3	§3.1

记号：本周默认行向量约定， $X\in\mathbb R^{T\times d}$ 的第 $i$ 行是第 $i$ 个 token 的表示 $x_i^\top$ 。这和 PyTorch 的 (batch, seq, d) 布局一致，写出来的公式可以直接翻译成代码。

§1 自注意力：让序列看自己

1.1 动机：两堵墙同时撞上

Week 0 结束时我们手里有两个没解决的问题：

静态词向量的天花板（Week 0 Day 3 §6）：“bank”只有一个向量，只能落在”银行”和”河岸”之间。我们需要上下文相关的表示：同一个词在不同句子里得到不同的向量。
RNN 的串行（Week 0 Day 4 §5.6）：第 $t$ 步必须等第 $t-1$ 步算完， $T$ 个位置要走 $T$ 步，GPU 的并行能力用不上；相隔 $k$ 个位置的信息要经过 $k$ 次状态更新才能传到，路径越长越容易丢。

RNN 其实解决了问题 1（ $h_t$ 就是上下文相关的表示），代价正是问题 2。有没有一种机制，既能让每个词的表示依赖整句话，又不需要逐步传递？

1.2 直觉：每个词同时向全句提问

回忆 Bahdanau 注意力：解码器在第 $t$ 步拿着自己的状态当查询，去编码器的所有位置软查表，取回一个加权平均（Week 0 Day 4 §5.3）。

现在把这个操作搬到同一个序列内部：

句子里的每个词都拿自己当查询，去同一句话的所有词里软查表，用查到的内容更新自己。

以”I sat on the bank of the river”为例。“bank”发出查询”我附近有什么能说明我是哪个意思的词？”，“river”的 key 与之高度匹配，于是”bank”取回了大量”river”的内容，更新后的表示就偏向”河岸”。换成”I deposited money at the bank”，同一个查询匹配上的是”money”和”deposited”。同一个词，不同的上下文，不同的输出。 问题 1 解决了。

再看问题 2：每个词的查表互不依赖。“bank”查它的，“river”查它的，全部可以同时进行，写成矩阵就是一次矩阵乘。任意两个位置之间都是一步直达。问题 2 也解决了。

这就是自注意力。它不是新发明，是把 2014 年的注意力换了一个作用对象。 真正的洞察在于意识到：注意力本身就足以承担 RNN 的全部职能，循环是多余的。

1.3 形式化

输入 $X\in\mathbb R^{T\times d}$ 。三个可学的投影矩阵 $W_Q,W_K\in\mathbb R^{d\times d_k}$ ， $W_V\in\mathbb R^{d\times d_v}$ ：

$Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V$

$\boxed{\ \mathrm{Attention}(Q,K,V)=\underbrace{\mathrm{softmax}\!\Big(\frac{QK^\top}{\sqrt{d_k}}\Big)}_{A\in\mathbb R^{T\times T}}V\ }$

softmax 对每一行做（对 key 那一维归一化）。逐元素写开，第 $i$ 个位置的输出是

$y_i=\sum_{j=1}^{T}A_{ij}\,v_j,\qquad A_{ij}=\frac{\exp(q_i^\top k_j/\sqrt{d_k})}{\sum_{l}\exp(q_i^\top k_l/\sqrt{d_k})}$

和 Week 0 Day 4 §5.2 的三步（打分 → 归一化 → 加权求和）逐项对应，唯一的区别是 $q_i$ 、 $k_j$ 、 $v_j$ 都来自同一个 $X$ 。

形状速查（单头，一个序列）：

量	形状	含义
$X$	$T\times d$	输入，每行一个 token
$Q,K$	$T\times d_k$	每个 token 的”问题”和”标签”
$V$	$T\times d_v$	每个 token 的”内容”
$S=QK^\top/\sqrt{d_k}$	$T\times T$	$S_{ij}$ ：第 $i$ 个词对第 $j$ 个词的打分
$A=\mathrm{softmax}(S)$	$T\times T$	每行和为 1
$Y=AV$	$T\times d_v$	每行是 value 的加权平均

1.4 为什么要三个投影，而不是直接用 $X$

动机。 最省事的写法是 $\mathrm{softmax}(XX^\top/\sqrt d)X$ ，一个参数都没有。为什么不行？

原因一：只看自己。 Week 0 Day 4 §5.6 已经验证过：不投影时 $x_i^\top x_i=\|x_i\|^2$ ，在范数相近时往往是一行里最大的，于是每个词都主要注意自己。w1day1_checks.py 第 1 组：20 个随机词向量，不投影时 100% 的词把最大权重给了自己，加上两个独立投影后降到 0%。

原因二：关系必须能有方向。 投影后第 $i$ 个词对第 $j$ 个词的打分是一个双线性型：

$S_{ij}=\frac{1}{\sqrt{d_k}}x_i^\top\underbrace{W_QW_K^\top}_{M}\,x_j$

如果 $W_Q=W_K$ （或者干脆不投影）， $M$ 是对称的，于是 $S_{ij}=S_{ji}$ ：“i 关注 j”和”j 关注 i”被迫一样强。但语言里的关系大多有方向：动词找主语、代词找先行词、形容词找它修饰的名词，反过来都不成立。两个独立的投影让 $M$ 可以不对称。第 1 组实测：不投影时 $\|S-S^\top\|/\|S\|=0$ ，投影后是 1.36，完全不对称。

原因三：三种角色需要三种表示。 这是 Week 0 Day 4 §5.4 讲过的：key 编码”什么让我和你相关”，value 编码”被选中后我该给你什么”，两者是不同的目标。query 又是第三种：“我在找什么”。用三个投影，让三种角色各学各的。

意义。 注意 $M=W_QW_K^\top$ 的秩至多是 $d_k$ 。所以 $d_k<d$ 时，每个头的打分只能利用输入的一个 $d_k$ 维子空间。这一点在 §4.5 会变得重要。

1.5 置换等变：自注意力不知道顺序

推导。 设 $P$ 是一个置换矩阵（打乱行的顺序）。把输入换成 $PX$ ：

$Q'=PXW_Q=PQ,\quad K'=PK,\quad V'=PV$

$S'=\frac{PQ(PK)^\top}{\sqrt{d_k}}=P\,S\,P^\top$

softmax 逐行做，对行和列同时置换不改变每行的内容，只改变位置，所以 $A'=PAP^\top$ 。于是

$Y'=A'V'=PAP^\top PV=PAV=PY\qquad(\text{因为 }P^\top P=I)$

$\boxed{\ f(PX)=P\,f(X)\ }$

打乱输入，输出只是被同样地打乱。 “dog bites man”和”man bites dog”在自注意力眼里是同一袋词，“dog”在两句里得到的输出完全相同。w1day1_checks.py 第 2 组：误差 $6.7\times10^{-16}$ ，机器精度。

意义。 这是自注意力最大的结构性缺陷，也是 Transformer 必须外挂位置编码的原因（Day 2 §3）。RNN 不需要位置编码，因为循环本身就是按顺序的；扔掉循环，顺序也一起扔掉了。

因果掩码（§3.1）会打破等变性（第 2 组：误差变成 2.89），因为第 $t$ 个位置只能看前 $t$ 个词，这本身就泄露了”我在第几个”。有研究利用这一点做不带位置编码的 decoder（NoPE），但主流模型仍然显式加位置信息。

1.6 收益与代价

收益	代价
每个词的表示依赖整句（上下文相关）	计算和显存随 $T^2$ 增长（§5）
所有位置并行计算	不知道顺序，必须外挂位置编码
任意两个位置一步直达（路径长度 $O(1)$ ）	本身只做线性混合，需要 FFN 补非线性（§2.2）
注意力矩阵可直接可视化	推理仍然是逐 token 的（§5.4）

§2 缩放点积注意力，拆开看

2.1 为什么是 $\sqrt{d_k}$ ，以及刚初始化时注意力在做什么

$\sqrt{d_k}$ 的完整推导在 Week 0 Day 4 §5.5：分量独立、单位方差时 $\mathrm{Var}[q^\top k]=d_k$ ；方差大导致 softmax 饱和，雅可比 $\mathrm{diag}(p)-pp^\top\to0$ ，梯度消失；除以 $\sqrt{d_k}$ 把方差拉回 1。这里补两个细节。

细节一：是 $d_k$ ，不是 $d_{\text{model}}$ 。 多头时每个头的点积只在 $d_k=d/h$ 维上求和，方差是 $d_k$ ，所以除的是 $\sqrt{d_k}$ 。PyTorch 的 F.scaled_dot_product_attention 默认 scale = 1/sqrt(query.size(-1))，用的正是每头的维度。

细节二：初始化时注意力几乎是均匀的。 “单位方差”只是推导的假设。实际中 $q,k$ 的尺度由投影矩阵的初始化决定。PyTorch 的 nn.Linear 默认从 $U(-1/\sqrt d,1/\sqrt d)$ 采样，方差 $1/(3d)$ （Week 0 Day 2 §5.2）。输入每维方差 1 时：

$\mathrm{Var}[q_i]=d\cdot1\cdot\frac{1}{3d}=\frac13,\qquad\mathrm{Var}\Big[\frac{q^\top k}{\sqrt{d_k}}\Big]=\frac{d_k\cdot\frac13\cdot\frac13}{d_k}=\frac19$

分数的标准差只有 $1/3$ 。w1day1_checks.py 第 11 组（ $d=512$ ，8 头， $T=128$ ）：实测标准差 0.333；每行最大权重平均 0.018，均匀分布是 0.008；熵 4.797，上限 $\ln128=4.852$ 。刚初始化的注意力几乎在做全局平均。 这不是坏事：分数小，softmax 远离饱和，梯度畅通；学会”看哪里”是训练的事。

旁注：有一派参数化方法（μP，Yang et al. 2022）主张按 $1/d_k$ 而不是 $1/\sqrt{d_k}$ 缩放，理由是训练后 $q$ 与 $k$ 会变得相关，点积按 $d_k$ 而不是 $\sqrt{d_k}$ 增长。这属于”如何在不同宽度间迁移超参”的问题，第 1 周不展开，知道有这回事即可。

2.2 输出是凸组合：注意力只混合，不变换

推导。 $A$ 的每一行非负且和为 1，所以 $y_i=\sum_jA_{ij}v_j$ 是 value 的凸组合，落在 $\{v_j\}$ 的凸包里。逐维看：

$\min_jv_{jk}\ \le\ y_{ik}\ \le\ \max_jv_{jk}$

w1day1_checks.py 第 3 组验证了这一点。

意义一：注意力需要一个搭档。 凸组合是一种平均。它能把信息从一个位置搬到另一个位置，但不能对单个位置的表示做逐元素的非线性变换（比如”如果第 3 维大于 0，就把第 7 维翻倍”）。这就是为什么每个 Transformer 块里注意力后面都跟一个逐位置的 FFN（Day 2 §1）：注意力负责在位置之间搬运信息，FFN 负责在每个位置上加工信息。

意义二：秩坍缩。 平均的平均还是平均。反复做下去，所有 token 会趋于同一个向量。Dong et al.（2021）证明了纯注意力堆叠会让表示以双指数速度坍缩成秩 1。w1day1_checks.py 第 4 组用”离所有行都相同有多远” $\|X-\mathbf 1\bar x^\top\|/\|X\|$ 来度量（20 个种子的中位数）：

层数	0	1	2	3	4	8	12
纯注意力	0.98	0.76	$2.1\times10^{-2}$	$2.2\times10^{-7}$	$4.1\times10^{-16}$	机器精度	机器精度
加残差	0.98	0.97	0.95	0.90	0.83	0.47	0.18
残差 + LN	0.98	0.97	0.95	0.91	0.84	0.38	0.10

纯注意力的指数每层至少翻倍（ $10^{-2}\to10^{-7}\to10^{-16}$ ），4 层就坍缩到机器精度。加上残差，双指数坍缩变成了缓慢的衰减。 残差连接对注意力来说不只是”给梯度开高速公路”（Week 0 Day 1 §7.2），它还在保住表示的秩： $x+\mathrm{Attn}(x)$ 里原封不动的那个 $x$ 不参与平均。随机权重下 12 层后仍有衰减，真实模型里 FFN 和训练会进一步维持多样性。

2.3 反向传播：索引级推导

动机。 为什么要手推注意力的反向？三个理由：(1) 面试会考；(2) 理解”反向为什么要存 $T\times T$ 的矩阵”，才理解 FlashAttention 在优化什么（Day 6 §3）；(3) 自定义注意力变体时要自己写 backward。

设定（单头，一个序列）：

$S=s\,QK^\top\ (s=1/\sqrt{d_k}),\qquad A=\mathrm{softmax}_{\text{row}}(S),\qquad Y=AV$

已知上游梯度 $\bar Y\equiv\partial L/\partial Y\in\mathbb R^{T\times d_v}$ ，求 $\bar Q,\bar K,\bar V$ 。按 Week 0 Day 1 §5.0 的五步法，从后往前一个节点一个节点地推。

第一步： $Y=AV$ ，求 $\bar V$ 和 $\bar A$ 。 这是一个矩阵乘。标量形式 $Y_{ik}=\sum_jA_{ij}V_{jk}$ ，于是

$\bar V_{jk}=\sum_i\frac{\partial L}{\partial Y_{ik}}\frac{\partial Y_{ik}}{\partial V_{jk}}=\sum_i\bar Y_{ik}A_{ij}\ \Rightarrow\ \boxed{\bar V=A^\top\bar Y}$

$\bar A_{ij}=\sum_k\bar Y_{ik}V_{jk}\ \Rightarrow\ \boxed{\bar A=\bar YV^\top}$

和 Week 0 Day 1 §5.8 练习 3 的矩阵乘反向完全一样：谁乘了谁，梯度就乘谁的转置。 形状核对： $\bar V$ 是 $(T\times T)^\top(T\times d_v)=T\times d_v$ ✓， $\bar A$ 是 $(T\times d_v)(d_v\times T)=T\times T$ ✓。

第二步： $A=\mathrm{softmax}(S)$ ，求 $\bar S$ 。 softmax 逐行独立，每一行用 Week 0 Day 1 §5.2.1 的雅可比 $\partial A_{ij}/\partial S_{il}=A_{ij}(\delta_{jl}-A_{il})$ ：

$\bar S_{il}=\sum_j\bar A_{ij}\frac{\partial A_{ij}}{\partial S_{il}}=\sum_j\bar A_{ij}A_{ij}(\delta_{jl}-A_{il})=A_{il}\bar A_{il}-A_{il}\sum_j\bar A_{ij}A_{ij}$

$\delta_{jl}$ 把第一项的求和塌缩到 $j=l$ 。定义每行一个标量 $D_i=\sum_j\bar A_{ij}A_{ij}$ ：

$\boxed{\ \bar S_{il}=A_{il}\big(\bar A_{il}-D_i\big)\ }$

直觉： $D_i$ 是 $\bar A$ 在第 $i$ 行按 $A$ 加权的平均。 $\bar S$ 就是”每个元素的梯度减去它们的加权平均，再乘以权重”。减均值来自 softmax 的平移不变性（Week 0 Day 1 §2.4）：所有分数同加一个常数不改变输出，所以梯度在”全体同加”的方向上必须为零，也就是 $\sum_l\bar S_{il}=0$ 。

一个重要的恒等式。 把 $\bar A=\bar YV^\top$ 代进 $D_i$ ：

$D_i=\sum_j\Big(\sum_k\bar Y_{ik}V_{jk}\Big)A_{ij}=\sum_k\bar Y_{ik}\underbrace{\sum_jA_{ij}V_{jk}}_{Y_{ik}}=\boxed{\ \bar y_i^\top y_i\ }$

$D_i$ 只需要 $\bar Y$ 和 $Y$ 的第 $i$ 行，不需要 $T\times T$ 的 $\bar A$ 。 这个看起来平淡的代数变形，是 FlashAttention 反向能省掉 $O(T^2)$ 显存的关键之一（Day 6 §3）。

第三步： $S=sQK^\top$ ，求 $\bar Q$ 和 $\bar K$ 。 又是矩阵乘， $S_{il}=s\sum_mQ_{im}K_{lm}$ ：

$\boxed{\ \bar Q=s\,\bar SK,\qquad\bar K=s\,\bar S^\top Q\ }$

掩码怎么处理？ 被掩的位置 $A_{il}=0$ ，由 $\bar S_{il}=A_{il}(\cdots)$ 立刻得到 $\bar S_{il}=0$ ：梯度不会流过被掩的位置，不需要额外处理。

验证。 w1day1_checks.py 第 5 组（带因果掩码）： $\bar Q,\bar K,\bar V$ 与数值微分的最大相对误差分别为 $1.1\times10^{-9}$ 、 $1.7\times10^{-9}$ 、 $2.4\times10^{-10}$ ； $D_i=\bar y_i^\top y_i$ 的误差 $8.9\times10^{-16}$ ；被掩位置的 $\bar S$ 恰为 0。

代价：反向要什么。 看三个框：算 $\bar V$ 要 $A$ ，算 $\bar S$ 要 $A$ ，算 $\bar Q,\bar K$ 要 $Q,K$ 。 $A$ 是 $T\times T$ 的，必须从前向保存到反向（或者重算）。这就是 §5.2 那笔显存账的来源。

完整的反向，一共五行：

dV = A.T @ dY                                   # (T,T)ᵀ(T,dv)
dA = dY @ V.T                                   # (T,dv)(dv,T)
D  = (dY * Y).sum(-1, keepdims=True)            # 每行一个标量，= (dA * A).sum(-1)
dS = A * (dA - D)                               # softmax 反向
dQ, dK = s * dS @ K, s * dS.T @ Q

2.4 三种角色，一张图

            X (T×d)
     ┌───────┼───────┐
   ·W_Q    ·W_K    ·W_V
     │       │       │
   Q(T×dk) K(T×dk) V(T×dv)
     └──┬────┘       │
   S = QKᵀ/√dk (T×T)  │        ← 谁和谁相关（打分）
        │            │
   A = softmax_row(S) │        ← 每行一个概率分布（查表权重）
        └─────┬──────┘
          Y = AV (T×dv)        ← 按权重取回内容

两个乘法的分工： $QK^\top$ 决定”看谁”（路由）， $AV$ 决定”拿什么”（搬运）。Elhage et al.（2021）把这两部分分别叫 QK 电路和 OV 电路，是机械可解释性研究的基本分解（§4.5）。

§3 掩码

3.1 因果掩码：为什么训练时必须”蒙住未来”

动机。 语言模型的训练目标是”根据前文预测下一个 token”。Week 0 Day 5 §2.3 讲过，为了效率，我们一次前向就算出全部 $T$ 个位置的预测：位置 $t$ 的输出去预测 $x_{t+1}$ 。

但自注意力让每个位置都能看到整句话，包括 $x_{t+1}$ 本身。位置 $t$ 只要学会”把下一个位置的内容抄过来”，损失就能降到几乎为零，而这个模型在生成时毫无用处：生成时下一个 token 还不存在。

直觉。 这就是考试时把答案放在考生眼前。RNN 天然不存在这个问题，因为 $h_t$ 只由 $x_{\le t}$ 算出；自注意力需要人为地把未来遮住。

做法。 在 softmax 之前，把 $j>i$ 的分数设成 $-\infty$ ：

$S_{ij}\leftarrow\begin{cases}S_{ij}&j\le i\\-\infty&j>i\end{cases}\qquad\Rightarrow\qquad A_{ij}=0\ \ (j>i)$

$\exp(-\infty)=0$ ，被掩的位置在分子和分母里都消失，剩下的权重自动重新归一化。掩码矩阵是一个严格上三角：

      key →  x1  x2  x3  x4
query x1  [  ·   ✗   ✗   ✗ ]
      x2  [  ·   ·   ✗   ✗ ]      ✗ = −∞（屏蔽）
      x3  [  ·   ·   ·   ✗ ]      · = 保留
      x4  [  ·   ·   ·   · ]

验证：掩码真的隔断了信息流吗？ 看梯度最直接：如果位置 $t$ 的输出不依赖 $x_s$ （ $s>t$ ），那么 $\partial y_t/\partial x_s=0$ 。w1day1_checks.py 第 6 组用数值微分算出完整的敏感度矩阵：不加掩码时，未来位置的最大敏感度是 1.1；加因果掩码后恰好是 0。

为什么必须加在 softmax 之前？ 如果在 softmax 之后把 $A_{ij}$ 置零，每行的和就不再是 1，而且分母里仍然包含了未来位置的 $\exp(S_{ij})$ ，信息照样通过分母泄露（改变 $x_{t+1}$ 会改变分母，从而改变 $y_t$ ）。

意义。 因果掩码 + teacher forcing = 一次前向、 $T$ 个监督信号、零泄露。这是 decoder-only 模型训练效率的根本来源，也是它比 RNN 快的另一半原因（另一半是 §5.3 的并行）。

3.2 Padding 掩码

动机。 一个 batch 里的序列长短不一，短的要补齐（padding）。padding 位置没有任何含义，不应该被任何 query 注意到，否则真实 token 的表示里会混进 pad 的内容，结果随补齐的多少而变。

做法。 屏蔽的是 key 这一维：如果第 $j$ 个位置是 pad，那么所有 query 的第 $j$ 列都设为 $-\infty$ 。

 key_padding = [F, F, F, T, T]           （后两个是 pad）
 掩码（广播到每一行）：
      [ ·  ·  ·  ✗  ✗ ]
      [ ·  ·  ·  ✗  ✗ ]
      ...

query 那一维要不要屏蔽？ 不需要在注意力里屏蔽。pad 位置作为 query 算出来的输出是垃圾，但它不会影响别的位置（别的位置不会看它），最后在损失里用 ignore_index 把这些位置排除即可（Week 0 Day 5 §6）。

两种掩码常常同时用。 decoder 训练时 = 因果掩码或 padding 掩码（逻辑或：任何一个说”屏蔽”就屏蔽）。

3.3 实现：加法掩码与布尔掩码

两种等价写法：

# 写法一：布尔掩码 + masked_fill（True = 屏蔽）
S = S.masked_fill(mask, float("-inf"))

# 写法二：加法掩码（0 = 保留，-inf = 屏蔽），直接加到分数上
S = S + additive_mask

加法掩码的好处是可以表达软偏置（不是 0 就是 $-\infty$ ，而是任意实数），ALiBi 位置编码就是往分数上加一个随距离线性下降的偏置（Day 5 §4）。

不同 API 的布尔约定相反，这是本周最容易踩的坑，见 §A.3。

3.4 整行被掩的陷阱：左补齐 + 因果掩码

Week 0 Day 4 §A.4 讲过：某一行被整行屏蔽时，softmax 算出 $0/0=$ NaN。当时的例子是”全是 padding 的样本”。在 decoder 里还有一个更隐蔽的触发方式。

场景。 批量生成时要左补齐（Week 0 Day 3 §A.6：生成靠”最后一个位置”，所以真实内容必须靠右）。假设一个样本左边补了 2 个 pad：

位置:        0    1    2    3    4
token:     pad  pad   A    B    C
因果 ∪ padding 掩码后，每行还剩几个可看的 key：
  query 0 (pad): 只能看位置 0，但位置 0 是 pad → 0 个
  query 1 (pad): 只能看 0、1，都是 pad         → 0 个
  query 2 (A):   能看 0、1、2，其中 2 可看      → 1 个

w1day1_checks.py 第 7 组：每行未被屏蔽的 key 数是 [0, 0, 1, 2, 3]，用 $-\infty$ 时第 0、1 行出现 NaN；用 $-10^9$ 时没有 NaN，这两行退化成均匀分布。

为什么这很危险？ 这两行是 pad 自己的输出，本来就会被丢掉，损失里也不会算它们。但即使损失只在真实位置上算，NaN 照样会流进梯度：§2.3 的 $\bar V=A^\top\bar Y$ 里，pad 行的 $\bar Y$ 是 0， $A$ 是 NaN，而 IEEE 规定 $\mathrm{NaN}\times0=\mathrm{NaN}$ ，于是所有位置的 $\bar V$ 都变成 NaN，然后污染所有参数。一次 NaN，整个模型报废。（w1day1_checks.py 第 7 组：pad 行的 $\bar Y$ 置 0 后， $\bar V$ 的 5 行全部是 NaN；code/w1d1_attention.py 自检第 9 项在 PyTorch 里复现同一件事。）

三种修法： (1) 用有限的大负数（torch.finfo(dtype).min 或 $-10^9$ ，fp16 下 $-10^4$ ）；(2) 让 pad query 至少能看到自己（放开对角线）；(3) 用 PyTorch 的 SDPA，它的融合内核对整行被掩的情况有专门处理（但不同后端行为不完全一致，别依赖它）。

3.5 其他常见掩码

掩码	形状	用在哪
因果	下三角	GPT 类 decoder 训练
padding	按 key 屏蔽	所有变长 batch
双向（不掩）	全 1	BERT 类 encoder（Day 6 §4）
prefix-LM	前缀部分双向、后半部分因果	T5 的一种变体、部分多模态模型：图像/指令部分互相全看，回答部分因果生成
文档块对角	块对角 × 因果	把多篇短文档打包进一个长序列训练时，防止文档 A 看到文档 B

最后一种在预训练里非常常见：为了不浪费算力在 padding 上，会把多篇文档首尾相接拼成固定长度。不加文档掩码的话，模型会学着从上一篇不相关的文档里”找线索”。很多预训练实现选择不加（认为影响不大），但近来的研究倾向于加上。

§4 多头注意力

4.1 动机：一个 softmax 只能给出一个分布

问题。 考虑”The animal didn’t cross the street because it was too tired”。处理”it”时，模型可能同时需要：

找它指代谁（→ “animal”，语义）
看前一个词是什么（→ “because”，局部句法）
找它的谓语（→ “was”）

单头注意力给”it”一个概率分布。如果把权重分给这三个位置，得到的是三者内容的同一个加权平均，每一维都按同样的比例混合。如果”指代”信息存在 value 的前几维、“句法”信息存在后几维，单头没办法从 A 取前几维、从 B 取后几维：它们被同一组权重搅在一起了。

验证。 w1day1_checks.py 第 8 组构造了一个最小的例子：目标是”从前一个词取前半维 + 从前两个词取后半维”。给单头最优的注意力权重和最优的线性映射（相当于允许 $W_V W_O$ 任意），10 个种子的相对误差中位数 0.690，接近 $1/\sqrt2=0.707$ ：最优解基本就是放弃其中一个位置，恰好丢掉一半能量。两个各 $d/2$ 维的头：误差 0，头 1 看前一个词取前半，头 2 看前两个词取后半。

4.2 直觉：并行地查多次表

既然一次查表只能表达一种”相关性”，那就同时查 $h$ 次，每次用不同的 Q/K/V 投影，让每个头学一种关系，最后把结果拼起来。

这和 CNN 的多通道是同一个思想：一个卷积核只检测一种模式，所以要很多个卷积核。

4.3 形式化与形状

$h$ 个头，每头维度 $d_k=d_v=d/h$ ：

$\mathrm{head}_i=\mathrm{Attention}(XW_Q^{(i)},\,XW_K^{(i)},\,XW_V^{(i)}),\qquad W_Q^{(i)},W_K^{(i)},W_V^{(i)}\in\mathbb R^{d\times d/h}$

$\mathrm{MHA}(X)=\mathrm{Concat}(\mathrm{head}_1,\dots,\mathrm{head}_h)\,W_O,\qquad W_O\in\mathbb R^{d\times d}$

实现上不会真的建 $h$ 组小矩阵。 把 $h$ 个 $W_Q^{(i)}$ 横向拼成一个 $d\times d$ 的 $W_Q$ ，一次矩阵乘算出所有头的 Q，再 reshape 切开：

B, T, d = x.shape
q = self.W_q(x)                                   # (B, T, d)
q = q.view(B, T, h, d // h).transpose(1, 2)       # (B, h, T, d_k)  ← 头维挪到前面
# k、v 同理
att = (q @ k.transpose(-2, -1)) / math.sqrt(d // h)   # (B, h, T, T)
att = att.masked_fill(mask, float("-inf")).softmax(-1)
y = att @ v                                       # (B, h, T, d_k)
y = y.transpose(1, 2).contiguous().view(B, T, d)  # (B, T, d)   ← 拼回去
y = self.W_o(y)

形状的每一步都值得停下来想：view(B, T, h, d_k) 把最后一维切成 $h$ 段，第 $i$ 段就是第 $i$ 个头；transpose(1, 2) 把头维挪到 batch 维旁边，于是后面的矩阵乘对 $(B,h)$ 两个维度自动广播， $h$ 个头的注意力一次算完；最后反过来拼回。.contiguous() 不能省，见 §A.5。

4.4 为什么是”切分”而不是”复制”

另一种可能：每个头都用完整的 $d$ 维， $h$ 个头参数量就是 $h$ 倍。

实际做法：总维度 $d$ 切成 $h$ 份，参数量和计算量与单头相同。w1day1_checks.py 第 9 组：MHA 参数量恒为 $4d^2+4d$ ，与 $h$ 无关。

这是一个预算决策：在固定算力下，用 $h$ 个低维的头换一个高维的头。换来的是 §4.1 的表达力（ $h$ 种独立的关系），付出的是每个头的维度（ $d_k$ 变小，每个头的打分只能用一个 $d_k$ 维子空间）。

$h$ 能不能无限加？ 不能。 $d_k=d/h$ 太小时， $W_QW_K^\top$ 的秩至多 $d_k$ ，头就无法表达复杂的匹配关系。Bhojanapalli et al.（2020）把这叫低秩瓶颈。实践中 $d_k$ 几乎总是 64 或 128：GPT-2 small $d=768,h=12$ ；Llama 2 7B $d=4096,h=32$ ；它们的 $d_k$ 分别是 64 和 128。

4.5 低秩视角：QK 电路与 OV 电路

把 $W_O$ 按行切成 $h$ 块 $W_O^{(i)}\in\mathbb R^{d_v\times d}$ ，拼接再乘 $W_O$ 等价于各头分别乘自己那块再求和：

$\mathrm{MHA}(X)=\sum_{i=1}^h\ \underbrace{\mathrm{softmax}\Big(\frac{XW_Q^{(i)}W_K^{(i)\top}X^\top}{\sqrt{d_k}}\Big)}_{\text{QK 电路：看谁}}\ X\ \underbrace{W_V^{(i)}W_O^{(i)}}_{\text{OV 电路：搬什么}}$

每个头由两个 $d\times d$ 的低秩矩阵刻画： $W_Q^{(i)}W_K^{(i)\top}$ （秩 $\le d_k$ ）决定注意力模式， $W_V^{(i)}W_O^{(i)}$ （秩 $\le d_k$ ）决定被选中的内容如何写回残差流。多头注意力 = $h$ 个独立的低秩”读-写”操作之和。 这是 Elhage et al.（2021）的分解，也是 Lecture 9 前沿话题里机械可解释性的起点。

一个直接推论： $W_Q$ 和 $W_K$ 单独都不可辨识，只有乘积 $W_QW_K^\top$ 有意义（对任意可逆 $R$ ， $W_QR$ 和 $W_KR^{-\top}$ 给出同样的打分）。所以别去分析单个 $W_Q$ 的结构。

4.6 收益与代价

收益： $h$ 种关系并行表达，参数和计算量不变。代价：(1) 每个头的秩受 $d_k$ 限制；(2) 推理时每个头都要缓存自己的 K 和 V，这笔显存随头数线性增长，是 Day 4 的 MQA/GQA 要解决的问题；(3) 很多头在训练后是冗余的，Michel et al.（2019）发现相当一部分头可以剪掉而几乎不掉点。

§5 复杂度与并行性

5.1 FLOPs

一层 MHA，一个长度 $T$ 的序列（乘加算 2 次浮点运算）：

步骤	矩阵乘形状	FLOPs
Q、K、V 投影	3 次 $(T\times d)(d\times d)$	$6Td^2$
$QK^\top$ （所有头合计）	$h$ 次 $(T\times d_k)(d_k\times T)$	$2T^2d$
$AV$ （所有头合计）	$h$ 次 $(T\times T)(T\times d_k)$	$2T^2d$
输出投影 $W_O$	$(T\times d)(d\times d)$	$2Td^2$
合计		$8Td^2+4T^2d$

两项相等时 $T=2d$ 。w1day1_checks.py 第 9 组：

$d$	$T$	注意力（ $T^2$ 项）占比
768	512	25.0%
768	1536	50.0%
768	8192	84.2%
4096	2048	20.0%
4096	8192	50.0%
4096	32768	80.0%

反直觉的结论：对 7B 级模型（ $d=4096$ ），序列长度 8k 以内，计算的大头是线性投影，不是注意力矩阵。再加上 FFN 的 $16Td^2$ （Day 2 §5）， $T^2$ 项要到更长才占主导。“注意力是 $O(T^2)$ 所以慢”这句话在短序列上是不准确的。

5.2 显存：真正先撞墙的地方

$A$ 是 $B\times h\times T\times T$ 的，而且 §2.3 说过反向要用它，必须存下来。w1day1_checks.py 第 10 组（bf16）：

配置	每层	全部层
GPT-2 small， $B=8,h=12,T=1024$	0.19 GiB	2.2 GiB
同上， $T=4096$	3.00 GiB	36.0 GiB
7B 级， $B=1,h=32,T=32768$	64.00 GiB	2048 GiB

$T$ 翻 4 倍，这一项翻 16 倍。长上下文训练首先撞上的不是算力，是显存。 FlashAttention 的解法是干脆不存 $A$ ，反向时分块重算（Day 6 §3），把这一项从 $O(T^2)$ 降到 $O(T)$ 。

5.3 路径长度与并行：Transformer 为什么能取代 RNN

Attention Is All You Need 的 Table 1 用三个指标比较了几种层（ $n$ 为序列长度）：

层类型	每层复杂度	串行操作数	最长路径
自注意力	$O(n^2d)$	$O(1)$	$O(1)$
循环	$O(nd^2)$	$O(n)$	$O(n)$
卷积（核宽 $k$ ）	$O(knd^2)$	$O(1)$	$O(\log_kn)$

串行操作数决定能不能并行：RNN 必须串行 $n$ 步。最长路径决定长程依赖好不好学：梯度在两个位置之间要穿过几层变换（Week 0 Day 4 §2.3 的连乘）。自注意力在两项上都是 $O(1)$ ，付出的是 $n^2$ 。在 GPU 上，用计算换并行极其划算。

5.4 为什么训练能并行，推理不能

训练时：整个目标序列是已知的（teacher forcing），因果掩码保证每个位置只看前文，于是 $T$ 个位置一次前向同时算完。

推理时：第 $t+1$ 个 token 要等第 $t$ 个 token 被采样出来才知道是什么，采样又依赖第 $t$ 步的输出。这是数据依赖，不是架构限制，任何自回归模型都逃不掉。

所以推理是 $T$ 次前向。朴素地做，第 $t$ 次前向要对整个长度为 $t$ 的前缀重算一遍，一层的代价是 $O(td^2+t^2d)$ ，生成 $T$ 个 token 累计 $O(T^2d^2+T^3d)$ 。KV Cache（Day 4）把前缀的 K、V 存下来，第 $t$ 步只为新 token 算投影、只让它一个 query 去看 $t$ 个 key，一层代价降到 $O(d^2+td)$ ，累计 $O(Td^2+T^2d)$ ，代价是显存。这就是 Week 0 Day 5 §5.1 埋下的那个锚点。

§A 注意力：API 用法与边界

A.1 `F.scaled_dot_product_attention`（SDPA）

import torch.nn.functional as F
y = F.scaled_dot_product_attention(
    q, k, v,                  # (..., L, E)、(..., S, E)、(..., S, Ev)，通常是 (B, h, T, d_k)
    attn_mask=None,           # bool（True = 参与注意力！）或 float（加到分数上）
    dropout_p=0.0,            # ★ 不会自动感知 train/eval
    is_causal=False,          # 与 attn_mask 不能同时给
    scale=None,               # 默认 1/sqrt(E)
)

PyTorch 2.0 起提供，会根据输入自动选 FlashAttention / memory-efficient / 数学实现三种后端之一。能用它就别手写，快且省显存。五个边界：

布尔 attn_mask 的 True 表示”参与”，和 masked_fill 的直觉（True = 填充 = 屏蔽）以及 nn.MultiheadAttention（True = 屏蔽）相反。见 A.3 的对照表。
dropout_p 不看 self.training。模块里必须写 dropout_p=self.p if self.training else 0.0，否则 eval 时仍在 dropout，结果随机且变差。
is_causal=True 用的是”左上角对齐”的下三角：官方参考实现是 torch.ones(L, S).tril(diagonal=0)。当 $L\ne S$ 时（比如带 KV Cache 的解码， $L=1$ 、 $S=$ 前缀长度），它只允许看第一个 key，完全错误。增量解码时要么 is_causal=False（单个新 token 本来就能看所有旧 key），要么自己传 attn_mask。这是 Day 4 写 KV Cache 时的头号坑。
attn_mask 和 is_causal=True 同时给会报错。要”因果 + padding”就自己合成一个 attn_mask。
融合内核对输入有要求（dtype、head_dim、是否有 mask），不满足时静默回退到数学实现，不报错但变慢。可以用 torch.nn.attention.sdpa_kernel(...) 上下文强制指定后端来检查。

A.2 `nn.MultiheadAttention`

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, dropout=0.0, batch_first=False)
out, weights = mha(query, key, value,
                   key_padding_mask=None,    # (B, S) bool，True = 忽略这个 key
                   need_weights=True,        # ★ 默认要返回注意力权重
                   attn_mask=None,           # (L, S) 或 (B·h, L, S)，bool 时 True = 不允许
                   average_attn_weights=True)

batch_first 默认是 False：输入要是 (T, B, d)。直接喂 (B, T, d) 不报错，它会把 batch 当成序列、序列当成 batch，结果完全错误。永远显式写 batch_first=True。
自注意力要把同一个张量传三次：mha(x, x, x)。它的接口是为交叉注意力设计的。
need_weights=True 是默认值，返回的权重还默认对头取了平均。要各头的权重得 average_attn_weights=False；不需要权重时设 need_weights=False，否则会阻止它走融合的快速路径。
is_causal 在这里只是”提示”：真正起作用的仍是 attn_mask，只给 is_causal=True 不给 mask，多数版本会直接报错。
Q/K/V 三个投影被合并存成一个 in_proj_weight，形状 (3·embed_dim, embed_dim)。加载别人的权重或做 LoRA 时要注意这个布局。

本周的建议：手写 MHA 用于学习（§B），生产用 SDPA，nn.MultiheadAttention 知道它的坑就行。

A.3 掩码约定对照表 ★

API	布尔 True 表示	加法（float）掩码
`Tensor.masked_fill(mask, v)`	填充（通常用来屏蔽）	—
`F.scaled_dot_product_attention(attn_mask=)`	参与注意力	加到分数上
`nn.MultiheadAttention(attn_mask=)`	不允许参与	加到分数上
`nn.MultiheadAttention(key_padding_mask=)`	忽略这个 key	加到分数上
HuggingFace `attention_mask`	1 = 真实 token，0 = pad	模型内部转成加法掩码

同一个布尔张量，传给 SDPA 和传给 nn.MultiheadAttention，含义正好相反。 这类错误不报错，只让模型悄悄学不好。最稳妥的习惯：统一用加法掩码（0 保留、 $-\infty$ 或极大负数屏蔽），它在所有 API 里含义一致。

A.4 构造因果掩码

# True = 屏蔽（配合 masked_fill）
causal = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
# True = 允许（配合 SDPA 的布尔 attn_mask）
allowed = torch.tril(torch.ones(T, T, dtype=torch.bool))
# 在模块里注册成 buffer：跟着 .to(device) 走，但不是参数
self.register_buffer("mask", causal, persistent=False)

triu(diagonal=1) 是不含对角线的上三角。写成 diagonal=0 会把对角线也屏蔽，第一个位置整行被掩，直接 NaN（§3.4）。
persistent=False 让它不进 state_dict：掩码可以重建，存进 checkpoint 只会在改了最大长度后导致加载失败。
预先建一个 max_len × max_len 的掩码，前向时切片 self.mask[:T, :T]。

A.5 形状操作：`view`、`transpose`、`contiguous`

y = y.transpose(1, 2)              # (B, h, T, dk) → (B, T, h, dk)，只改 stride，内存没动
y = y.contiguous().view(B, T, d)   # 必须先 contiguous 才能 view
y = y.reshape(B, T, d)             # 或者直接 reshape（不连续时它会自动复制）

transpose 之后张量在内存里不再是连续的，view 会报 “view size is not compatible with input tensor’s size and stride”（Week 0 Day 5 §A.2）。最隐蔽的错误是 reshape 的顺序写错：y.view(B, T, d) 而不先 transpose，形状对了，内容却把不同位置的头拼在了一起，不报错，只是学不好。检查方法：对一个 $h=1$ 的情况，多头实现应该与单头实现逐元素相等（§B 第 3 题）。

A.6 `einsum`：让形状一目了然

att = torch.einsum("bhqd,bhkd->bhqk", q, k) / math.sqrt(dk)
y   = torch.einsum("bhqk,bhkd->bhqd", att, v)

下标就是文档：q、k 两个不同字母提醒你 query 长度和 key 长度可以不同（交叉注意力、KV Cache）。性能上 einsum 通常会被转成同样的 bmm，可以放心用。

§B 实操

跑通 w1day1_checks.py，11 组结论都看懂。第 4 组（秩坍缩）和第 8 组（单头 vs 多头）最值得停下来想：它们回答的是”为什么要残差”和”为什么要多头”。
纸笔推一遍 §2.3 的反向（1h，闭卷），推完对照五行代码。重点是 softmax 那一步的 $\delta_{jl}$ 塌缩和 $D_i=\bar y_i^\top y_i$ 的来历。
手写多头注意力（2h），不许用 nn.MultiheadAttention 和 SDPA：
- 支持因果掩码和 padding 掩码（合成一个掩码）
- assert 每行权重和为 1，被掩位置权重为 0
- 验证 $h=1$ 时与单头实现逐元素相等；验证置换等变（无掩码时）
- 再与 F.scaled_dot_product_attention 对拍，亲手处理一次布尔掩码的约定翻转
手写注意力的 autograd.Function（1h）：用 §2.3 的公式写 backward，gradcheck 通过（float64）。
复现左补齐的 NaN（0.5h）：构造左补齐 + 因果掩码的 batch，用 -inf 看到 NaN 流进梯度，再用两种修法修好。

参考实现：code/w1d1_attention.py（--selftest 覆盖以上全部）。

§C 自测

动机与直觉层

自注意力同时解决了 Week 0 留下的哪两个问题？它和 Bahdanau 注意力的区别只有什么？
为什么需要 $W_Q$ 、 $W_K$ 、 $W_V$ 三个投影？至少说出三个理由。 $W_Q=W_K$ 会怎样？
什么是置换等变？它为什么是自注意力的缺陷？因果掩码能完全弥补吗？
注意力输出为什么落在 value 的凸包里？这对 Transformer 块的设计意味着什么？
为什么要因果掩码？在 softmax 之后置零为什么不行？
单头注意力为什么无法同时”从 A 取一些维度、从 B 取另一些维度”？多头如何解决？
多头为什么是”切分”而不是”复制”？ $h$ 能无限增大吗？

推导层

证明 $f(PX)=Pf(X)$ 。
索引级推出 $\bar V,\bar A,\bar S,\bar Q,\bar K$ ，并证明 $D_i=\sum_j\bar A_{ij}A_{ij}=\bar y_i^\top y_i$ 。
用 $\bar S_{il}=A_{il}(\bar A_{il}-D_i)$ 说明被掩位置的梯度为什么自动为零，以及为什么 $\sum_l\bar S_{il}=0$ 。
PyTorch 默认初始化下，初始注意力分数的标准差是多少？推出来。
写出一层 MHA 的 FLOPs，求投影项与注意力项相等时的 $T$ 。
把 MHA 写成 $h$ 个”QK 电路 × OV 电路”之和，说明为什么 $W_Q$ 单独不可辨识。

工程与 API 层

SDPA 的布尔 attn_mask 与 nn.MultiheadAttention 的布尔 attn_mask，True 各表示什么？
SDPA 的 dropout_p 为什么要手动根据 self.training 设置？
带 KV Cache 解码时（ $L=1$ ），is_causal=True 会发生什么？应该怎么写？
nn.MultiheadAttention 不写 batch_first=True 会怎样？报错吗？
左补齐 + 因果掩码为什么会 NaN？给出三种修法。
多头 reshape 时为什么必须先 transpose 再合并？漏了会报错吗？
torch.triu(..., diagonal=0) 和 diagonal=1 做因果掩码，差别是什么？