""" ============================================================================= 第 3 讲 · 代码 2：位置编码 —— 从正弦编码到 RoPE（旋转位置编码） ============================================================================= 对应讲义：lectures/lecture03_architectures_hyperparameters.md 第 5 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 注意力本身是"不看位置"的：把输入序列打乱顺序，每个 token 得到的注意力输出只是跟着换了位置， 数值完全不变。"狗咬人"和"人咬狗"在它眼里是同一袋词。所以必须把位置信息"注入"模型。 我们想要的性质：**注意力分数只依赖两个 token 的相对距离 m - n**，而不是它们的绝对位置。 （"前一个词"对理解当前词的作用，不应该因为这句话出现在文章开头还是结尾而不同。） 这个文件： 1. 实现原始 Transformer 的正弦位置编码，并用数值实验说明它**不满足**上面的性质； 2. 实现 RoPE（CS336 作业 1 的接口），验证它**满足**这个性质； 3. 从复数的角度理解 RoPE，看看它的"多种转速"； 4. 边界：两种常见的 RoPE 写法（相邻配对 vs 前后半配对）不能混用。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 02_rope.py """ import math import torch def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：正弦位置编码（原始 Transformer，2017） # ============================================================================= def sinusoidal_positions(max_len, d_model): """ 正弦位置编码表。 【公式】 PE[pos, 2i] = sin(pos / 10000^(2i/d)) PE[pos, 2i+1] = cos(pos / 10000^(2i/d)) 使用时直接**加**到词嵌入上：输入 = 词嵌入 + PE[位置]。 【返回】(max_len, d_model) 的张量 【边界】d_model 必须是偶数（sin 和 cos 成对出现）。 """ if d_model % 2 != 0: raise ValueError("d_model 必须是偶数") pos = torch.arange(max_len).float()[:, None] # (max_len, 1) freq = 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model)) # (d/2,) pe = torch.zeros(max_len, d_model) pe[:, 0::2] = torch.sin(pos * freq) pe[:, 1::2] = torch.cos(pos * freq) return pe # ============================================================================= # 第 2 部分：RoPE # ============================================================================= class RotaryPositionalEmbedding(torch.nn.Module): """ 旋转位置编码 RoPE（Su et al. 2021；CS336 作业 1 的接口）。 【思路】 把 d_k 维向量的分量两两配对：(x0, x1), (x2, x3), ...，每一对看成二维平面上的一个点。 位置 i 的 token，第 k 对要逆时针旋转角度 θ_{i,k} = i × Θ^(-2k/d_k)（k 从 0 开始）。 不同的对转速不同：第 0 对转得最快（每个位置转 1 弧度），最后一对转得极慢。 旋转一个二维点 (a, b) 角度 φ： a' = a cos φ - b sin φ b' = a sin φ + b cos φ **只对 query 和 key 做旋转，不对 value 做。** 为什么这样能编码相对位置？见讲义 5.4 节或下面的 demo_relative_property。 【参数】 theta : Θ，常用 10000（Llama 3 用 500000，能更好地处理长上下文） d_k : 每个头的维度，必须是偶数 max_seq_len : 最大序列长度（预先计算好这么多位置的 cos/sin） 【形状】 forward(x, token_positions)： x : (..., seq_len, d_k) token_positions : (..., seq_len) 整数，每个 token 的位置 返回 : 与 x 同形状 【边界】 - token_positions 的值必须在 [0, max_seq_len) 之间。 - cos/sin 表用 register_buffer(persistent=False) 保存：它会随模型一起搬到 GPU， 但不是可训练参数，也不会被存进检查点（因为随时可以重新算出来）。 """ def __init__(self, theta, d_k, max_seq_len, device=None): super().__init__() if d_k % 2 != 0: raise ValueError("d_k 必须是偶数") inv_freq = 1.0 / (theta ** (torch.arange(0, d_k, 2, device=device).float() / d_k)) # (d_k/2,) positions = torch.arange(max_seq_len, device=device).float() angles = torch.outer(positions, inv_freq) # (max_seq_len, d_k/2) # 【PyTorch 小知识】register_buffer：登记一个"不是参数但属于模型"的张量 self.register_buffer("cos", angles.cos(), persistent=False) self.register_buffer("sin", angles.sin(), persistent=False) def forward(self, x, token_positions): cos = self.cos[token_positions] # (..., seq, d_k/2) sin = self.sin[token_positions] x_even = x[..., 0::2] # 每对的第 1 个分量 x_odd = x[..., 1::2] # 每对的第 2 个分量 out_even = x_even * cos - x_odd * sin out_odd = x_even * sin + x_odd * cos # 把两部分重新交错排列回 (x0', x1', x2', x3', ...) return torch.stack([out_even, out_odd], dim=-1).flatten(-2) def rope_half_split(x, positions, theta=10000.0): """ 另一种常见的 RoPE 写法（GPT-NeoX、Hugging Face Llama 实现等采用）： 不是把相邻的 (x0, x1) 配对，而是把 (x_i, x_{i + d/2}) 配对（前一半和后一半配对）。 【边界】 两种写法在数学上等价（只是维度的排列顺序不同），但**用一种写法训练的权重， 必须用同一种写法推理**，否则位置编码就全乱了。 从一种格式的检查点转换到另一种时，需要对 W_Q、W_K 的行做相应的重排。 """ d = x.shape[-1] inv_freq = 1.0 / (theta ** (torch.arange(0, d, 2).float() / d)) angles = positions[..., None].float() * inv_freq # (..., seq, d/2) cos, sin = angles.cos(), angles.sin() x1, x2 = x[..., : d // 2], x[..., d // 2:] return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) # ============================================================================= # 第 3 部分：演示 # ============================================================================= def demo_attention_is_permutation_invariant(): """演示 0：没有位置编码的注意力不关心顺序。""" print_header("演示 0：没有位置编码时，注意力看不到顺序") torch.manual_seed(0) x = torch.randn(5, 16) # 5 个 token wq, wk, wv = (torch.randn(16, 16) for _ in range(3)) def attn(inp): q, k, v = inp @ wq, inp @ wk, inp @ wv return torch.softmax(q @ k.T / 4, dim=-1) @ v perm = torch.tensor([3, 0, 4, 1, 2]) out, out_perm = attn(x), attn(x[perm]) print(f" 打乱输入顺序后，输出只是跟着换了位置：{torch.allclose(out[perm], out_perm, atol=1e-5)}") print(" → 注意力是'一袋词'模型：'狗咬人' 和 '人咬狗' 对它来说没有区别。必须注入位置信息。") def demo_sinusoidal(): """演示 1：正弦位置编码不满足"只依赖相对位置"。""" print_header("演示 1：正弦位置编码（加到词嵌入上）") torch.manual_seed(0) d = 64 pe = sinusoidal_positions(100, d) x, y = torch.randn(d), torch.randn(d) # 两个词的嵌入 print(" 同样两个词、相对距离都是 3，只改变绝对位置：") for m in [0, 10, 50]: n = m + 3 score = torch.dot(x + pe[m], y + pe[n]).item() print(f" 位置 ({m:>2}, {n:>2})：点积 = {score:8.3f}") print(""" 分数随绝对位置变化。原因：展开后有交叉项 = + + + 其中 和 依赖绝对位置 m、n，而不只是 m - n。""") def demo_rope_relative_property(): """演示 2：RoPE 的点积只依赖相对位置。""" print_header("演示 2：RoPE —— 点积只依赖相对距离") torch.manual_seed(0) d = 64 rope = RotaryPositionalEmbedding(theta=10000.0, d_k=d, max_seq_len=2048) q, k = torch.randn(d), torch.randn(d) def score(m, n): qm = rope(q[None], torch.tensor([m]))[0] kn = rope(k[None], torch.tensor([n]))[0] return torch.dot(qm, kn).item() print(" 同样的 q、k，相对距离固定为 3，只改变绝对位置：") base = score(0, 3) for m in [0, 10, 50, 1000]: print(f" 位置 ({m:>4}, {m + 3:>4})：点积 = {score(m, m + 3):9.5f}") assert all(abs(score(m, m + 3) - base) < 1e-3 for m in [10, 50, 1000]) qr = rope(q[None], torch.tensor([7]))[0] print(f"\n 旋转不改变长度：|q| = {q.norm().item():.5f}，|RoPE(q, 7)| = {qr.norm().item():.5f}") print(""" 为什么？旋转的性质：把 q 转 mθ、k 转 nθ 后再做点积，等于把 q 转 (m - n)θ 后与原来的 k 做点积： = 它只依赖 m - n。RoPE 对每一对分量都这样做（每对转速不同），所以整体点积也只依赖 m - n。""") def demo_complex_view(): """演示 3：用复数理解 RoPE。""" print_header("演示 3：复数视角 —— 旋转 = 乘以 e^(iφ)") torch.manual_seed(0) d, pos = 8, 5 rope = RotaryPositionalEmbedding(theta=10000.0, d_k=d, max_seq_len=16) x = torch.randn(d) # 把相邻两个实数看成一个复数：z_k = x_{2k} + i x_{2k+1} z = torch.view_as_complex(x.view(-1, 2)) phi = pos * (1.0 / (10000 ** (torch.arange(0, d, 2).float() / d))) z_rot = z * torch.exp(1j * phi) # 乘以单位复数 = 旋转 via_complex = torch.view_as_real(z_rot).flatten() via_rope = rope(x[None], torch.tensor([pos]))[0] print(f" 复数乘法 与 RoPE 实现 的最大差异：{(via_complex - via_rope).abs().max().item():.2e}") print(" 复数点积 Re(z_q · conj(z_k)) 中，e^(imφ) · conj(e^(inφ)) = e^(i(m-n)φ)，只剩下相对位置。") print("\n 各对分量的转速（d_k = 64，Θ = 10000）：") d = 64 inv_freq = 1.0 / (10000 ** (torch.arange(0, d, 2).float() / d)) for k in [0, 1, 4, 8, 16, 31]: wavelength = 2 * math.pi / inv_freq[k].item() print(f" 第 {k:>2} 对：每个位置转 {inv_freq[k].item():.2e} 弧度，转一整圈需要 {wavelength:>10,.0f} 个位置") print(""" 像一个有很多根指针的钟：秒针（高频）区分相邻的位置，时针（低频）区分很远的位置。 Θ 越大，最慢的指针越慢，能区分的距离越远 —— 所以 Llama 3 把 Θ 从 10000 调到了 500000 以支持长上下文。""") def demo_long_range_decay(): """演示 4：RoPE 让"相似的 q、k"的分数随距离整体衰减。""" print_header("演示 4：距离越远，分数的上界越小") d = 128 rope = RotaryPositionalEmbedding(theta=10000.0, d_k=d, max_seq_len=4096) v = torch.ones(d) # q = k = 全 1 向量（最"相似"的情况） base = rope(v[None], torch.tensor([0]))[0] for dist in [0, 1, 4, 16, 64, 256, 1024]: rotated = rope(v[None], torch.tensor([dist]))[0] print(f" 距离 {dist:>5}：点积 = {torch.dot(base, rotated).item():8.2f}") print(""" 同一个向量放在两个位置上，点积随距离（总体趋势上）减小：不同转速的分量逐渐"失去同步"、相互抵消。 这给了模型一个温和的归纳偏置：近处的 token 更容易被关注。（数值会上下波动，不是单调的。）""") def demo_two_conventions(): """演示 5：两种 RoPE 写法等价但不能混用。""" print_header("演示 5：【边界】相邻配对 vs 前后半配对") torch.manual_seed(0) d, L = 8, 6 x = torch.randn(L, d) pos = torch.arange(L) interleaved = RotaryPositionalEmbedding(10000.0, d, 16)(x, pos) half = rope_half_split(x, pos) # 排列：把相邻配对的维度顺序 (0,1,2,3,...) 变成前后半配对的顺序 (0,2,4,...,1,3,5,...) perm = torch.cat([torch.arange(0, d, 2), torch.arange(1, d, 2)]) print(f" 直接比较两种写法的输出：最大差异 {(interleaved - half).abs().max().item():.3f}（不同！）") same = torch.allclose(interleaved[:, perm], rope_half_split(x[:, perm], pos), atol=1e-6) print(f" 先按排列重排输入维度，再比较：{same}（数学上等价）") print(" → 加载别人的检查点时，一定要确认它用的是哪一种 RoPE 写法。") def main(): demo_attention_is_permutation_invariant() demo_sinusoidal() demo_rope_relative_property() demo_complex_view() demo_long_range_decay() demo_two_conventions() print("\n全部演示运行完毕。") if __name__ == "__main__": main()