第 1 课 · 约 15–20 分钟 · 2026-10-25（北京时间）

缩放点积注意力：亲手写出 GPT 的心脏

Scaled dot-product self-attention —— Q/K/V、√d_k、softmax、因果掩码，以及每一步的张量形状。

为什么从这里开始？你的目标第一条就是"能在白板上推导 scaled dot-product attention 并解释张量形状"，而 nanoGPT 里整个 CausalSelfAttention 的核心就是这 4 行见 nanoGPT model.py 中 CausalSelfAttention.forward 的 "manual implementation" 分支。。后面的多头、KV cache、FlashAttention 全都是在这几行上做文章。

1 · 一句话模型

注意力就是加权平均：每个位置拿自己的 query 去和所有位置的 key 比相似度，相似度经 softmax 变成权重，再用权重对 value 求平均论文原话："The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key." —— Vaswani et al. §3.2。

在自注意力里，Q、K、V 都来自同一个输入 X：Q = X W_Q，K = X W_K，V = X W_V。三个 W 是可学习参数——这就是"注意力学到了什么"的全部所在。

2 · 五步，盯住形状

设序列长度 T=4，d_model=8，d_k=d_v=16（就是练习里的数字）：

口诀：(T, d) 进，(T, T) 中间，(T, d) 出。那个 T×T 的中间矩阵就是注意力 O(T²) 复杂度的来源论文 Table 1：self-attention 每层复杂度 O(n²·d)。长上下文、FlashAttention 都是围绕它展开的，后面的课会回来。。

3 · 为什么除以 √d_k

假设 q、k 的每个分量独立、均值 0、方差 1。那么 q·k = Σ q_ik_i 是 d_k 个"均值 0、方差 1"的项之和，方差是 d_k这正是论文 §3.2.1 的脚注。Karpathy 在视频里用 head_size=16 现场演示了同样的事。。d_k=512 时，分数的标准差约 22，几个分数之间差几十，softmax 就几乎变成 one-hot——论文的说法是把 softmax 推进了"梯度极小的区域"。除以 √d_k，方差拉回 1。

你等下会在练习输出里亲眼看到：未缩放时 8 个 key 中最大的那个权重平均占 0.958，缩放后只有 0.365。

动手玩：拖动 d_k，看 softmax 怎么变尖

下面是 6 个 token 的注意力权重矩阵（行 = query，列 = key，颜色越深权重越大）。Q、K 的每个分量都是固定种子的 N(0,1) 随机数，d_k 取前 d_k 维。

看完记住一句话：不缩放时分数标准差 ≈ √d_k，d_k 越大 softmax 越接近 one-hot；缩放后标准差 ≈ 1，分布形状和 d_k 基本无关。打开掩码后，被挡住的格子权重变成 0，每行剩下的权重重新归一化，和仍然是 1。

4 · 因果掩码：不许偷看答案

GPT 训练时把整段文本一次喂进去，并行预测每个位置的下一个 token。如果位置 2 能看到位置 3，它就是在抄答案。所以在 softmax 之前，把分数矩阵对角线以上全部设为 −∞；e^−∞ = 0，softmax 后这些位置权重恰好为 0，其余权重仍然和为 1论文 §3.2.3："masking out (setting to −∞) all values in the input of the softmax which correspond to illegal connections." 3Blue1Brown 的视频 11:08 有很直观的动画。。

5 · 动手：补完 3 个 TODO 实战

1. 确认有 NumPy：python3 -c "import numpy; print(numpy.__version__)"（没有就 pip install numpy）。
2. 把下面代码存成 0001_attention_todo.py（或直接下载 0001_attention_todo.py），补完 TODO 1/2/3（每个一行）。卡住时可对照 参考解答。
3. 运行 python3 0001_attention_todo.py，和"预期输出"逐行对照。随机种子固定，数字应完全一致。

# 第 1 课练习：补完 3 个 TODO，然后 python 0001_attention_todo.py
import numpy as np

def softmax(x):
    x = x - x.max(axis=-1, keepdims=True)          # 数值稳定：先减最大值
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V, causal=False):
    d_k = Q.shape[-1]
    scores = ...      # TODO 1：Q 和 K 的点积，再缩放。形状应为 (T, T)
    if causal:
        T = scores.shape[-1]
        mask = ...    # TODO 2：布尔矩阵，True 表示"未来位置"（提示：np.triu(..., k=1)）
        scores = np.where(mask, -np.inf, scores)
    weights = ...     # TODO 3：对哪一个轴做 softmax？
    return weights @ V, weights                        # (T, d_v), (T, T)

rng = np.random.default_rng(0)
T, d_model, d_k, d_v = 4, 8, 16, 16
X = rng.standard_normal((T, d_model))
W_q, W_k, W_v = (rng.standard_normal((d_model, d)) / np.sqrt(d_model) for d in (d_k, d_k, d_v))
Q, K, V = X @ W_q, X @ W_k, X @ W_v

out, w = attention(Q, K, V, causal=True)
print("Q", Q.shape, "K", K.shape, "V", V.shape)
print("out", out.shape, "weights", w.shape)
print("row sums:", w.sum(axis=-1).round(3))
print("causal weights:\n", w.round(2))

# 实验：为什么要除以 sqrt(d_k)
for d in (16, 512):
    q, k = rng.standard_normal((2, 10000, d))
    dots = (q * k).sum(-1)
    print(f"d_k={d:4d}  var(q·k)={dots.var():7.1f}  var(q·k/sqrt(d_k))={(dots/np.sqrt(d)).var():.2f}")
S = rng.standard_normal((1000, 8)) * np.sqrt(512)     # 1000 行未缩放分数，每行 8 个 key
print("avg max weight  unscaled:", softmax(S).max(-1).mean().round(3),
      " scaled:", softmax(S / np.sqrt(512)).max(-1).mean().round(3))

Q (4, 16) K (4, 16) V (4, 16)
out (4, 16) weights (4, 4)
row sums: [1. 1. 1. 1.]
causal weights:
 [[1.   0.   0.   0.  ]
 [0.23 0.77 0.   0.  ]
 [0.19 0.33 0.49 0.  ]
 [0.16 0.37 0.21 0.26]]
d_k=  16  var(q·k)=   16.3  var(q·k/sqrt(d_k))=1.02
d_k= 512  var(q·k)=  511.0  var(q·k/sqrt(d_k))=1.00
avg max weight  unscaled: 0.958  scaled: 0.365

已在 NumPy 2.4 / 2.5（Python 3.11 / 3.13）上实测。前两行检查形状，第三行检查 softmax 轴，矩阵的上三角全 0 检查掩码，最后三行就是第 3 节的 √d_k 实验。

scores = Q @ K.swapaxes(-2, -1) / np.sqrt(d_k)
mask = np.triu(np.ones((T, T), dtype=bool), k=1)
weights = softmax(scores)

加分题（5 分钟，需要 PyTorch）

用 PyTorch 写批量多头版，并和官方 F.scaled_dot_product_attention 对拍——这正是 nanoGPT 在 PyTorch ≥ 2.0 时走的那条路。

下载：0001_bonus_torch_check.py

# 加分题：用 PyTorch 官方实现对拍（pip install torch）
import math, torch
import torch.nn.functional as F

torch.manual_seed(0)
B, H, T, d_k = 2, 4, 5, 16                       # batch, heads, seq_len, head_dim
q, k, v = torch.randn(3, B, H, T, d_k).unbind(0)

def my_attention(q, k, v):
    scores = q @ k.transpose(-2, -1) / math.sqrt(k.size(-1))          # (B, H, T, T)
    causal = torch.tril(torch.ones(T, T, dtype=torch.bool))
    scores = scores.masked_fill(~causal, float("-inf"))
    return F.softmax(scores, dim=-1) @ v                              # (B, H, T, d_k)

mine = my_attention(q, k, v)
ref = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print("shape:", tuple(mine.shape))
print("allclose:", torch.allclose(mine, ref, atol=1e-6))

shape: (2, 4, 5, 16)
allclose: True

注意形状从 (T, d) 变成 (B, nh, T, hs)，但代码一个字都不用改——@ 和 transpose(-2, -1) 只作用在最后两维，前面的维度自动当作批量。这就是下一课"多头注意力"的伏笔。

6 · 检索练习

不看上文，凭记忆作答。点一次就锁定。

7 · 面试会怎么问

Q1. 请写出 self-attention 的公式并标注每一步的形状。为什么要除以 √d_k？如果改成除以 d_k 会怎样？

Q2. Decoder 为什么需要因果掩码？推理时每次只生成一个 token，还需要掩码吗？

8 · 推荐主要来源

间隔复习

两天后回来，只做第 6 节的测验（点"再测一次"会打乱选项），然后不看资料在纸上默写第 5 节的 3 行答案。能写出来，这节课才算进了长期记忆。