位置编码¶
没有位置机制时,自注意力对输入 token 的排列具有置换等变性。位置编码向模型注入顺序、距离或多维坐标,使相同内容在不同位置产生不同交互。
位置机制需要回答四个问题:
- 位置进入 hidden state,还是只进入 attention score;
- 使用绝对坐标还是相对距离;
- prefill、增量 decode 和 cache 怎样保持同一坐标系;
- 超出训练范围时,频率或偏置如何变化。
绝对位置¶
最直接的做法是把 token embedding 与位置向量相加:
Learned absolute embedding 令 \(p_t\) 为参数表,接口简单但天然受表长限制。原始 Transformer 使用正弦位置:
正弦函数可计算任意位置,但“能计算”不等于模型在训练范围外仍会正确使用。
RoPE¶
Rotary Position Embedding 把每两个通道视为二维平面,并对位置 \(m\) 应用旋转:
对 query 和 key 使用
由于
点积自然依赖相对位移:
最小语义实现¶
apply_rope 接收 [T,D] 张量和每个 token 的逻辑 position ID,按相邻偶/奇通道配对后返回同 shape 张量。两个断言分别验证旋转保持向量范数,以及 query/key 同时平移时点积只依赖相对位移。
import torch
def apply_rope(x, position, base=10_000.):
assert x.ndim == 2 and x.shape[-1] % 2 == 0
half = x.shape[-1] // 2
index = torch.arange(half, device=x.device)
inv_freq = base ** (-index / half)
angle = torch.as_tensor(position, device=x.device)[:, None] * inv_freq[None, :]
even, odd = x[:, 0::2], x[:, 1::2]
cos, sin = angle.cos(), angle.sin()
return torch.stack(
(even * cos - odd * sin, even * sin + odd * cos), dim=-1
).flatten(-2)
torch.manual_seed(0)
q, k = torch.randn(3, 8), torch.randn(3, 8)
position = torch.arange(3)
torch.testing.assert_close(apply_rope(q, position).norm(dim=-1), q.norm(dim=-1))
lhs = (apply_rope(q[:1], [7]) * apply_rope(k[:1], [11])).sum()
rhs = (apply_rope(q[:1], [0]) * apply_rope(k[:1], [4])).sum()
torch.testing.assert_close(lhs, rhs)
真实模型还要固定 split-half/interleaved 约定、partial rotary dimension、batch broadcast、角度计算精度与 cache 中 K 的存储形态;仅凭公式无法保证 checkpoint 对齐。支持高维张量的实现见张量原语:Rotary Position Embedding,prefill/decode 的 position 与 mask 应同 Decoder-only Transformer:Attention 联合测试。
不同通道使用不同角频率。常见定义为
其中 \(d_r\) 是参与旋转的维度,\(b\) 是 frequency base。实现必须保存 \(d_r\)、base、缩放规则和 position ID;仅保存“最大长度”不足以复现。
增量解码¶
若 cache 已含位置 \(0,\ldots,T-1\) 的 key,新 token 必须使用位置 \(T\)。批处理中的 padding、prefix cache、speculative branch 和截断窗口都会使“数组下标”不再等于“逻辑位置”。
需要明确 cache 存储:
- 未旋转 K:读取时按历史位置旋转,灵活但增加计算;
- 已旋转 K:decode 快,但 cache 复用要求位置语义完全一致。
ALiBi¶
ALiBi 不修改 hidden state,而在每个 head 的 score 上加入距离偏置:
不同 head 使用不同斜率 \(m_h\),形成多个距离尺度。它不需要位置表,也可计算训练范围外的距离;实际外推质量仍取决于训练分布、任务和注意力模式。
相对位置偏置¶
另一类方法把相对距离映射为 bucket:
短距离可用细粒度 bucket,长距离用对数 bucket。其优点是直接控制相对距离,代价是 bucket 定义、方向性和跨长度行为都成为模型接口。
多维位置¶
图像和视频 token 不只有一维序号。一个视频 patch 可以具有
三个坐标。Qwen2-VL 的 M-RoPE 将旋转维度分配给时间、高度和宽度;文本片段则让这些坐标共同前进。
实现要明确:
- 每个轴占多少 rotary dimension;
- dynamic tiling 后 patch 的原始二维坐标;
- 多图之间是否重置坐标;
- padding、缩放和 frame sampling 是否同步更新位置;
- 文本、图像和视频交错时逻辑时间如何推进。
把二维或三维网格直接 flatten 成一维索引虽然简单,却让相邻关系依赖具体扫描顺序。
长度扩展¶
设训练长度为 \(L_0\)、目标长度为 \(L_1\)。Position Interpolation 把目标位置压回训练范围:
YaRN 对不同频率采取更细粒度的插值与修正;LongRoPE 搜索非均匀缩放并处理短上下文恢复。它们改变的是位置频谱,不会自动解决:
- full attention 的二次计算;
- KV Cache 容量;
- 长距离训练数据不足;
- 模型在中部长证据上的利用能力。
这些系统与评测问题见长上下文。
Shape 与实现契约¶
设
若 rotary dimension 为 \(d_r\le d_h\),则只旋转 [..., :d_r],其余通道原样保留。实现应固定:
- interleaved 还是 split-half 配对;
- 正余弦缓存的 shape 与 broadcast 轴;
- position IDs 是 \([T]\)、\([B,T]\) 还是多维坐标;
- Q 与 K 是否使用相同 rotary dimension;
- dtype:角度计算与最终张量是否同精度;
- cache 保存旋转前还是旋转后 K;
- 长度扩展配置是否进入 checkpoint metadata。
失效模式¶
- Off-by-one:prefill 最后位置与首个 decode 位置重复。
- Padding 漂移:用物理下标而不是每个样本的有效位置。
- 配对约定不一致:训练与推理分别使用 interleaved、split-half。
- Cache 复用错误:已旋转 K 被放到不同绝对位置复用。
- 低精度角度误差:极长位置下相位计算失真。
- 只改配置长度:模型能分配 cache,却未学会范围外位置。
- 短上下文回退:扩展频率破坏原训练区间性能。
- 多维坐标丢失:动态切片后仍按简单一维顺序编码。
验证¶
| 测试 | 判据 |
|---|---|
| 相对性 | 同时平移 Q/K 位置后点积保持 |
| Prefill–decode | 全序列计算与逐 token cache 结果对齐 |
| Batch padding | 左右 padding 不改变有效 token 结果 |
| 配对约定 | 与 checkpoint 官方实现逐元素对齐 |
| 长度边界 | \(L_0-1,L_0,L_0+1,L_1-1\) 均无突变 |
| 位置扫描 | 相同证据在头、中、尾的性能曲线 |
| 短上下文 | 扩展前后原训练长度内无不可接受回退 |
| 多模态 | crop、tile、frame 顺序变化后坐标仍一致 |
RoPE 与 GQA 的组合实现见注意力家族,可执行的最小序列实验见序列模型手撕实现。