第 9 课：注意力机制 —— 让 token 互相"看"

代码位置：src/attention.rs（MultiHeadAttention::forward 中的投影 / 打分 / softmax / 掩码部分） 底层算子：src/tensor.rs（softmax_last_dim / matmul / permute / mul_scalar） 掩码构造：src/model.rs（GPT::forward）

1. 本课要搞懂的问题

1. 为什么一个 token 需要"看"别的 token？注意力到底在算什么？
2. Query / Key / Value 是什么？为什么不能直接用 x 打分？
3. 为什么要除以 √d_k？softmax 在干什么？
4. 为什么 GPT 只能看过去？-inf 掩码怎么实现？

2. 从词向量到"上下文向量"

第 12 课会把每个 token 嵌入成一个向量，但 Embedding 查表得到的向量只代表这个 token 自己：它不知道周围发生了什么。"bank" 到底是"河岸"还是"银行"，单看它自己无法判断，必须结合上下文。

注意力要做的事：让每个位置根据内容，从其他位置"收集"信息。输出不再是"我自己的向量"，而是"我 + 我关注的别人的向量的加权混合"。

3. 检索的视角：Query / Key / Value

注意力最直观的类比是图书馆检索：

角色	类比	数学含义
Query（查询）Q	你输入的"关键词"	我想要什么样的信息
Key（键）K	每本书的"标签"	我能提供什么样的信息
Value（值）V	书的内容	真正被取走的信息

流程：拿 query 去和所有 key 比对相似度 → 得到权重 → 按权重把 value 加权求和。

我们的模型里 Q、K、V 都来自同一个输入 x（每个位置既当"提问者"又当"被看者"），这叫自注意力（Self-Attention）。

4. Q / K / V 投影：三个角色各学各的表示

直接用 x 当 Q、K、V 行不行？可以，但不好——三个角色应该学不同的表示。所以每个角色先过一层可学习的线性投影（Linear: y = xW + b）：

// 1. 投影得到 Q、K、V
let q = self.c_q.forward(x).reshape(vec![b, t, d]); // [B, T, D]
let k = self.c_k.forward(x).reshape(vec![b, t, d]);
let v = self.c_v.forward(x).reshape(vec![b, t, d]);

• c_q / c_k / c_v 是三个独立的 Linear::new(D, D)（MultiHeadAttention 结构体里的字段），各有自己的权重和偏置，在训练中学习
• Linear::forward 对 3D 输入会自动展平成 [B*T, D] 计算再还原 3D；这里的 reshape(vec![b, t, d]) 保证形状明确是 [B, T, D]
• D = n_embd。模型配置 GPTConfig::tiny：n_embd = 64，n_head = 4 → head_dim = 16（第 10 课讲怎么拆头；本课先假设只有一个头，公式层面 H = 1）

5. 打分：scores = Q·Kᵀ / √d_k

query 和 key 的相似度用点积衡量：q_i · k_j 越大越相似。一次性算所有两两相似度，就是矩阵乘法：

S = Q · Kᵀ / √d_k

• Q: [T, d_k]，K: [T, d_k] → Kᵀ: [d_k, T] → S: [T, T]
• S[i][j] = (q_i · k_j) / √d_k：第 i 个查询对第 j 个键的"原始分数"

对应代码：

// 4. 注意力分数：scores = Q·Kᵀ / √d_k
let scale = 1.0 / (head_dim as f32).sqrt();
let kt = k.permute(&[0, 2, 1]);                     // K 转置：[B*H, head_dim, T_total]
let scores = q.matmul(&kt).mul_scalar(scale);       // [B*H, T, T_total]

• permute(&[0, 2, 1])：把 K 的 [T, head_dim] 两维交换，得到数学上的 Kᵀ
• q.matmul(&kt)：3D 批量矩阵乘法（第 1 课学的 matmul，这里按"批"逐矩阵相乘）
• mul_scalar(scale)：每个元素乘以 1/√d_k

5.1 手算例子（d_k = 2，两个 token）

Q = [ [1, 0],          K = [ [1, 0],          V = [ [1, 0],
      [0, 1] ]               [0, 1] ]               [0, 1] ]

Q·Kᵀ = [ [1, 0],
         [0, 1] ]

除以 √2 后：S = [ [0.707, 0    ],
                 [0,     0.707] ]

softmax 按行归一化（见第 7 节）后，第 1 行 ≈ [0.670, 0.330]，第 2 行 ≈ [0.330, 0.670]——即"每个位置主要关注自己"（对角线），很符合直觉。

6. 为什么要除以 √d_k

假设 q、k 每个维度独立、均值 0、方差 1。两个 d_k 维向量的点积：

Var(q·k) = Var(Σ q_i·k_i) = d_k          →  标准差 = √d_k

d_k	点积的标准差	后果
小（如 2）	1.4	softmax 分布温和，梯度健康
大（如 64）	8	点积数值大、分布"尖"，softmax 趋向 one-hot（一个 1 其余 0），梯度趋近 0，学不动

• 除以 √d_k 把方差拉回 1，softmax 保持"温和"的分布
• 所以这个公式叫 Scaled Dot-Product Attention——缩放点积注意力

7. softmax 归一化：把分数变成权重

原始分数可正可负、量纲不定，不能直接当权重。softmax 把每一行变成非负且和为 1 的概率分布：

a_ij = exp(S_ij) / Σ_j' exp(S_ij')

含义：第 i 个查询给第 j 个键的"注意力权重"，一行加起来等于 1。

数值稳定技巧

exp 对大的正数会爆炸（如 exp(1000) = 无穷大），所以先减去每行最大值——数学上不改变结果（分子分母同除 exp(max)）：

// softmax_last_dim：沿最后一维逐行独立归一化
for r in 0..rows {
    let mut maxv = f32::NEG_INFINITY;
    for j in 0..d {
        maxv = maxv.max(sd[r * d + j]);                    // 1. 找行最大值
    }
    let mut sum = 0.0;
    for j in 0..d {
        out_data[r * d + j] = (sd[r * d + j] - maxv).exp(); // 2. 减 max 再 exp
        sum += out_data[r * d + j];
    }
    for j in 0..d {
        out_data[r * d + j] /= sum;                        // 3. 归一化
    }
}

稳定版公式：

softmax(x)_i = exp(x_i - max x) / Σ_j exp(x_j - max x)

反向传播公式（softmax_last_dim 的 backward 里已实现）：∂x_i = s_i·(g_i − Σ_j g_j·s_j)，其中 s 是 softmax 输出、g 是上游梯度。这正是第 2 课自动微分的用武之地。

8. 加权求和：out = A·V

把权重矩阵 A 作用到 V 上：

out_i = Σ_j a_ij · v_j

即"每个输出 = 所有位置 V 的凸组合"（权重非负、和为 1）。代码只有两行：

// 6. softmax 得到注意力权重，加权求和
let attn = scores.softmax_last_dim();  // [B*H, T, T_total]   注意力权重
let out = attn.matmul(&v);             // [B*H, T, head_dim]  加权求和

至此，注意力 = 内容相关的可学习加权平均：每个位置根据 query 与各 key 的相似度，从所有 value 里"取"信息。

9. 因果掩码：为什么 GPT 只能看过去

GPT 是自回归语言模型：预测第 t 个 token 时，条件只能是它之前的 token：

P(token_t | token_1, ..., token_{t-1})

• 训练时我们把整段文本并行喂入（这样才快！）。如果不加限制，位置 t 会"偷看"位置 t+1 的答案——信息泄漏，模型退化成"抄答案"
• 解决办法：让位置 i 的查询看不到 j > i 的键——这就是因果掩码（causal mask）

掩码长什么样

GPT::forward 里构造（t_total = t + base，训练时 base = 0）：

let mut mask_data = vec![0.0f32; t * t_total];
for i in 0..t {
    for j in 0..t_total {
        if j > i + base {
            mask_data[i * t_total + j] = f32::NEG_INFINITY;
        }
    }
}
let mask = Tensor::from_vec(mask_data, vec![t, t_total]);

以 t = 4 为例，mask 是这样一个矩阵：

查询 i \ 键 j	j=0	j=1	j=2	j=3
i=0	0	-inf	-inf	-inf
i=1	0	0	-inf	-inf
i=2	0	0	0	-inf
i=3	0	0	0	0

对角线及以下（j ≤ i）是 0，未来位置（j > i）是 -inf。

为什么是 -inf 而不是把列删掉

softmax 里 exp(-inf) = 0，被掩码位置的注意力权重精确为 0——"看不见"等价于"权重为 0"。同时张量形状保持不变，可以继续走统一的批量矩阵运算。

let scores = scores.add(mask);              // [B*H, T, T_total] + 广播 [T, T_total]
let attn = scores.softmax_last_dim();       // -inf -> exp(-inf) = 0 -> 权重为 0

scores.add(mask) 用到了第 3 课的广播：mask 形状 [T, T_total] 自动广播到 [B*H, T, T_total] 的每一个"批"（每个头共享同一张掩码）。 （代码里的 base 是 KV cache 引入的偏移，第 18 课讲解；训练时 base = 0。）

10. 公式 ↔ 代码 总对照

Attention(Q, K, V) = softmax( (Q·Kᵀ / √d_k) + M ) · V

数学	Rust 代码（MultiHeadAttention::forward）
Q = x·W_Q	self.c_q.forward(x).reshape(vec![b, t, d])
K = x·W_K，V = x·W_V	c_k / c_v 同理
Kᵀ	k.permute(&[0, 2, 1])
Q·Kᵀ / √d_k	q.matmul(&kt).mul_scalar(scale)，scale = 1.0 / (head_dim as f32).sqrt()
+ M（因果掩码）	scores.add(mask)
softmax	scores.softmax_last_dim()
·V	attn.matmul(&v)

11. 运行与测试

cargo test   # tensor 的 test_softmax / test_matmul_2d / test_matmul_3d 验证本课用到的底层算子
cargo run    # 演示 3 训练小 GPT：每一步前向/反向都跑在本课的注意力计算上

12. 动手练习

1. 手算：取第 5.1 节的 2×2 例子，写出 softmax 后每一行的权重，再算出最终的加权求和结果，对照公式逐项核对。
2. 修改 GPT::forward 里的掩码条件 j > i + base 为 j >= i + base（禁止"看自己"），思考对训练会有什么影响（提示：每个位置的预测将只能依赖更早的 token）。
3. 临时改 softmax_last_dim：去掉"减最大值"那一步，用大数（如 1000）做输入，观察溢出/精度问题，验证数值稳定技巧的必要性。
4. 阅读 softmax_last_dim 的反向代码，对照公式 ∂x_i = s_i·(g_i − Σ_j g_j·s_j) 逐行核对。

13. 本课总结

• 注意力 = 内容相关的加权平均：Q 问、K 答、V 给内容
• scores = Q·Kᵀ/√d_k：点积打分，缩放防止 d_k 大时 softmax 饱和
• softmax：分数 → 权重（非负、和为 1）；先减 max 保证数值稳定
• 因果掩码：未来位置填 -inf → softmax 后权重为 0 → 模型只能看过去
• 下一课：把注意力"切成多份"并行计算——多头注意力！