""" ============================================================================= CS349E · 第 2 讲 · 02_attention_from_scratch.py —— 注意力从零开始：软查表、√d 缩放、因果掩码、多头、位置编码 ============================================================================= 【这个脚本要回答的问题】 Attention Is All You Need（本讲阅读材料）里的每一个设计都有理由。这里一个一个用数字验证： 演示 1：注意力 = "软"的字典查找：query 和哪个 key 像，就主要取哪个 value；缩放决定"软"到什么程度。 演示 2：为什么除以 √d：不缩放时点积的方差随 d 线性增长，softmax 饱和、梯度几乎为 0（论文脚注 4）。 演示 3：因果掩码：改掉第 t 个 token，前面位置的输出一个比特都不变——这是 KV cache 能成立的前提。 演示 4：多头：把 d 切成 h 份并行做注意力，运算量和单头一样；GQA 让几个 query 头共用一组 K/V。 演示 5：位置信息：没有位置编码、没有掩码时注意力对顺序"视而不见"（置换等变）； 正弦位置编码和 RoPE 都让"相对位置"变成一个线性 / 旋转关系；用 RoPE 手工搭一个"看前一个 token"的头。 演示 6：RNN vs 自注意力：每层运算量 n·d² vs n²·d、必须顺序执行的步数 n vs 1（论文表 1），以及 numpy 上的实测。 【怎么读输出】演示 1–5 是确定性的计算；演示 6 的墙钟时间和机器有关（编写时用的是 Apple M4 Pro），只看趋势。 【运行】python3 02_attention_from_scratch.py （约 3 秒） 【依赖】numpy；同目录的 transformer_lab；从 ../lecture01 导入 console。 """ import math import os import sys import time import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_num, fmt_time, print_header, print_observations, print_table # noqa: E402 from transformer_lab import (TINY, apply_rope, attention_core, forward, init_params, # noqa: E402 rope_angles, softmax) # --------------------------------------------------------------------------- # 演示 1：软查表 # --------------------------------------------------------------------------- def demo1_soft_lookup(): """ 一个 Python 字典是"硬"查找：key 完全相等才取出 value。注意力把它变"软"： 权重 w_i = softmax_i(q·k_i / τ)，输出 = Σ_i w_i · v_i q 和哪个 k 的点积大，就主要取哪个 v。τ（温度）越小越接近硬查找，越大越接近"所有 value 的平均"。 """ print_header("演示 1：注意力 = 软的字典查找") rng = np.random.default_rng(0) d = 16 keys = rng.standard_normal((5, d)) values = np.array([10.0, 20.0, 30.0, 40.0, 50.0]) q = keys[2] + 0.3 * rng.standard_normal(d) # 一个"有点像 key 2"的 query rows = [] for tau in (0.25, 1.0, math.sqrt(d), 4 * math.sqrt(d), 1e6): w = softmax(keys @ q / tau) label = "1e6（≈ 平均）" if tau == 1e6 else f"{tau:.2f}" + ("（= √d）" if tau == math.sqrt(d) else "") rows.append([label] + [f"{x:.3f}" for x in w] + [f"{w @ values:.2f}"]) print("5 个 key 对应的 value 是 10、20、30、40、50；query ≈ key 2（value = 30）+ 噪声") print_table(["温度 τ", "w0", "w1", "w2", "w3", "w4", "输出 Σ w·v"], rows) print_observations([ "τ 小时几乎所有权重都在 key 2 上，输出 ≈ 30：接近普通字典的硬查找。", "τ 很大时权重趋于均匀，输出 ≈ 所有 value 的平均 30（这里恰好也是 30，因为 value 是等差的）。", "和硬查找不同，软查找对 q、k、v **处处可导**：模型可以用梯度下降学会'查什么'（q）、'用什么当索引'（k）、" "'取出什么'（v）。Transformer 里这三者都是同一个输入乘三个不同的矩阵得到的。", ]) # --------------------------------------------------------------------------- # 演示 2：为什么除以 √d # --------------------------------------------------------------------------- def demo2_scaling(): """ q、k 的每个分量独立、均值 0、方差 1 时，点积 q·k = Σ q_i k_i 的方差是 d（论文脚注 4）。 d 很大时分数的差距很大，softmax 输出接近 one-hot，雅可比矩阵 J = diag(p) − p pᵀ 接近 0：梯度传不回去。 除以 √d 把方差拉回 1，和 d 无关。 """ print_header("演示 2：为什么除以 √d——点积方差、softmax 饱和与梯度") rng = np.random.default_rng(0) S, trials = 64, 200 rows = [] for d in (4, 16, 64, 256, 1024): stats = {"std": [], "pmax_raw": [], "pmax_s": [], "j_raw": [], "j_s": []} for _ in range(trials): q = rng.standard_normal(d) K = rng.standard_normal((S, d)) s = K @ q stats["std"].append(np.std(s)) for key_p, key_j, scores in (("pmax_raw", "j_raw", s), ("pmax_s", "j_s", s / math.sqrt(d))): p = softmax(scores) J = np.diag(p) - np.outer(p, p) stats[key_p].append(p.max()) stats[key_j].append(np.linalg.norm(J)) m = {k: float(np.mean(v)) for k, v in stats.items()} rows.append([d, f"{m['std']:.1f}", f"{math.sqrt(d):.1f}", f"{m['pmax_raw']:.3f}", f"{m['pmax_s']:.3f}", f"{m['j_raw']:.4f}", f"{m['j_s']:.4f}"]) print(f"每次：1 个 query、{S} 个 key，分量 ~ N(0, 1)，取 {trials} 次平均") print_table(["d", "点积的标准差", "√d", "最大权重(不缩放)", "最大权重(÷√d)", "‖J‖(不缩放)", "‖J‖(÷√d)"], rows) print_observations([ "点积的标准差 ≈ √d：d = 1024 时分数之间差几十，exp 之后最大的那个压倒其余所有。", "不缩放时，d 从 4 增大到 1024，最大权重从 0.27 涨到 0.96（几乎是 one-hot），softmax 的雅可比范数从 d = 16 时的 0.30 " "降到 0.06，而且 d 越大越小：梯度传不回 q 和 k，训练很难进行。这就是论文说的'softmax 进入梯度极小的区域'。", "除以 √d 之后两列都不随 d 变化。缩放因子只是一个常数乘法（可以并进 Wq 里），对推理的代价是零。", "【系统联系】低精度推理时，分数先乘 1/√d 再取 exp 还有一个好处：数值范围小，不容易在 FP16 / FP8 里溢出（前置讲 P3 §5.4）。", ]) # --------------------------------------------------------------------------- # 演示 3：因果掩码 # --------------------------------------------------------------------------- def demo3_causal_mask(): """ 语言模型逐个生成 token，训练时第 t 个位置只能看位置 ≤ t 的 token（不能偷看答案）。 做法：把 QKᵀ 里 j > t 的分数设成 −∞，softmax 之后权重就是 0。 """ print_header("演示 3：因果掩码——改掉后面的 token，前面的输出完全不变") cfg = TINY params = init_params(cfg, seed=0) rng = np.random.default_rng(3) toks = rng.integers(0, cfg.vocab, size=(1, 12)) toks2 = toks.copy() t = 6 toks2[0, t] = (toks[0, t] + 1) % cfg.vocab rows = [] for causal in (True, False): a = forward(params, cfg, toks, causal=causal) b = forward(params, cfg, toks2, causal=causal) diff = np.max(np.abs(a - b), axis=-1)[0] rows.append(["因果掩码" if causal else "双向（不掩码）", f"{diff[:t].max():.1e}", f"{diff[t:].min():.2f}"]) print(f"12 个 token 的序列，只改第 {t} 个（从 0 数），比较每个位置的 logits：") print_table(["注意力", f"位置 0–{t - 1} 的最大变化", f"位置 {t}–11 的最小变化"], rows) print_observations([ f"有因果掩码时，位置 0–{t - 1} 的输出**一个比特都没变**（差恰好是 0）：它们根本看不到位置 {t}。", "双向注意力（BERT、原论文的编码器）里，改一个 token 会影响所有位置。", "【系统联系】这条性质是 KV cache 的基础：位置 j 的 K、V 只依赖 token 0 … j，生成新 token 不会改变旧的 K、V，" "所以可以存起来复用（04 脚本）。双向模型做不到这一点。", ]) # --------------------------------------------------------------------------- # 演示 4：多头与 GQA 的运算量 # --------------------------------------------------------------------------- def demo4_multihead(): """ 单头：Q、K、V 都是 [T, d]，分数 [T, T]。多头：切成 h 份，每份 d/h 维，h 个 [T, T] 的分数矩阵。 投影矩阵的总大小相同（d×d），注意力核心的运算 h · 4·T·S·(d/h) = 4·T·S·d 也相同（论文 §3.2.2）。 GQA：K、V 只有 G 组头，投影矩阵变小，KV cache 变小；注意力核心的运算不变（query 头数没变）。 """ print_header("演示 4：多头注意力的运算量与单头相同；GQA 只减少 K/V") d, T = 512, 1024 rows = [] for name, h, g in (("单头", 1, 1), ("8 头（MHA）", 8, 8), ("64 头（MHA）", 64, 64), ("8 头 GQA（2 组 KV）", 8, 2), ("8 头 MQA（1 组 KV）", 8, 1)): dh = d // h proj = 2 * T * d * (h * dh + 2 * g * dh + h * dh) core = 4 * T * T * h * dh kv_per_tok = 2 * g * dh rows.append([name, h, dh, g, fmt_num(proj), fmt_num(core), kv_per_tok]) print(f"d = {d}，序列长度 T = {T}（不计因果掩码省下的一半）：") print_table(["配置", "头数 h", "每头 d/h", "KV 组数", "投影 FLOPs", "核心 FLOPs", "每 token 每层 KV 个数"], rows) # 用 transformer_lab 的实现检查：GQA 的 G = H 时就是普通多头 rng = np.random.default_rng(0) q = rng.standard_normal((1, 7, 8, 16)) k = rng.standard_normal((1, 7, 8, 16)) v = rng.standard_normal((1, 7, 8, 16)) out_mha = attention_core(q, k, v, 0) # 逐头单独算一遍（最朴素的写法）对照 ref = [] for hh in range(8): s = q[0, :, hh] @ k[0, :, hh].T / 4.0 s = np.where(np.tril(np.ones((7, 7), bool)), s, -np.inf) ref.append(softmax(s) @ v[0, :, hh]) ref = np.concatenate(ref, axis=-1) print(f"\n检查：attention_core（einsum 一次算所有头）与逐头循环的最大差 = {np.max(np.abs(out_mha[0] - ref)):.1e}") print_observations([ "把 d 切成多少个头，投影和注意力核心的 FLOPs 都不变：多头几乎是'免费'的，换来的是同时关注好几个地方的能力" "（每个头的 softmax 各自归一化，可以看不同的位置）。", "GQA / MQA 只缩小了 K、V 的投影和每 token 的 KV（8 组 → 2 组 → 1 组），注意力核心的运算不变——" "所以它们省的是**字节**而不是 FLOPs，正好打在 decode 的瓶颈上（05 脚本）。", "头越多，每个头越窄（64 头时每头只有 8 维），单个头的矩阵乘法越'瘦'，硬件效率会下降——真实模型的 d_head 通常是 64–128。", ]) # --------------------------------------------------------------------------- # 演示 5：位置信息 # --------------------------------------------------------------------------- def sinusoidal_pe(positions, d): """原论文的正弦位置编码：PE(pos, 2i) = sin(pos / 10000^(2i/d))，PE(pos, 2i+1) = cos(…)。""" pos = np.asarray(positions, dtype=np.float64)[:, None] freq = 10000.0 ** (-np.arange(0, d, 2) / d) pe = np.zeros((len(positions), d)) pe[:, 0::2] = np.sin(pos * freq) pe[:, 1::2] = np.cos(pos * freq) return pe def demo5_positions(): print_header("演示 5：位置信息——置换等变、正弦编码、RoPE 与一个'看前一个 token'的头") cfg = TINY params = init_params(cfg, seed=0) rng = np.random.default_rng(5) toks = rng.integers(0, cfg.vocab, size=(1, 10)) perm = rng.permutation(10) # (a) 没有位置编码、没有掩码：打乱输入，输出只是跟着打乱 rows = [] for use_rope, causal, label in ((False, False, "无位置编码、双向"), (False, True, "无位置编码、因果掩码"), (True, False, "RoPE、双向")): a = forward(params, cfg, toks, use_rope=use_rope, causal=causal) b = forward(params, cfg, toks[:, perm], use_rope=use_rope, causal=causal) rows.append([label, f"{np.max(np.abs(a[:, perm] - b)):.1e}"]) print("(a) 打乱 10 个 token 的顺序，比较'先算再打乱'和'先打乱再算'：") print_table(["设置", "最大差"], rows) # (b) 正弦编码：PE(pos + k) 是 PE(pos) 的线性函数（每一对数旋转 k·ω_i） d, k = 16, 7 pe = sinusoidal_pe(np.arange(0, 50), d) freq = 10000.0 ** (-np.arange(0, d, 2) / d) M = np.zeros((d, d)) for i, w in enumerate(freq): c, s = math.cos(k * w), math.sin(k * w) # (sin(a+b), cos(a+b)) = [[cos b, sin b], [−sin b, cos b]] · (sin a, cos a) M[2 * i:2 * i + 2, 2 * i:2 * i + 2] = [[c, s], [-s, c]] err = np.max(np.abs(pe[k:] - pe[:-k] @ M.T)) print(f"\n(b) 正弦位置编码：PE(pos + {k}) = M_{k} · PE(pos)，M_{k} 是一个和 pos 无关的旋转矩阵，最大误差 {err:.1e}") # (c) RoPE：分数只依赖 m − n dh = 16 q = rng.standard_normal((1, 1, 1, dh)) kk = rng.standard_normal((1, 1, 1, dh)) rows = [] for m, n in ((5, 2), (105, 102), (10005, 10002), (5, 4)): cq, sq = rope_angles([m], dh) ck, sk = rope_angles([n], dh) score = float(np.sum(apply_rope(q, cq, sq) * apply_rope(kk, ck, sk))) rows.append([m, n, m - n, f"{score:.10f}"]) print("\n(c) RoPE：同一对 q、k 放在不同的绝对位置上，点积只取决于相对位置 m − n：") print_table(["query 位置 m", "key 位置 n", "m − n", "q·k（加 RoPE 之后）"], rows) # (d) 用 RoPE 手工搭一个"看前一个 token"的头 # 让 k 的每一对数 = q 的对应一对数预先旋转 +ω_i：分数 ∝ Σ_i |q_i|² cos((m − n − 1)·ω_i)，在 m − n = 1 时最大 base = rng.standard_normal(dh // 2) ** 2 + 1.0 amp = 3.0 qv = np.zeros(dh) qv[0::2] = amp * np.sqrt(base) c1, s1 = rope_angles([1], dh) kv = apply_rope(qv[None, None, None, :], c1, s1)[0, 0, 0] T = 12 cos, sin = rope_angles(np.arange(T), dh) Q = apply_rope(np.tile(qv, (1, T, 1, 1)), cos, sin) K = apply_rope(np.tile(kv, (1, T, 1, 1)), cos, sin) V = np.zeros((1, T, 1, dh)) V[0, :, 0, :T] = np.eye(T) # value = 位置的独热向量（补零到 dh 维），输出就是注意力权重本身 W = attention_core(Q, K, V, 0)[0, :, :T] # [T, T]：第 t 行 = 位置 t 的 query 对各位置的权重 rows = [] for t in (3, 7, 11): top = np.argsort(-W[t])[:2] rows.append([t, int(top[0]), f"{W[t, top[0]]:.2f}", int(top[1]), f"{W[t, top[1]]:.2f}"]) print("\n(d) 手工搭的'前一个 token 头'（所有 token 的内容都一样，只靠 RoPE 区分位置）：") print_table(["query 位置", "权重最大的位置", "权重", "第二大的位置", "权重"], rows) print_observations([ "(a) 没有位置编码、没有掩码时，打乱输入只会让输出跟着打乱（差为 0）：注意力本身不知道顺序，" "'猫追狗'和'狗追猫'得到同一组向量。所以必须注入位置信息。（有趣的是因果掩码本身也泄露了位置：只能看前面，" "第二行的差不为 0；有些模型干脆不用位置编码，就是靠这一点。）", "(b) 正弦编码的妙处：'往后挪 k 个位置'是一个固定的线性变换，模型容易学会'看相对位置'（论文 §3.5 的动机）。", "(c) RoPE（Su et al. 2021）更进一步：不是把位置加到输入上，而是把 q、k 按位置旋转。点积只依赖 m − n——" "位置 10005 和位置 5 看到的'相对 3 个位置'完全一样。Llama、Qwen、DeepSeek 等都用它。", "(d) 只用 RoPE 就能做出一个'总是看前一个 token'的头，而且在任何绝对位置都一样。真实模型里确实能找到这样的头" "（previous-token head），它是复制、归纳等能力的基础部件。", "【系统联系】RoPE 加在 K 上之后才写进 KV cache；用 MLA 压缩 KV 时，带 RoPE 的那部分不能被压缩吸收，" "DeepSeek 只好单独存一份 64 维的'RoPE key'（05 脚本）。", ]) # --------------------------------------------------------------------------- # 演示 6：RNN vs 自注意力 # --------------------------------------------------------------------------- def demo6_rnn_vs_attention(): """ RNN：h_t = tanh(h_{t−1}·W + x_t·U)。第 t 步必须等第 t−1 步算完：n 步顺序执行，每步是一个 [1, d]×[d, d] 的小乘法。 自注意力：Q·Kᵀ、P·V 是两个大矩阵乘法，所有位置同时算。 论文表 1：每层运算量 RNN O(n·d²)、自注意力 O(n²·d)；顺序步数 O(n) vs O(1)；最长路径 O(n) vs O(1)。 """ print_header("演示 6：RNN vs 自注意力——运算量、顺序步数与实测") d = 256 rows = [] for n in (64, 256, 1024, 4096): rnn = 2 * 2 * n * d * d # 每步两个 [1,d]×[d,d] attn_core = 4 * n * n * d # QKᵀ 和 PV attn_proj = 2 * n * d * 4 * d # Q/K/V/O 四个投影 rows.append([n, fmt_num(rnn), fmt_num(attn_core), fmt_num(attn_proj), n, 1]) print(f"d = {d}：") print_table(["n", "RNN FLOPs", "注意力核心 FLOPs", "注意力投影 FLOPs", "RNN 顺序步数", "注意力顺序步数"], rows) rng = np.random.default_rng(6) W = (rng.standard_normal((d, d)) / math.sqrt(d)).astype(np.float32) rows = [] for n in (256, 1024, 4096): X = rng.standard_normal((n, d)).astype(np.float32) def rnn_sequential(): # 循环部分：h_t 依赖 h_{t−1}，只能一步一步算（输入部分 x_t·U 不依赖 h，可以事先一次算完，这里略去） h = np.zeros(d, dtype=np.float32) for t in range(n): h = np.tanh(h @ W + X[t]) return h def same_flops_one_matmul(): # 同样的运算量，如果没有顺序依赖：把 n 个向量摞成一个矩阵，一次乘完 return np.tanh(X @ W + X) def attn(): s = X @ X.T / math.sqrt(d) s = np.where(np.tril(np.ones((n, n), bool)), s, -np.inf) return softmax(s) @ X ts = [] for fn in (rnn_sequential, same_flops_one_matmul, attn): best = float("inf") for _ in range(3): t0 = time.perf_counter() fn() best = min(best, time.perf_counter() - t0) ts.append(best) f_rec = 2 * n * d * d rows.append([n, fmt_time(ts[0]), fmt_num(f_rec / ts[0]) + "FLOP/s", fmt_time(ts[1]), fmt_num(f_rec / ts[1]) + "FLOP/s", f"{ts[0] / ts[1]:.0f}×", fmt_time(ts[2])]) print("\nnumpy 实测（float32，d = 256，取 3 次最快；墙钟时间，只看趋势）：") print_table(["n", "RNN 循环（逐步）", "速度", "同样运算一次乘完", "速度", "逐步慢几倍", "因果注意力核心"], rows) print_observations([ "按 FLOPs 算，n < d 时自注意力核心更省，n > d 时 RNN 更省（论文表 1 的 n²·d vs n·d²）；再加上 Q/K/V/O 四个投影，" "自注意力层的总运算量其实比 RNN 层多。论文当时的句子长度远小于 d，所以核心部分不是负担。", "但同样的运算，逐步算比'摞成一个矩阵一次乘完'慢约 10 倍（编写时的机器上 9–11 倍）：每一步只是向量乘矩阵，" "每个权重只用一次，还要付每步的调用开销。在有几千个计算单元的 GPU 上差距更大——一个向量乘矩阵根本喂不饱它。", "这正是 Transformer 胜出的系统原因：训练时所有位置一起算，全是大矩阵乘法，**能不能并行比 FLOPs 多少更重要**。" "讽刺的是，**推理的 decode 阶段又退回了'一步一个 token'**：每一步只有 1 个新 token，又变成了向量乘矩阵——" "这就是第 1 讲说 decode 是 memory-bound 的根源，也是第 4–7 讲各种批处理、投机解码技巧要解决的问题。", ]) if __name__ == "__main__": demo1_soft_lookup() demo2_scaling() demo3_causal_mask() demo4_multihead() demo5_positions() demo6_rnn_vs_attention()