""" ============================================================================= 第 7 讲 · 代码 2：位置编码 —— 正弦编码、RoPE、可学习的位置嵌入 ============================================================================= 对应讲义：lectures/lecture07_self_attention_transformer.md 第 5 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 01 号文件演示 4 看到：自注意力本身不知道词序，"dog bites man" 和 "man bites dog" 对它来说是一样的。 要让模型知道每个词在第几个位置，必须把位置信息"告诉"它。三种主流做法： 1. 正弦位置编码（Vaswani 等 2017，原始 Transformer）：用一组不同频率的 sin/cos 表示位置，加到词向量上 2. 可学习的位置嵌入（BERT、GPT-2）：每个位置一个向量，和词向量一样通过训练学出来 3. 旋转位置编码 RoPE（Su 等 2021，LLaMA 等今天的大部分模型）：不加到输入上，而是在每一层把查询和键"旋转" 演示 1：正弦位置编码长什么样（打印数值 + 字符热图） 演示 2：两个位置的编码做点积，只取决于它们的距离 演示 3："往后挪 k 个位置"是一个固定的线性变换（旋转） 演示 4：RoPE —— 旋转之后，查询和键的点积只取决于相对位置 演示 5：可学习的位置嵌入 —— 简单，但超出训练长度就没有了 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 02_positional_encoding.py 只需要 numpy，约 1 秒。装了 matplotlib 时会另存一张图 figures/sinusoidal_pe.png； 装了 PyTorch 时会检查这里的 numpy 实现和 transformer.py 里的实现是否一致。 """ from pathlib import Path import numpy as np def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def sinusoidal(max_len, d, base=10000.0): """ 正弦位置编码（numpy 版，和 transformer.py 中的 sinusoidal_encoding 完全相同）。 PE[p, 2i] = sin(p · ω_i) PE[p, 2i+1] = cos(p · ω_i) 其中 ω_i = 1 / base^(2i/d) 【参数】 max_len : 位置个数（0 ~ max_len-1） d : 维度（必须是偶数） base : 频率的底数 【返回】(max_len, d) 的数组 【边界】 - 第 0 对维度的频率 ω_0 = 1（周期 2π ≈ 6.3 个位置，转得最快）； 最后一对的频率接近 1/base（周期接近 2π·base ≈ 62832 个位置，几乎不动）。 - d 是奇数时报错。 """ if d % 2: raise ValueError("d 必须是偶数") omega = base ** (-np.arange(0, d, 2) / d) # (d/2,) angles = np.arange(max_len)[:, None] * omega[None, :] # (max_len, d/2) pe = np.zeros((max_len, d)) pe[:, 0::2] = np.sin(angles) pe[:, 1::2] = np.cos(angles) return pe def shift_matrix(k, d, base=10000.0): """ 构造矩阵 M_k，使得对任意位置 p：PE[p + k] = M_k · PE[p]。 【思路】（讲义第 5.2 节） 每一对 (sin(pω), cos(pω)) 往后挪 k 个位置变成 (sin((p+k)ω), cos((p+k)ω))。由和角公式： sin((p+k)ω) = cos(kω)·sin(pω) + sin(kω)·cos(pω) cos((p+k)ω) = -sin(kω)·sin(pω) + cos(kω)·cos(pω) 这是一个 2×2 的旋转矩阵，只和 k 有关、和 p 无关。M_k 就是 d/2 个这样的小矩阵排在对角线上。 【返回】(d, d) 的数组 """ omega = base ** (-np.arange(0, d, 2) / d) M = np.zeros((d, d)) for i, w in enumerate(omega): c, s = np.cos(k * w), np.sin(k * w) M[2 * i:2 * i + 2, 2 * i:2 * i + 2] = [[c, s], [-s, c]] return M def rope(x, positions, base=10000.0): """ RoPE（numpy 版）：把向量 x 在位置 p 处旋转。和 transformer.py 的 apply_rope 用同样的配对方式： 前一半维度的第 j 维和后一半维度的第 j 维组成一对，旋转角度 p·θ_j，θ_j = base^(-2j/d)。 【参数】 x : (n, d) 的数组，n 个向量（d 必须是偶数） positions : 长度为 n 的整数序列，每个向量所在的位置 【返回】(n, d) 的数组 【边界】旋转不改变向量的长度；位置 0 处不旋转（角度为 0）。 """ n, d = x.shape half = d // 2 theta = base ** (-np.arange(0, d, 2) / d) # (d/2,) ang = np.asarray(positions)[:, None] * theta[None, :] # (n, d/2) c, s = np.cos(ang), np.sin(ang) x1, x2 = x[:, :half], x[:, half:] return np.concatenate([x1 * c - x2 * s, x1 * s + x2 * c], axis=1) def ascii_heatmap(M, chars=" .:-=+*#%@"): """ 把取值在 [-1, 1] 的矩阵画成字符图：越接近 -1 用越"淡"的字符，越接近 1 用越"浓"的字符。 【返回】list[str]，每个元素是一行 """ idx = np.clip(((M + 1) / 2 * (len(chars) - 1)).round().astype(int), 0, len(chars) - 1) return ["".join(chars[i] for i in row) for row in idx] def main(): rng = np.random.default_rng(0) # ---------------- 演示 1：长什么样 ---------------- print_header("演示 1：正弦位置编码长什么样？（d = 16）") pe = sinusoidal(64, 16) print("\n 位置 " + " ".join(f"{'维' + str(j):>6s}" for j in [0, 1, 2, 3, 14, 15])) for p in [0, 1, 2, 3, 10, 50]: print(f" {p:>4d} " + " ".join(f"{pe[p, j]:>7.3f}" for j in [0, 1, 2, 3, 14, 15])) print("\n 字符热图（每行一个位置 0~39，每列一个维度 0~15；'@' ≈ +1，' ' ≈ -1）：") for p, row in enumerate(ascii_heatmap(pe[:40])): if p % 4 == 0: print(f" {p:>3d} |{row}|") print(""" 观察： - 前面的维度（左边）变化很快：第 0、1 维每过约 6 个位置就转完一圈； 后面的维度（右边）变化很慢：第 14、15 维在 0~50 个位置里几乎不变（sin ≈ 0，cos ≈ 1）。 - 就像钟表的秒针、分针、时针：每根针单独看会"重复"，但所有针合在一起，每个位置的组合都不同。 - 所有值都在 [-1, 1] 之间，和词向量的量级相当，加上去不会"淹没"词的信息。""") try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt fig_dir = Path(__file__).parent / "figures" fig_dir.mkdir(exist_ok=True) plt.figure(figsize=(8, 4)) plt.imshow(sinusoidal(100, 64).T, aspect="auto", cmap="RdBu", vmin=-1, vmax=1) plt.xlabel("position") plt.ylabel("dimension") plt.colorbar() plt.title("Sinusoidal positional encoding (d = 64)") plt.tight_layout() plt.savefig(fig_dir / "sinusoidal_pe.png", dpi=120) plt.close() print(" （已另存 figures/sinusoidal_pe.png：100 个位置 × 64 维的热图）") except ImportError: print(" （没有安装 matplotlib，跳过画图）") # ---------------- 演示 2：点积只取决于距离 ---------------- print_header("演示 2：两个位置编码的点积，只取决于它们的距离 k") pe = sinusoidal(200, 64) print(f"\n {'距离 k':>6s}" + "".join(f"{'p=' + str(p):>10s}" for p in [0, 20, 100])) for k in [0, 1, 2, 4, 8, 16, 32, 64]: print(f" {k:>6d}" + "".join(f"{pe[p] @ pe[p + k]:>10.3f}" for p in [0, 20, 100])) print(""" 观察： - 同一行的三个数完全相同：PE[p]·PE[p+k] = Σ_i cos(k·ω_i)，只和距离 k 有关，和起点 p 无关。 - 距离越远，点积大体上越小（k = 0 时是 d/2 = 32）。模型可以通过这样的点积"感受到"两个词离得多远。""") # ---------------- 演示 3：平移是线性变换 ---------------- print_header("演示 3：\"往后挪 k 个位置\"是一个固定的线性变换") d = 16 pe = sinusoidal(100, d) for k in [1, 5]: M = shift_matrix(k, d) ok = all(np.allclose(M @ pe[p], pe[p + k]) for p in range(0, 90)) print(f"\n k = {k}：对所有 p = 0..89，M_k · PE[p] == PE[p + k] 都成立？{ok}") print(""" 观察：存在一个只依赖 k 的矩阵 M_k（d/2 个 2×2 旋转矩阵），把任何位置的编码变成往后 k 个位置的编码。 原论文的设想：这样模型的线性层就有可能学会"关注往前 k 个位置的词"这样的相对位置模式。 （注意：正弦编码是**加**在词向量上的，经过 W_Q、W_K 之后，这个性质并不能被完美地利用 —— 这正是 RoPE 的出发点。）""") # ---------------- 演示 4：RoPE ---------------- print_header("演示 4：RoPE —— 旋转之后，点积只取决于相对位置") dh = 8 q, k = rng.normal(size=dh), rng.normal(size=dh) print(f"\n 同一个查询 q 和同一个键 k，放在不同的位置 (m, n)：") print(f" {'(m, n)':>10s} {'m - n':>6s} {'RoPE 点积':>10s} {'加正弦编码后的点积':>18s}") pe = sinusoidal(200, dh) for m, n in [(5, 2), (10, 7), (100, 97), (5, 5), (50, 50), (2, 5), (60, 63)]: qr = rope(q[None], [m])[0] kr = rope(k[None], [n])[0] additive = (q + pe[m]) @ (k + pe[n]) print(f" {str((m, n)):>10s} {m - n:>6d} {qr @ kr:>10.4f} {additive:>18.4f}") lens = np.linalg.norm(rope(np.tile(q, (5, 1)), [0, 1, 10, 100, 1000]), axis=1) print(f"\n q 在位置 0、1、10、100、1000 旋转后的长度：{np.round(lens, 4)}（原长度 {np.linalg.norm(q):.4f}）") print(""" 观察： - RoPE：m - n 相同的行，点积完全相同（3、0、-3 三组）—— 注意力打分只取决于**相对位置**。 原因：q 旋转 mθ、k 旋转 nθ，两个旋转后的向量的点积 = q 和 "k 旋转 (n-m)θ" 的点积。 - 把正弦编码**加**到向量上再做点积：结果和 m、n 各自的值都有关，不只取决于 m - n。 - 旋转不改变长度：位置信息只改变方向，不会"放大"或"缩小"向量。 - RoPE 在**每一层**的注意力里都作用于 Q 和 K（不作用于 V），不需要额外的参数。""") try: import torch from transformer import apply_rope, rope_cache, sinusoidal_encoding x = rng.normal(size=(6, dh)) cos, sin = rope_cache(16, dh) xt = torch.tensor(x, dtype=torch.float32)[None, None] # (1, 1, 6, dh) ours = apply_rope(xt, cos, sin)[0, 0].numpy() same_rope = np.allclose(ours, rope(x, range(6)), atol=1e-5) same_sin = np.allclose(sinusoidal_encoding(50, 16).numpy(), sinusoidal(50, 16), atol=1e-5) print(f" 自我检查：numpy 实现与 transformer.py 一致？RoPE {same_rope}，正弦编码 {same_sin}") except ImportError: print(" （没有安装 PyTorch，跳过和 transformer.py 的一致性检查）") # ---------------- 演示 5：可学习的位置嵌入 ---------------- print_header("演示 5：可学习的位置嵌入 —— 简单，但超出训练长度就没有了") max_len = 8 pos_table = rng.normal(scale=0.02, size=(max_len, 4)) # 训练前：随机初始化的 8 个位置向量 print(f"\n 位置嵌入表：{max_len} 行（位置 0~{max_len - 1}），每行一个可学习的向量。") for length in [5, 8, 12]: try: _ = pos_table[np.arange(length)] print(f" 长度 {length:>2d} 的句子：可以") except IndexError: print(f" 长度 {length:>2d} 的句子：出错 —— 表里没有位置 {max_len}~{length - 1} 的向量") print(""" 观察：可学习的位置嵌入最灵活（每个位置的向量随便学），BERT、GPT-2 都用它。缺点： - 最长只能处理 max_len 个位置（BERT 是 512，GPT-2 是 1024）； - 即使把表做得很大，训练时从没出现过的位置，它们的向量也没有被训练过（还是随机初始化的值）。 这会影响模型能否推广到比训练时更长的句子 —— 05 号文件用实验比较四种做法。""") if __name__ == "__main__": main()