时长 8 小时:§1 相对位置的谱系(1h)· §2 RoPE 推导与实现(2.5h)· §3 长度外推:PI / NTK / YaRN / ALiBi(2h)· §4 现代版 mini-GPT 与外推实验(1.5h)· §A–§C(1h)
对应正课 Lecture 2 的位置编码部分:固定、可学习、RoPE 及其应用。学习计划要求能说清”RoPE 相比绝对位置编码的好处、为什么能外推、NTK-aware / YaRN 在改什么”。
配套代码
w1day5_checks.py(纯 numpy):RoPE 的相对性(含投影)、旋转的性质、两种配对约定、远程衰减、哪些维度会外推失败、PI / NTK / YaRN 的角频率、ALiBi 的斜率、RoPE 下缓存不失效、bf16 算角度的误差§B 参考实现(
code/目录):w1d5_rope.py(RoPE / ALiBi / 三种扩展;现代版 mini-GPT;四种位置编码的外推实验)。今天的目标:能从”打分只依赖相对位置”这一个要求推出 RoPE;能指出 RoPE 外推失败的具体维度;能说清 PI、NTK-aware、YaRN 各自改了什么、为什么;手里有一个 RoPE + GQA + RMSNorm + SwiGLU 的现代版 mini-GPT。
§0 今天的主线:还前几天欠下的三笔债
| 欠在哪 | 问题 | 今天怎么还 |
|---|---|---|
| Day 2 §3.6 | 正弦编码本身有完美的相对性,经过 投影后就没了 | RoPE:把旋转放在投影之后(§2) |
| Day 4 §1.6 | 绝对位置编码下窗口一滑,KV Cache 全部失效 | RoPE:打分只看位置差,缓存永远有效(§2.6) |
| Day 3 §2.4 | 可学习位置嵌入有硬上限,超过 block_size 无法处理 | RoPE 能算任意位置,但”能算”不等于”算得好”(§3) |
第三笔债只还了一半,这就引出了今天的后半段:长度外推。
今天要用到的旧结论:Day 2 §3.3–§3.5(正弦编码、、点积只依赖距离);Day 2 §3.6(投影破坏相对性);Day 4 §1.6(缓存失效);Day 1 §4.5( 与 只有乘积有意义)。
§1 相对位置编码的谱系
1.1 动机:模型真正需要的是相对位置
“它”指代的名词通常在几个词之前;形容词修饰的名词通常在后一个。这些规律和”它在第 137 个位置”无关,只和距离有关。一个在位置 10 学会”看前一个词”的头,应该在位置 1000 也会。
绝对位置编码要求模型在每一个位置上分别学会这件事(Day 2 §3.6 的第 7 组:同样相距 3,不同位置打分完全不同)。相对位置编码则让打分天然只依赖距离。
1.2 两条路线
| 路线 | 做法 | 代表 |
|---|---|---|
| 加法偏置 | 在打分上加一个只依赖 的项: | Shaw et al.(2018,学习一组相对位置向量,距离截断到 );T5(每个头一组标量偏置,距离按对数分桶);ALiBi(固定的线性偏置,§3.6) |
| 乘法旋转 | 按位置旋转 和 ,让内积本身只依赖 | RoPE(Su et al. 2021) |
加法偏置直观,但每一对 都要额外加一项,和 FlashAttention 这类融合内核配合起来麻烦一些;RoPE 只改 和 ,注意力本身一行不动,任何注意力实现都能直接用。这是 RoPE 在工程上胜出的重要原因之一。
§2 RoPE:从一个要求推出来
2.1 目标
找一个函数 ,把位置 注入向量 ,使得
内积只依赖内容和相对位置,与 各自的值无关。
2.2 二维:复数一行就解决
先看 ,把二维向量当成复数 。复数的”内积”是 。
试试乘一个单位复数:。
和 只以差的形式出现。 乘 就是逆时针旋转 角,写成实矩阵:
为什么是旋转,不是别的? 旋转矩阵满足 (第 2 组验证:误差 ),所以
两个位置的旋转”相减”了。 这就是全部的数学。
2.3 高维: 个不同频率的时钟
把 维切成 对,第 对用自己的角频率 旋转:
每一对都满足 ,块对角拼起来仍然成立。打分:
w1day5_checks.py 第 1 组:先用随机的 投影、再旋转,位置对 (10,13)、(110,113)、(200,203)、(5000,5003) 的打分到小数点后 11 位都相同。对照 Day 2 第 7 组:加法式正弦编码在同样的设置下给出 4.935、0.302、2.819。
2.4 与正弦编码的关系
频率完全相同: 就是 Day 2 §3.3 的 。Day 2 §3.4 推出正弦编码满足 ,Vaswani 等人”假设”这能让模型学会按相对位置注意。
区别只在旋转作用在哪里:
- 正弦编码:把位置向量加到输入上,再投影。投影 夹在两个旋转中间,一般不与旋转可交换,相对性被破坏(Day 2 §3.6)。
- RoPE:先投影,再把旋转乘在 上。投影在旋转之外,不干扰 。
RoPE 是把 2017 年那个”假设”变成了一个”保证”。 这是一个很好的例子:好的设计往往不是新想法,而是把一个旧想法放在对的位置上。
2.5 实现:逐元素,不需要矩阵
块对角矩阵乘法可以写成逐元素运算。对一对 :
两种配对约定,都常见,不能混用:
| 约定 | 哪两维配成一对 | 谁在用 |
|---|---|---|
| 相邻配对 | RoFormer 原文、Meta 的 Llama 原始代码(用复数实现)、GPT-J | |
两半配对(rotate_half) | GPT-NeoX、HuggingFace 的 Llama 实现 |
def rotate_half(x): # 两半配对的写法
x1, x2 = x[..., :d//2], x[..., d//2:]
return torch.cat([-x2, x1], dim=-1)
q_rot = q * cos + rotate_half(q) * sin # cos、sin 形状 (T, d),前后两半重复
第 3 组:两种约定只差一个维度置换,数学上等价;但同一对 用错约定,打分从 2.349 变成 13.069,不报错。HuggingFace 把 Meta 的 Llama 权重转换过来时,会对 、 的行做一次置换,正是为了弥合这两种约定。加载别人的权重时,先确认约定。
2.6 性质
- 保范数:旋转不改变长度(第 2 组)。 的尺度不受位置影响。
- 没有参数:不增加任何可学习参数,也没有
max_len的表。 - 只转 和 ,不转 :位置只应该影响”看谁”(路由),不应该影响”拿什么”(内容)。用 Day 1 §2.4 的话说,RoPE 只改 QK 电路,不碰 OV 电路。
- 每一层都加:绝对位置编码只在输入处加一次,靠残差流传到深层;RoPE 在每一层的注意力里都重新注入位置。
- 远程衰减:RoFormer 证明了一个随相对距离下降的上界。第 4 组():该上界从距离 0 的 32.5 降到 100 的 10.2、1000 的 4.5,但有起伏,4000 时是 5.0。与 Day 2 §3.5 的正弦点积是同一个形状,因为频率相同。这是一个弱先验,不是硬约束:模型完全可以学出关注远处的头。
- 缓存友好(Day 4 §1.6 的债):缓存里存的是按绝对位置旋转过的 。后来的 query 按它自己的绝对位置旋转,内积只看差值。第 10 组:用旧缓存算的打分与”窗口重新编号后重算”的打分差 。位置编号可以一直往上数,缓存永远不用重建。
2.7 收益与代价
收益:相对性有保证;无参数;与任何注意力内核兼容;缓存友好;原则上能处理任意位置。代价:每层多一次逐元素运算(很便宜);两种约定的兼容性坑;以及最重要的一点:原则上能处理,不等于实际上处理得好(§3)。
§3 长度外推
3.1 问题:RoPE 为什么也会外推失败
在长度 上训练,在 上推理。可学习位置嵌入直接没有向量(Day 2 §3.7);RoPE 能算出任意位置的旋转,可实验上一旦超过训练长度,困惑度照样会急剧上升。为什么?
看每一对维度在训练中见过哪些角度。 第 对在位置 的角度是 ,训练时 ,见过的角度范围是 。
- 高频维度( 大):训练长度内已经转了很多圈, 的每个角度都见过。位置再大,也只是”再转几圈”,不是新东西。
- 低频维度( 小):波长 超过 ,训练中连一圈都没转完。推理到 时,这些维度转到了从没见过的角度。对模型而言,这是分布外输入。
第 5 组数出了这样的维度有多少(波长 > 训练长度的维度对数):
| base | 训练 2048 | 训练 4096 | 训练 8192 | |
|---|---|---|---|---|
| 10000 | 64 | 11 / 32 | 9 / 32 | 7 / 32 |
| 10000 | 128 | 23 / 64 | 18 / 64 | 14 / 64 |
| 500000 | 128 | 35 / 64 | 32 / 64 | 29 / 64 |
问题出在低频维度。 接下来的三种方法,本质上都在回答同一个问题:“怎样让低频维度在更长的序列上,仍然只看到训练时见过的角度?”
3.2 位置插值(PI):把位置压缩回训练范围
Chen et al.(2023):要把上下文从 扩展到 ,就把位置 换成 :
所有位置都落回 ,所有角度都在训练范围内(第 6 组:插值后最大角度 / 训练时最大角度 = 1.000)。配合约 1000 步微调,他们把 Llama 扩展到了 32k。
代价:所有维度被均匀压缩。低频维度需要压缩,高频维度本来不需要,却也丢了分辨率:第 6 组, 时相邻 token 在最高频维度上的角度差从 1.000 rad 变成 0.250 rad,相邻的 token 更难区分了。这就是 PI 必须微调的原因。
3.3 NTK-aware:不改位置,改 base
直觉:既然只有低频维度需要插值,那就只压缩低频、不动高频。有没有一个简单的办法做到?
改 base。 ,把 base 调大,所有 都会变小,但变小的程度随 增大: 时 ,与 base 无关,最高频永远不变。
推导 base 该调多大。 要求最低频那一对()恰好插值 倍:
第 7 组(,):base’ = 40,889.9;最高频 ,最低频 0.2500,中间第 8、16、32、48 对分别是 0.839、0.703、0.495、0.348。插值量从高频到低频平滑过渡。
这个方法最早以一篇 Reddit 帖子的形式出现(2023 年,名字来自神经正切核理论的一个类比),不微调就能明显改善外推,很快被广泛采用。“动态 NTK”会随当前序列长度实时调整 。
直接调大 base 是同一思路的”训练时版本”:Code Llama 把 base 设为 ,Llama 3 设为 。更大的 base 意味着更多维度的波长远超训练长度(上表第三行),给后续扩展留了余地。
3.4 YaRN:按”转了几圈”分段处理
Peng et al.(2023)把”高频不动、低频插值”做得更精确:按每对维度在训练长度内转了几圈 来决定:
- (转了很多圈,高频):完全不插值
- (不到一圈,低频):完全插值,
- 中间:线性过渡
对 Llama,论文取 。第 8 组(训练 4096,扩展 16 倍,):21 对完全不插值、18 对完全插值、25 对过渡, 落在 。
第二个修正:注意力温度。 序列变长后,softmax 要在更多 key 上分配权重,分布会被摊平。YaRN 给 logits 乘一个随扩展倍数增长的系数,(q、k 各乘一次,logits 乘它的平方)。第 8 组: 时 。
YaRN 用很少的微调数据就达到了很好的长上下文效果,是目前最常用的扩展方法之一。Llama 3.1 用了一个思路相近的分段方案(HuggingFace 里叫 rope_type="llama3"),把 8k 扩展到 128k。
3.5 三种方法对照
| 方法 | 改了什么 | 高频维度 | 低频维度 | 需要微调吗 |
|---|---|---|---|---|
| PI | 位置 | 被压缩(丢分辨率) | 被压缩 ✓ | 需要(约 1000 步) |
| NTK-aware | base → base· | 不动 ✓ | 被压缩 ✓ | 不微调也能用一部分 |
| YaRN | 按圈数分段 + 温度 | 不动 ✓ | 被压缩 ✓ | 少量微调,效果最好 |
3.6 ALiBi:换一条路
Press et al.(2022)的论文标题就是答案:Train Short, Test Long。不用任何位置编码,直接往打分上加一个线性惩罚:
每个头一个固定斜率 ,取几何级数: 时是 (第 9 组)。斜率大的头只看最近几个词,斜率小的头看得很远。
为什么外推好? 没有任何依赖绝对位置的东西;更远的 key 只是被扣更多的分,不存在”没见过的输入”。代价:很强的近因偏好,远处的信息天然被压制。BLOOM、MPT 用了 ALiBi;但近两年的主流大模型大多回到了 RoPE + 扩展方法。
3.7 “能算”不等于”能用”
即使位置编码完美外推,模型在长上下文里用好信息的能力也是另一回事。Lecture 3 会讲”lost in the middle”:信息放在长上下文的中间时,模型最容易忽略它。上下文窗口的名义长度、困惑度不崩的长度、真正能利用信息的长度,是三个不同的数。
§4 现代版 mini-GPT 与外推实验
4.1 现代版的配方
把本周所有零件换成现代选择:
| 零件 | Day 3 的版本 | 现代版 | 出处 |
|---|---|---|---|
| 位置编码 | 可学习绝对位置 | RoPE | 今天 §2 |
| 注意力 | MHA | GQA( 或 ) | Day 4 §2 |
| 归一化 | LayerNorm | RMSNorm | Day 2 §2.5 |
| FFN | GELU, | SwiGLU, | Day 2 §1.4 |
| bias | 有 | 无 | Llama 系 |
| 推理 | 朴素生成 | KV Cache,且超过训练长度也不用重建 | Day 4、今天 §2.6 |
这就是 Llama 系列的块结构。参考代码里 modern_config() 一行给出这套配置;--train-modern 用与 Day 3 相同的配方训练,便于对照。
4.2 外推实验
学习计划的动手任务:“把正弦位置编码换成 RoPE,观察在超出训练长度的序列上两者困惑度的差异。” 参考代码 --extrapolate 把它扩展成一个完整的对照:
- 用相同的配方、训练长度 128,训练四个模型:learned、sin、RoPE、ALiBi。
- 在长度 64、128、256、512 上评估验证 loss,并单独统计超出训练长度那部分位置的 loss。
- 对训好的 RoPE 模型,不微调,分别换上 PI、NTK、YaRN 的角频率再评估。
先写下你的预测,再看结果。可以参考的公开结论(不是本仓库实测):ALiBi 论文报告正弦编码和 RoPE 在超出训练长度后困惑度迅速上升,ALiBi 平稳;NTK 类方法不微调就能缓解一部分;PI 不微调时效果有限。字符级、小模型、短上下文的设置和论文差别很大,结论不一定成立,这正是值得亲手跑的原因。跑完请把表发回来。
§A RoPE 与长上下文:API 用法与边界
A.1 HuggingFace 的配置字段
config.rope_theta # base;Llama 2 是 10000,Llama 3 是 500000
config.max_position_embeddings # 名义上下文长度(RoPE 本身没有表,这只是"训练时的长度")
config.rope_scaling = {"rope_type": "yarn", "factor": 4.0,
"original_max_position_embeddings": 4096} # 也可以是 "linear"(PI)、"dynamic"(动态 NTK)、"llama3"
- 旧版本里这个键叫
"type",新版本叫"rope_type"。 - 改了
rope_scaling后,max_position_embeddings通常也要相应改大,否则有些代码路径会按旧长度截断或报警。 - 改 rope 配置等于改了模型:同一份权重换一种 scaling 就是另一个模型,评估结果要标清楚。
A.2 配对约定与权重转换
§2.5 的两种约定。自己实现 RoPE 去加载 HuggingFace 的 Llama 权重,要用 rotate_half(两半配对);加载 Meta 原始 checkpoint 要用相邻配对,或者先做 HF 转换脚本里那个行置换。验证方法:拿官方实现跑一次前向,与你的实现逐位比对 logits。
A.3 精度:角度必须用 fp32 算
freqs = 1.0 / (base ** (torch.arange(0, d, 2, dtype=torch.float32) / d))
t = torch.arange(max_len, dtype=torch.float32)
ang = torch.outer(t, freqs) # fp32
cos, sin = ang.cos().to(dtype), ang.sin().to(dtype) # 最后再转成模型的 dtype
w1day5_checks.py 第 11 组:bf16 只有 8 位有效二进制数字,连位置整数本身都存不准,1000 附近间隔为 4,30000 附近间隔为 128。用 bf16 算角度,位置 1001 的误差 1.17 rad,位置 4097 的误差 4.14 rad,已经超过 ,cos/sin 完全不对。这类错误不报错,只让长序列上的效果莫名变差。 Meta 的原始实现用 torch.polar 和复数,复数不支持半精度,所以它会先把 转成 fp32 再旋转,也是出于同样的原因。
A.4 左补齐时的位置编号
批量生成要左补齐(Day 1 §3.4)。如果位置编号简单地用 arange(T),真实内容的第一个 token 会被当成在位置 2、3……而不是 0。绝对位置编码下这会改变输出;RoPE 下打分只看相对距离,同一个样本内部的相对距离不受影响,但 pad 也占了编号,与训练时的分布略有不同。HuggingFace 生成时按 attention_mask.cumsum(-1) - 1 计算 position_ids,让每个样本的真实内容都从 0 开始编号。自己写批量生成时照做。
A.5 SDPA 与 ALiBi
ALiBi 要把一个 的浮点偏置作为 attn_mask 传给 SDPA。传了浮点掩码后,部分融合内核(尤其是 FlashAttention 后端)不能用,会回退到其他后端,速度变慢(Day 1 §A.1 边界 5)。这是加法偏置路线的一个实际代价,也解释了 §1.2 说的”RoPE 与任何内核兼容”为什么重要。
§B 实操
- 跑通
w1day5_checks.py,11 组结论都看懂。第 1 组对照 Day 2 第 7 组;第 5 组是理解外推失败的关键。 - 推导(1h,纸笔):(a) 用复数推出二维 RoPE 的相对性;(b) 推出 ;(c) 推出 NTK-aware 的 base’ 公式;(d) 对 、base 10000、训练长度 4096,算出哪几对维度的波长超过训练长度。
- 实现 RoPE(1.5h):两种约定都写;验证相对性和置换等价;接到 Day 4 的注意力上(只覆盖位置钩子);确认增量解码可以超过训练长度一直进行而不重建缓存,且与全量前向一致。
- 现代版 mini-GPT(1h):RoPE + GQA + RMSNorm + SwiGLU + 无 bias。与 Day 3 的 GPT 同配方训练,比较验证 loss 和参数量。
- 外推实验(2h):
--extrapolate。先写预测,再填表,再解释:哪种位置编码外推得最好?RoPE 的三种免训练扩展谁更好?和论文的结论一致吗?不一致的话,可能的原因是什么?
参考实现:code/w1d5_rope.py。
§C 自测
动机与直觉层
- 为什么模型真正需要的是相对位置?绝对位置编码在这件事上的代价是什么?
- 加法偏置和乘法旋转两条路线各有什么优劣?
- RoPE 和正弦编码的频率相同,区别在哪里?为什么这个区别决定了相对性能否保住?
- 为什么 RoPE 只转 和 ,不转 ?
- RoPE 能算任意位置,为什么外推还会失败?具体是哪些维度出了问题?
- PI 的代价是什么?NTK-aware 为什么能避免这个代价?
- ALiBi 为什么外推好?它的代价是什么?
推导层
- 用复数推出二维 RoPE:为什么乘 能让内积只依赖 ?
- 写出 ,并由此写出 RoPE 的打分公式。
- 推出 NTK-aware 的 。为什么最高频不受影响?
- YaRN 的分段依据是什么?温度系数的作用是什么?
- 解释为什么 RoPE 下缓存的 K 在窗口滑动后依然有效。
工程与 API 层
rotate_half与相邻配对有什么区别?用错了会怎样?怎么检查?- 为什么 cos/sin 表必须用 fp32 算?bf16 在位置 4097 附近的误差有多大?
- HuggingFace 的
rope_scaling有哪几种类型,各对应今天的哪种方法? - 左补齐时位置编号应该怎么算?
- 为什么 ALiBi 会让 SDPA 退出 FlashAttention 后端?
出关标准
- 能从”内积只依赖相对位置”推出 RoPE(复数二维 → 块对角高维)
- 能说清 RoPE 与正弦编码的关系:同样的频率,旋转放在投影之后
- 能列出 RoPE 的六条性质,特别是”缓存友好”
- 能指出 RoPE 外推失败的具体维度(波长 > 训练长度的低频维度)
- 能推出 NTK-aware 的 base 公式,说清 PI / NTK / YaRN 各改了什么
- 能说出 ALiBi 的做法、外推优势和代价
- 实现 RoPE,接进带缓存的 GQA 注意力,超过训练长度仍与全量前向一致
- 训练现代版 mini-GPT,完成外推实验并解释结果