跳到主要内容

RoPE — 用旋转编码位置,以及撑长上下文的三级火箭

这一章讲三件事: 注意力为什么天生是「不识数」的(不知道谁在前谁在后); RoPE 怎么用「旋转」这一个动作同时把绝对位置和相对距离都编码进去; 以及当输入比训练时长了 30 倍,位置编码怎么不散架(NTK/YaRN)。 这是本书第一个「数学多于代码」的零件——但主走查只需要中学三角。

1. 先看现象:注意力分不清「狗咬人」和「人咬狗」

第 01 章说过,注意力把整段文字「一起看」。一起看有个副作用: 在注意力眼里,换序之后的两句话是同一个东西——它天生对顺序视而不见(原书用的说法:排列不变)1。 「狗咬人」和「人咬狗」,词一样、关注度一样,含义天差地别。

位置编码(position encoding,给每个词附上「我在第几位」的信息)就是补这个洞的。 原书把历史路线摆成一张进退两难的对照表2:

路线做法死穴
绝对式(正弦函数,2017 原版)每个位置配一串固定的正弦/余弦——余弦,就是一种波浪形的数学曲线——数,加到词表示上超出训练长度的位置没见过,撑不长
绝对式(可学习,GPT-2/3)每个位置的可训练向量同样撑不长;还多出一大块参数
相对式(Transformer-XL、T5)直接把「隔多远」加进注意力分数远程好,但推理时相对距离一直在变,缓存(存起来免得重算的仓库)没法用了

注意表里没有一条路同时满足:能外推(在没见过的长度上还能工作) + 不占参数 + 能缓存。RoPE 就是冲着「全都要」来的3

2. 主走查:旋转怎么把「距离」变成「分数」

RoPE 的核心动作一句话:把每个头的 64 维两两配对,当成 32 个二维平面上的点,按位置转个角度。 角度 = 位置 × 频率。先在二维里手算(角度 30° 是为演示选的;真实频率由第 04 章的 rope_theta=10000 派生)4:

取一个二维向量 k = [1, 0] (一个「键」)。

位置 0:转 0° → [1.00, 0.00]
位置 1:转 30° → [0.87, 0.50]
位置 2:转 60° → [0.50, 0.87]

现在拿「查询」去点乘「键」,看分数怎么走:

查询@位置1(转30°) · 键@位置0(转0°) = 0.87×1.00 + 0.50×0.00 = 0.87
查询@位置2(转60°) · 键@位置0(转0°) = 0.50×1.00 + 0.87×0.00 = 0.50
查询@位置2(转60°) · 键@位置1(转30°) = 0.87 ← 又是 0.87!

图说:三个分数只由「两个位置隔多远」决定——隔 1 格永远 0.87,隔 2 格永远 0.50。
这正是中学三角的恒等式 cos(A−B) 在干活:两个旋转向量的点积 = cos(角度差)。
顺便注意:分数随距离衰减(0.87 → 0.50),越远的词天然越不受重视。

这一张表同时兑现了 RoPE 的两个承诺:分数只看相对距离(绝对位置被旋转「吸收」掉了), 远处的词自动衰减——后者恰好符合语言的常识:相邻的词更相关5

还有一个白送的:旋转不改变向量长度(转 30° 的 [1,0] 还是 0.87²+0.50²=1), 词义本身的信号一点没被稀释——这就是原书说的「保模长」6

3. 高维怎么推广:粗细两套频率

真实的一个头有 64 维,配成 32 个平面,每个平面用不同的转速: 低维转得慢(粗钟:管「大概在前半段」这种远距离粗关系), 高维转得快(细钟:管「正好隔 3 个词」这种近距离精关系)7

为什么要两套钟?只用慢钟,相邻的词分数几乎一样,分不出「谁紧跟谁」; 只用快钟,远距离上角度转了无数圈,分数成了乱数。粗细搭配,近处精、远处糙,都有据可依。 这 32 个转速由一条公式从 rope_theta=10000 派生(第 04 章埋的那个数在这里接线)8

代码实现反而朴素:把张量(可套多维的大数表)沿最后一维切成两半 x1, x2, 输出 (x1·cos − x2·sin, x2·cos + x1·sin) 拼回去——就是 §2 那个二维旋转的批量版9。 cos/sin 表按位置预计算好,推理时来一个新词只需算一行10

4. 三级火箭:训练 4k,用到 128k

RoPE 本身只解决「位置有编码」。真正的难题在外推(在超出训练长度的输入上还能工作): 快钟的高维在远距离上转了太多圈,分数乱掉。原书给了三级方案,层层递进11:

名字做法直觉
直接外推什么都不做快钟乱掉,质量崩——这是「不做会怎样」的对比基准
NTK 感知缩放高维的钟拨慢(频率除以缩放系数),低维保持原速远处的词用慢钟凑合看,近处的准头不牺牲
YaRN在 NTK 基础上把「拨慢多少」按维度做成平滑过渡,再加一个缩幅系数稳住数值精细版的二级,128k+ 上下文的主流做法

第二、三级共用同一套实现(本书 RotaryEmbedding 类里的那段):设缩放系数 s>1 时, 按「低维保持原频率(外推)、高维除以 s(插值:按比例往回收)」的带状区域做平滑过渡, 过渡带的位置由第 04 章那两个参数 ntk_alpha、ntk_beta 圈定12; 另外乘一个随 s 微增的「聚拢系数」,把 cos/sin 的幅度往 1 收,防数值漂移13

原书给的两个直觉比喻值得留下:ntk_alpha 像琴弦的张力(太紧会断=旋转过冲,太松没音=位置感消失), ntk_beta 定校正多早介入(早=中长文本就修,晚=专攻超长)14

效果口径(书里引的数字):LLaMA-3 这类模型把上下文撑到 128k token; 对比实验里 RoPE 系的模型收敛快 20-30%,困惑度(第 10 章细讲,越低越准)3.5 对 4.015

5. 作者的判断与证据

说法书里的证据我们的标注
RoPE 同时拿下「外推+零参数+可缓存」机制推导(§2 的三角恒等式)+ 与三条旧路线的逐项对照推导成立,是这三个性质的结构性后果
RoPE 已是行业标准点名 LLaMA-2/3、Mistral、Mixtral、PaLM、Gemma、CodeGen 采用16领域共识
YaRN 能撑 128k+引 LLaMA-3 的上下文长度厂商口径,机制本身可信

两处原书自己的笔误,本章顺手修正(读原书时别被带偏):

判断(我们的,不是书里的): 原书章末总结把提出 RoPE 的论文《RoFormer》误写成 "Reformer"—— Reformer 是另一篇 2020 年的、处理别的问题的论文,不相干; 正文里还把 NTK 展开成 "Neuro Tangent Kernel",正确展开是 Neural Tangent Kernel(神经正切核,2018 年一篇讲无限宽网络训练动力学的论文)17。 两处都不影响机制讲解,但引用时要以论文原题为准。 如果错,会错在: 如果 Apress 后续勘误已修正这两处,我们的「笔误」标注就过时了——核对纸面原文即可裁决。

6. 边界与局限

  • RoPE 的频率表(粗细搭配)靠 rope_theta 人工定;原书也承认「频率选择要小心平衡」18, 后来的自适应底数研究它没展开;
  • NTK/YaRN 是不重训的应急术:拨慢钟是让模型「将就」没见过的长度, 远距离的准头仍逊于从头在长文本上训过的模型——原书没有量化和从头训练的差距;
  • 二维直觉到高维的推广依赖「两两配对」这个实现约定,换一种配对方式(如相邻维配对)公式要重推,书里未讨论。

7. 可带走的

  1. 注意力天生不知道词序;位置编码是补丁,不是装饰;
  2. 三条旧路线各有死穴:绝对式外推差,相对式没法缓存;
  3. RoPE = 把维度两两配对按位置旋转;点积(对应位相乘再相加的乘法)只看角度差 = 只看相对距离;
  4. 旋转保长度,词义信号不被位置信号稀释;
  5. 粗细两套钟:低维管远距离粗关系,高维管近距离精关系;
  6. 距离衰减是白送的:远词天然降权,符合语言直觉;
  7. 撑长上下文 = 把高维的钟拨慢(NTK),按维度平滑过渡(YaRN);这是不重训的将就,不是免费午餐;
  8. 128k 上下文不是架构奇迹,是「旋转 + 缩放」这套几何手艺的直接后果。

8. 原文地图

主题原书章原文位置
RoFormer 2021 与 RoPE 定位Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:17(搜「RoFormer」)
排列不变性与「狗咬人」Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:31(搜「the dog bites the man」)
三条旧路线与死穴Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:55(搜「Sinusoidal Positional Embeddings」) · text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:75(搜「Learned Absolute」) · text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:111(搜「relative positional bias」)
二维旋转与相对距离Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:221(搜「rotation matrix」) · text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:231(搜「depends only on relative distance」)
保模长Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:275(搜「Norm Preservation」)
距离衰减Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:253(搜「cosine near 1」)
高维块对角与粗细频率Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:245(搜「block-diagonal」) · text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:249(搜「coarse」)
预计算 cos/sinRotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:263(搜「Precomputing」)
收敛 20-30% 与困惑度 3.5/4.0Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:427(搜「20-30%」)
128k 与 LLaMA-3What Is YaRN (Yet Another RoPE Extension)text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:1053(搜「128k+」)
工业采用名单Rotary Positional Embeddings: An In-Depth and Comprehensive Explorationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:373(搜「LLaMA Family」)
_apply_rotary_emb 代码RoPE Embeddings Implementationtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:449(搜「torch.chunk」) · text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:455(搜「torch.cat」)
频率表与缩放混合The RotaryEmbedding Class Definition and Methodstext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:609(搜「geometrically」) · text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:631(搜「extrapolation」) · text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:635(搜「interpolation」)
low/high 与 clampThe RotaryEmbedding Class Definition and Methodstext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:623(搜「clamp」)
concentration 聚拢系数What Is NTK (Neural Tangent Kernel) in the Context of RoPEtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:989(搜「concentration factor」)
NTK 出处与外推问题What Is NTK (Neural Tangent Kernel) in the Context of RoPEtext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:961(搜「Jacot」) · text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:975(搜「oscillate rapidly」)
YaRN 定义What Is YaRN (Yet Another RoPE Extension)text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:1019(搜「YaRN extends」)
琴弦张力比喻Model Parameters Configuration and Mathematical Foundations(第 04 章)text/39-ch04-4-rms-normalization-and-model-configuration.txt:267(搜「guitar string」)
原书笔误:ReformerSummarytext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:1063(搜「Reformer」)
原书笔误:Neuro Tangent KernelThe RotaryEmbedding Class Definition and Methodstext/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:773(搜「Neuro Tangent Kernel」)

Footnotes

  1. 出处:「Rotary Positional Embeddings: An In-Depth and Comprehensive Exploration」第 31 段(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:31,搜「permutation-invariant」)。原文:自注意力是排列不变的,把输入 token 当无序集合,「the dog bites the man」与「the man bites the dog」无法区分。

  2. 出处:同章「Early Transformer Positional Encodings」「Relative Positional Encodings」两组(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:55,搜「Sinusoidal Positional Embeddings」;:75,搜「Learned Absolute Positional Embeddings」;:97,搜「Transformer-XL」;:111,搜「T5 relative positional bias」)。正弦式外推差;可学习式参数多;T5 相对偏置在自回归推理时因相对距离移动而无法缓存注意力分数。

  3. 出处:「Emergence of Rotary Positional Embeddings (RoPE)」第 127 段(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:127,搜「unifying the strengths」)。原文:RoPE 用旋转矩阵编码绝对位置、同时让注意力分数只依赖相对距离,统一了先前方法的长处;零参数、保模长、可缓存(第 131-139 段)。

  4. 演示角频率 30° 为演示编造;真实频率由 base^(2i/d) 派生(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:967,搜「block-diagonal rotation matrices」)。旋转-点积恒等式见原文:text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:231(搜「depends only on relative distance」)——「由于旋转矩阵满足该性质,分数只依赖相对距离」。

  5. 出处:「Attention Decay and Frequency Selection」第 253 段(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:253,搜「cosine near 1」)。原文:小距离保持 cos 接近 1(强注意力),大距离振荡衰减(弱注意力),与「邻近词更相关」的语言直觉一致。

  6. 出处:「Key Properties」第 3 条(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:275,搜「Norm Preservation」)。原文:正交旋转矩阵保持嵌入的模长,位置编码不扭曲语义内容。

  7. 出处:「Generalizing to High Dimensions」(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:245,搜「block-diagonal」;:249,搜「Lower frequencies capture coarse」)。原文:独立二维旋转块、频率各不相同;低频抓粗位置关系,高频抓细粒度。

  8. 出处:「The RotaryEmbedding Class Definition and Methods」(_compute_concentration_and_inv_freq 第 1 步)(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:609,搜「geometrically」)。原文代码:freq = base ** (arange(0, d, 2) / d),base 即第 04 章配置里的 rope_theta=10000。

  9. 出处:「RoPE Embeddings Implementation」(_apply_rotary_emb 代码)(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:449,搜「torch.chunk」;:455,搜「torch.cat」)。原文:切半后按 (x1·cos − x2·sin, x2·cos + x1·sin) 拼接,并做 dtype 安全转换。

  10. 出处:「Efficient Implementation」(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:263,搜「Precomputing」)。原文:预计算所有位置的 cos/sin 进一步优化推理。

  11. 出处:「NTK Scaling Mechanism」(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:979,搜「blends two approaches」)与「What Is YaRN」(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:1019,搜「YaRN extends」)。原文:外推用原频率、插值用 1/(s·θ),ramp 平滑过渡;YaRN 在 NTK 之上精细化,支撑 128k+。

  12. 出处:text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:617(搜「ntk_beta」)与 :623(搜「clamp」)。原文代码:low/high 由 initial_context_length 与 ntk_beta/ntk_alpha 定出,clamp 到安全带,mask=1−ramp 在两档频率间线性过渡。

  13. 出处:text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:615(搜「concentration」)。原文代码:concentration = 0.1·log(scaling_factor) + 1.0,乘在 cos/sin 上稳定数值。

  14. 出处:第 04 章「The Geometry of Position」一节(text/39-ch04-4-rms-normalization-and-model-configuration.txt:267,搜「guitar string」;:271,搜「kick in early」)。原文:ntk_alpha 像琴弦张力(太紧会断、太松失声),ntk_beta 决定校正介入的早晚。

  15. 出处:text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:427(搜「20-30%」)与 :133(搜「128k tokens」)。原文:RoPE 实验中收敛快 20-30%、困惑度 3.5 对 4.0;LLaMA-3 借助缩放支撑 128k 上下文。

  16. 出处:「Applications in Modern Models」(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:373,搜「LLaMA Family」;:375,搜「Mistral AI」;:377,搜「PaLM and Gemma」)。

  17. 两处笔误:章末 Summary 称 RoPE 出自 2021 年论文 "Reformer"(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:1063,搜「Reformer」),而同章正文与脚注都写《RoFormer》(Su et al., 2021,:17:1083);NTK 被展开为 "Neuro Tangent Kernel"(:773,搜「Neuro Tangent Kernel」),Jacot 2018 年原题为 Neural Tangent Kernel(:961,搜「Jacot」)。

  18. 出处:「Summary of Traditional Positional Embedding Approaches」后的 RoPE 弱点条(text/40-ch05-5-rotary-positional-embeddings-integrating-ntk-a.txt:141,搜「frequency selection」)。原文:RoPE 需要小心选择频率以平衡粗细两种位置尺度。