5-A2S2-attention.ipynb 逐步说明

这份 notebook 在 2/1-class-prog.md 里已经讲过整体结构。这里补的是更具体的“每一步到底在实现什么”。

它和后面第 3 课的最小 LoRA 实验有非常清楚的对应关系：

• 这里你手写 pairwise_similarities、mask、split_heads
• 后面在 course_lora_qwen_src 里你虽然不再手写 attention，但你改的 q_proj / k_proj / v_proj / o_proj 正是这些操作背后的模块

2.1 部分：从零实现 Attention

Cell 8：init_qkv_proj(...) 和 self_attention(...)

init_qkv_proj(n_embd:int)

return (nn.Linear(n_embd, n_embd), nn.Linear(n_embd, n_embd), nn.Linear(n_embd, n_embd))

这说明：

• Q、K、V 一开始都来自同一个 embedding 维度
• 但会通过 3 个不同的线性层投影到不同的表示空间

这和你们第 3 课里看到的：

target_modules=['q_proj', 'k_proj', 'v_proj', ...]

可以直接对上。

self_attention(...)

这一格是整份作业的总装函数。你会反复回来补它。

它的结构本质上是：

1. 如果多头，就先 split heads
2. 算 QK^T
3. 做 scaling
4. 如果是语言模型，就加 causal mask
5. 做 softmax
6. 和 V 相乘得到输出
7. 如果多头，就 merge heads

这就是 Transformer block 里 self-attention 的主链。

Cell 10：Step 1 的核心函数

pairwise_similarities(Q, K)

这一格让你实现：

A = QK^T

它真正想让你注意的是：

• K 不是直接乘，而是最后两维转置
• attention matrix 的 shape 来自 query positions x key positions

attn_scaled(A, n_embd, n_heads)

这里要求你实现标准 scaling：

$$\frac{A}{\sqrt{d_k}}$$

其中：

• d_k = n_embd / n_heads

为什么要除以它？

• 维度越大，点积值越容易变大
• softmax 太尖锐，训练会不稳定

这和前面 A2S1 故意把 scaling 去掉，刚好形成对照。

attn_softmax(A)

这一步很简单，但有一个最容易犯错的点：

• softmax 要作用在最后一个维度 dim=-1

因为你是在“对每个 query，看所有 key 的分布”。

compute_outputs(A, V)

这里要求实现：

O = AV

也就是用注意力权重去加权求和 value。

Cell 14：Step 2，Causal Mask

make_causal_mask(n_tok)

目标是构造一个下三角 mask，让每个位置只能看自己和过去。

为什么一定要这样？

• 语言模型做 next-token prediction
• 不能偷看未来 token

apply_causal_mask(mask, A)

通常的做法是：

• 合法位置保留原分数
• 非法位置改成 -inf

这样 softmax 后，未来位置的概率就变成 0。

这一步和后面第 3 课的所有 causal LM 微调完全同源。

Cell 18：Step 3，多头注意力

split_heads_qkv(Q, K, V, n_heads)

这里的真正难点不是数学，而是 shape。

你要学会把：

• (B, T, C)

变成类似：

• (B, H, T, C/H)

也就是说：

• batch 维保留
• token 维保留
• 新增一个 head 维
• 每个 head 拿到更小的子维度

多头为什么不是“多做几次 attention”？

因为还涉及：

• reshape
• transpose
• merge

这些张量组织方式的变化，决定了每个 head 真正处理的是 embedding 的哪一部分。

2.2 部分：在 mini Transformer 中做实验

从这里开始，notebook 从“手写 attention 小练习”进入“放进最小 Transformer 里跑训练”。

Cell 27：下载数据

这部分和课程目录里的：

• N-gram.zip
• mingpt-cse447.zip

直接对应。你们在课程页里已经准备好了本地副本，就是为了让这一步不依赖在线下载。

Cells 28-35：数据与词表处理

• 导入 mingpt.model.GPT
• 读训练 / 开发文本
• 建词表
• 准备 X 和 Y

这部分很值得和第 3 课 course_lora_qwen_src/data/*.json 对照：

• 这里是经典语言模型 token 序列数据
• 第 3 课是 chat / instruction 数据

形式不同，但本质都是“把原始文本组织成 next-token prediction 数据”。

Cells 37-40：模型与训练器

• model_config = GPT.get_default_config()
• train_config = Trainer.get_default_config()
• model.to(DEVICE)
• plt.plot(log)

你可以把它理解成第 2 课版本的“最小训练闭环”。

Cells 42-47：评估和生成

• 单句 loss / perplexity
• 训练集与开发集 loss
• 调用生成接口

这和后面你们第 3 课里训练完后立刻做推理，是同样的节奏：

训练不是终点，验证行为才是。

这份 notebook 最值得学生掌握的 4 个点

1. QK^T、scaling、softmax、加权和是 attention 的四步主链
2. causal mask 是语言模型和普通 attention 的关键分界线
3. multi-head attention 的难点主要在 shape
4. 最终这些实现要能回到一个真实可训练的 Transformer 里