""" ============================================================================= 第 7 讲 · 代码 1：用 numpy 一步一步算自注意力 ============================================================================= 对应讲义：lectures/lecture07_self_attention_transformer.md 第 2、3、4、10 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 第 6 讲的注意力是"解码器看编码器"：查询来自一个序列，键和值来自另一个序列。 **自注意力**（self-attention）是"一个序列看它自己"：句子里的每个词都当一次查询，去看句子里所有的词。 输入 X：(n, d)，n 个词，每个词一个 d 维向量 Q = X W_Q, K = X W_K, V = X W_V 每个词各自投影出查询、键、值 输出 = softmax(Q Kᵀ / √d_k) V 每个词的输出 = 所有词的值的加权平均 这个文件不用 PyTorch，只用 numpy，把每一步的中间结果都打印出来： 演示 1：注意力 = "软的"字典查找 演示 2：一步一步算自注意力（循环写法 vs 矩阵写法） 演示 3：为什么要除以 √d_k？—— 不缩放时 softmax 会"饱和"，梯度消失 演示 4：自注意力不知道词序（置换等变）—— 所以需要位置编码 演示 5：因果掩码 —— 让每个词只能看到前面的词 演示 6：计算量 —— 注意力矩阵的 O(n²) 什么时候成为瓶颈？ ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 01_self_attention_numpy.py 只需要 numpy，约 1 秒。 """ import numpy as np def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def softmax(z, axis=-1): """ 数值稳定的 softmax：先减去最大值，再取指数、归一化。 【参数】 z : 任意形状的数组 axis : 沿着哪一维归一化（默认最后一维） 【返回】和 z 形状相同的数组，沿 axis 的元素之和为 1 【边界】 - 减去最大值不改变结果（分子分母同乘 e^(-max)），但能防止 e^z 溢出成 inf。 - z 中的 -inf 会变成 0 权重（e^(-inf) = 0）—— 掩码就是利用这一点。 如果某一行全是 -inf，结果是 nan（0/0），使用掩码时要保证每行至少有一个有限值。 【Python 小知识】 keepdims=True 让求最大值、求和之后保留那一维（长度变成 1），这样才能和原数组广播相减、相除。 """ z = z - np.max(z, axis=axis, keepdims=True) e = np.exp(z) return e / np.sum(e, axis=axis, keepdims=True) def self_attention_loop(X, W_Q, W_K, W_V, mask=None): """ 自注意力的"循环写法"：对每个词 i 单独算一遍。最容易看懂，但慢。 【参数】 X : (n, d) 输入，n 个词 W_Q, W_K, W_V : (d, d_k)、(d, d_k)、(d, d_v) 投影矩阵 mask : None，或 (n, n) 的布尔数组，mask[i, j] = True 表示"词 i 可以看词 j" 【返回】(输出 (n, d_v), 注意力权重 (n, n)) 【思路】对每个词 i： 1. 它的查询 q_i = x_i W_Q 2. 和每个词 j 的键 k_j = x_j W_K 做点积，再除以 √d_k，得到打分 e_ij 3. 对 j 做 softmax，得到权重 α_ij 4. 输出 o_i = Σ_j α_ij v_j，其中 v_j = x_j W_V """ n, d_k = X.shape[0], W_Q.shape[1] out = np.zeros((n, W_V.shape[1])) weights = np.zeros((n, n)) for i in range(n): q_i = X[i] @ W_Q scores = np.array([q_i @ (X[j] @ W_K) / np.sqrt(d_k) for j in range(n)]) if mask is not None: scores = np.where(mask[i], scores, -np.inf) alpha = softmax(scores) weights[i] = alpha out[i] = sum(alpha[j] * (X[j] @ W_V) for j in range(n)) return out, weights def self_attention(X, W_Q, W_K, W_V, mask=None, scale=True): """ 自注意力的"矩阵写法"：所有词一起算。结果和循环写法完全一样，但快得多（也是 GPU 上的写法）。 【参数】 X, W_Q, W_K, W_V, mask : 同 self_attention_loop scale : 是否除以 √d_k（演示 3 用 False 看不缩放的后果） 【返回】(输出 (n, d_v), 注意力权重 (n, n)) 【思路】 Q = X W_Q，K = X W_K，V = X W_V (n, d_k)、(n, d_k)、(n, d_v) S = Q Kᵀ / √d_k (n, n)：第 i 行第 j 列 = 词 i 对词 j 的打分 A = softmax(S)（按行） (n, n)：每一行是一个概率分布 输出 = A V (n, d_v) 【边界】mask 中某一行全为 False 时，这一行的 softmax 是 nan。 """ Q, K, V = X @ W_Q, X @ W_K, X @ W_V S = Q @ K.T if scale: S = S / np.sqrt(W_Q.shape[1]) if mask is not None: S = np.where(mask, S, -np.inf) A = softmax(S, axis=1) return A @ V, A def print_matrix(M, row_labels, col_labels, fmt="{:6.2f}"): """把矩阵连同行、列标签一起打印出来。""" width = max(len(fmt.format(0.0)), max(len(c) for c in col_labels)) + 1 print(" " + "".join(f"{c:>{width}s}" for c in col_labels)) for r, row in zip(row_labels, M): print(f" {r:<6s}" + "".join(f"{fmt.format(x):>{width}s}" for x in row)) def softmax_jacobian_norm(p): """ softmax 的雅可比矩阵 ∂p/∂z = diag(p) - p pᵀ 的 Frobenius 范数（所有元素平方和再开方）。 【思路】这个范数衡量"打分变一点，权重能变多少"，也就是梯度能不能从权重传回打分。 p 是均匀分布时它比较大；p 接近 one-hot（一个位置几乎拿走全部权重）时它接近 0 —— 梯度消失。 """ J = np.diag(p) - np.outer(p, p) return np.sqrt(np.sum(J ** 2)) def main(): rng = np.random.default_rng(0) # ---------------- 演示 1：软的字典查找 ---------------- print_header("演示 1：注意力 = \"软的\"字典查找") print(""" Python 的字典是"硬"查找：查询必须和某个键**完全相等**，才返回那个键对应的值。 注意力是"软"查找：查询和**每个**键都算一个相似度，按相似度给所有的值加权平均。""") keys = np.array([[1.0, 0.0], [0.0, 1.0], [0.7, 0.7]]) values = np.array([10.0, 20.0, 30.0]) for q in [np.array([10.0, 0.0]), np.array([1.0, 0.0]), np.array([0.5, 0.5])]: w = softmax(keys @ q) print(f"\n 查询 {q}：权重 {np.round(w, 3)}，输出 = {w @ values:.2f}") print(""" 观察：查询 (10, 0) 和第 1 个键方向相同、而且"很长"，约 95% 的权重给了第 1 个值（接近硬查找，输出接近 10）； 查询 (1, 0) 同样的方向但短得多，权重分散开了；查询 (0.5, 0.5) 和第 3 个键最像，但另外两个也分到了不少。 打分的**大小**决定了 softmax 有多"尖" —— 演示 3 会看到这一点为什么重要。""") # ---------------- 演示 2：一步一步算 ---------------- print_header("演示 2：一步一步算自注意力（4 个词，d = 4，d_k = d_v = 3）") words = ["the", "cat", "sat", "down"] n, d, d_k = 4, 4, 3 X = rng.normal(size=(n, d)).round(1) W_Q, W_K, W_V = (rng.normal(size=(d, d_k)).round(1) for _ in range(3)) print("\n 输入 X（每行一个词的向量）：") print_matrix(X, words, [f"x{j}" for j in range(d)]) Q, K, V = X @ W_Q, X @ W_K, X @ W_V print("\n 第 1 步：Q = X W_Q（每个词的查询）") print_matrix(Q, words, [f"q{j}" for j in range(d_k)]) S = Q @ K.T / np.sqrt(d_k) print(f"\n 第 2 步：打分 S = Q Kᵀ / √{d_k}（第 i 行第 j 列 = 词 i 对词 j 的打分）") print_matrix(S, words, words) A = softmax(S, axis=1) print("\n 第 3 步：按行 softmax，得到注意力权重 A（每行之和为 1）") print_matrix(A, words, words) out = A @ V print("\n 第 4 步：输出 = A V（每个词的输出是所有词的值向量的加权平均）") print_matrix(out, words, [f"o{j}" for j in range(d_k)]) out_loop, A_loop = self_attention_loop(X, W_Q, W_K, W_V) out_mat, A_mat = self_attention(X, W_Q, W_K, W_V) assert np.allclose(out_loop, out_mat) and np.allclose(A_loop, A_mat) assert np.allclose(A.sum(axis=1), 1.0) print("\n 自我检查：循环写法和矩阵写法的结果完全一致 ✓；每行权重之和为 1 ✓") print(""" 注意：这里的参数是随机的，所以权重没有什么"语言学含义"。训练之后，模型会学到有意义的模式， 比如动词关注它的主语、代词关注它指代的名词（04 号文件演示 4 会看到一个例子）。""") # ---------------- 演示 3：为什么要缩放 ---------------- print_header("演示 3：为什么要除以 √d_k？") print(""" 如果 q 和 k 的每一维都是均值 0、方差 1 的独立随机数，点积 q·k = Σ q_j k_j 是 d_k 个随机数之和， 方差是 d_k，标准差是 √d_k。d_k 越大，打分的"量级"越大，softmax 越"尖"（越接近 one-hot）。 下面对每个 d_k 做 2000 次实验：1 个查询、10 个键，都是标准正态随机向量。""") print(f"\n {'d_k':>6s} 点积的标准差 │ 不缩放：最大权重 梯度范数 │ 缩放：最大权重 梯度范数") for dk in [4, 16, 64, 256, 1024]: dots, mx_raw, g_raw, mx_s, g_s = [], [], [], [], [] for _ in range(2000): q = rng.normal(size=dk) Kmat = rng.normal(size=(10, dk)) z = Kmat @ q dots.append(z[0]) p_raw, p_s = softmax(z), softmax(z / np.sqrt(dk)) mx_raw.append(p_raw.max()) g_raw.append(softmax_jacobian_norm(p_raw)) mx_s.append(p_s.max()) g_s.append(softmax_jacobian_norm(p_s)) print(f" {dk:>6d} {np.std(dots):>12.1f} │ {np.mean(mx_raw):>16.3f} {np.mean(g_raw):>9.4f} │ " f"{np.mean(mx_s):>14.3f} {np.mean(g_s):>9.4f}") print(""" 观察： - 点积的标准差 ≈ √d_k（4 → 2，1024 → 32），和上面的推导一致。 - 不缩放时，d_k 越大，最大的那个权重越接近 1（softmax 几乎变成 argmax）， softmax 的雅可比矩阵的范数越接近 0 —— 梯度很难从注意力权重传回 Q 和 K，训练变慢甚至停滞。 - 除以 √d_k 之后，打分的方差回到 1，无论 d_k 多大，权重的"尖锐程度"和梯度的大小都基本不变。""") # ---------------- 演示 4：置换等变 ---------------- print_header("演示 4：自注意力不知道词序 —— 置换等变") words = ["dog", "bites", "man"] emb = {w: rng.normal(size=4) for w in words} W_Q, W_K, W_V = (rng.normal(size=(4, 4)) for _ in range(3)) X1 = np.stack([emb[w] for w in ["dog", "bites", "man"]]) X2 = np.stack([emb[w] for w in ["man", "bites", "dog"]]) o1, _ = self_attention(X1, W_Q, W_K, W_V) o2, _ = self_attention(X2, W_Q, W_K, W_V) print("\n 句子 1：dog bites man 句子 2：man bites dog") for w in words: i1 = ["dog", "bites", "man"].index(w) i2 = ["man", "bites", "dog"].index(w) print(f" \"{w}\" 的输出在两个句子里相同吗？{np.allclose(o1[i1], o2[i2])}") print(""" 观察：把输入的词打乱顺序，输出只是跟着打乱，每个词得到的向量**完全一样**。 对自注意力来说，句子只是一个"词的集合"：它分不清"狗咬人"和"人咬狗"。 原因：打分只取决于两个词的向量 q_i·k_j，和它们在句子里的位置无关。 解决办法：在输入中加入**位置信息**（02 号文件、讲义第 5 节）。""") P = np.stack([np.sin(np.arange(4) * (p + 1)) for p in range(3)]) # 随便一种位置向量，只要每个位置不同 o1p, _ = self_attention(X1 + P, W_Q, W_K, W_V) o2p, _ = self_attention(X2 + P, W_Q, W_K, W_V) print(f" 给每个位置加上不同的位置向量之后，\"dog\" 的输出还相同吗？{np.allclose(o1p[0], o2p[2])}") # ---------------- 演示 5：因果掩码 ---------------- print_header("演示 5：因果掩码 —— 只看前面的词") words = ["", "the", "cat", "sat"] X = rng.normal(size=(4, 4)) W_Q, W_K, W_V = (rng.normal(size=(4, 3)) for _ in range(3)) causal = np.tril(np.ones((4, 4), dtype=bool)) # 下三角（含对角线）为 True out, A = self_attention(X, W_Q, W_K, W_V, mask=causal) print("\n 掩码（1 = 可以看，0 = 不能看）：") print_matrix(causal.astype(float), words, words, fmt="{:4.0f}") print("\n 加了掩码之后的注意力权重（右上角全是 0）：") print_matrix(A, words, words) X_changed = X.copy() X_changed[3] = rng.normal(size=4) # 改变最后一个词 out2, _ = self_attention(X_changed, W_Q, W_K, W_V, mask=causal) print(f"\n 把最后一个词换掉，前 3 个位置的输出变了吗？{not np.allclose(out[:3], out2[:3])}") print(""" 观察：第 i 个词只能关注第 1~i 个词。语言模型要用第 i 个位置的输出预测第 i+1 个词， 如果它能看到第 i+1 个词，就等于"偷看答案"。有了掩码，**整个句子可以一次性并行地算完**， 每个位置的输出都和"只给模型看前 i 个词"时完全相同 —— 这是 Transformer 训练快的关键。""") # ---------------- 演示 6：计算量 ---------------- print_header("演示 6：计算量 —— 注意力矩阵的 O(n²) 什么时候成为瓶颈？") print(""" 一个 Transformer 层（d = d_model，FFN 隐藏层 4d），处理长度为 n 的序列，乘加运算次数（1 次乘加 = 2 次浮点运算）： Q、K、V、输出投影：4 个 d×d 矩阵乘 → 4·n·d² 前馈网络：d→4d→d → 8·n·d² 打分 Q Kᵀ 和加权求和 A V → 2·n²·d 两者之比 = 2n²d / 12nd² = n / (6d)：序列长度超过 6d 时，注意力矩阵的部分开始占大头。""") d = 768 print(f"\n d = {d}（GPT-2 small）") print(f" {'n':>7s} {'线性层（十亿次乘加）':>18s} {'注意力矩阵':>12s} {'注意力矩阵占比':>12s}") for n in [128, 512, 1024, 4096, 16384, 65536]: lin, att = 12 * n * d * d / 1e9, 2 * n * n * d / 1e9 print(f" {n:>7d} {lin:>18.2f} {att:>14.2f} {att / (lin + att):>14.1%}") print(""" 观察：n = 1024（GPT-2 的上下文长度）时注意力矩阵只占约 18%；n = 65536 时占到 90% 以上。 而且 n×n 的注意力矩阵本身要存下来（每个头、每一层一个），显存也按 n² 增长。 这就是"长上下文"难的原因之一（FlashAttention、稀疏注意力、线性注意力等方法见 CS336 第 4、6 讲）。""") if __name__ == "__main__": main()