AI 技术AI 推理 · 6/10
#RoPE#位置编码#长上下文#YaRN#LongRoPE#Transformer

AI 推理系列(六):RoPE 如何编码 Token 距离:从旋转位置到长上下文扩展

从 Self-Attention 不具备天然顺序感这一问题出发,解释 RoPE 如何通过旋转 Query 和 Key 注入位置关系,并梳理 Position Interpolation、YaRN、LongRoPE 等上下文扩展方法的机制、工程边界与长文本系统中的真实瓶颈。

把一句“猫吃鱼”改成“鱼吃猫”,词表没有变化,语义却完全不同。Transformer 如果只拿到三个 Token 的内容表示,并没有一个天然的“第一个、第二个、第三个”坐标轴。Self-Attention 擅长比较表示之间的相关性,却需要额外机制告诉模型这些表示出现在什么位置,以及两个位置之间相隔多远。

RoPE,也就是 Rotary Position Embedding,解决的正是这个坐标问题。它不把一个独立的位置向量直接加到 Token embedding 上,而是在 Attention 计算前,根据位置对 Query 和 Key 的部分维度做旋转。这样一来,位置不再只是输入表示上的一个附加标签,而会直接进入 Query 与 Key 的匹配关系。后来出现的 Position Interpolation、YaRN、LongRoPE 等长上下文方法,大多也是围绕这套旋转频率和位置映射继续做文章。

Attention 为什么必须额外获得位置

Self-Attention 的基本操作是把每个 Token 映射成 Query、Key 和 Value。某个 Query 与所有 Key 计算匹配分数,再用这些分数汇总对应的 Value。这个过程很适合回答“当前 Token 应该关注哪些其他 Token”,但如果没有任何位置相关信息,模型无法仅凭 Attention 运算区分同一组 Token 的不同排列。

例如句子里有两个相同词,内容 embedding 可以完全一致,但一个位于句首,另一个位于句尾。对于语法关系、指代、局部搭配和长距离依赖来说,“出现了什么”与“出现在哪里”经常缺一不可。位置编码因此承担两类信息:一个位置本身在哪里,以及两个位置之间是什么相对关系。

早期 Transformer 可以把一个位置向量直接加入 Token embedding。这样做很直观:第 10 个 Token 得到第 10 个位置表示,第 11 个 Token 得到另一个表示。但位置一旦直接混进输入向量,后续 Attention 需要自己从表示里重新学习如何把两个绝对位置组合成相对距离。RoPE 采用了另一条路径:先产生普通 Query 和 Key,再依据各自位置旋转它们,使相对位置信息自然进入内积。

这里需要避免一个常见误解。RoPE 不是给模型一个单独的“距离 = 5”数值,也不是提前建立一张距离查找表。它通过不同位置对应的旋转相位,让两个旋转后向量的匹配结果随相对位置变化。模型最终学到的仍然是 Attention 权重,只是这些权重现在具有可利用的位置结构。

RoPE 如何把位置变成旋转关系

理解 RoPE 可以先只看二维平面。假设 Query 的两个维度构成一个二维向量。位置为 0 时不旋转,位置为 1 时旋转一个角度,位置为 2 时旋转两倍角度。Key 使用同样的位置规则。于是,两个位置的 Query 和 Key 做内积时,结果不仅受原始内容向量影响,也受两次旋转之间的角度差影响。

真实模型的 head dimension 往往远大于 2。RoPE 会把维度两两分组,每一对维度看成一个独立二维平面,不同维度对使用不同旋转频率。某些维度随位置变化得更快,某些变化得更慢。这样,同一个位置会同时作用于多种尺度的位置模式。

可以把一层 Attention 中的位置处理抽象成下面的数据流:

flowchart TD
  A[Token 隐状态] --> B[线性投影 Q K V]
  B --> C[读取 Token 位置]
  C --> D[按频率旋转 Q]
  C --> E[按频率旋转 K]
  D --> F[旋转后的 Q]
  E --> G[旋转后的 K]
  F --> H[计算 Attention 分数]
  G --> H
  B --> I[V 保持内容表示]
  H --> J[Softmax 权重]
  J --> K[加权汇总 V]
  I --> K
  K --> L[Attention 输出]

这个流程里有一个重要边界:RoPE 主要作用于 Query 和 Key,而不是通过改写 Value 来保存位置信息。位置影响“谁应该和谁匹配”,随后 Attention 权重决定从哪些 Value 读取内容。也正因为旋转直接发生在匹配空间里,相对位置信息能够和 Attention 内积紧密结合。

RoFormer 提出的关键性质是,旋转后的 Query 与 Key 内积能够同时依赖内容和相对位置。工程上最值得记住的不是具体三角函数,而是两个状态:每个 Token 先拥有自己的绝对位置相位;两个 Token 进行匹配时,起作用的是它们相位之间的关系。

为什么训练长度会成为 RoPE 的边界

假设一个模型只在 4096 Token 以内完成训练。训练过程中,它只反复看到这段范围内的旋转相位组合和相对距离模式。如果部署时直接把位置扩展到数万甚至数十万,新的相位关系可能落到训练分布之外。数学公式当然仍然可以计算,但“公式能算”不等于“模型学会了如何使用这些新位置模式”。

这就是长上下文扩展里的核心矛盾。模型配置可以允许更大的 position index,推理代码也可以成功执行,却可能在超出原始训练范围后出现检索、复制、位置判断或长距离组合能力退化。问题不是整数位置放不下,而是模型的参数没有在这些位置模式上形成可靠行为。

对于在线长文档问答,这种退化很容易被“请求没有报错”掩盖。一个 128K Prompt 能完整进入模型,并不表示尾部问题可以稳定调用开头证据。系统层看到的是输入成功、显存正常、Attention kernel 正常;任务层却可能发现远距离证据命中率下降。这也是为什么评估长上下文时不能只检查最大允许长度。

长度外推还会和频率结构发生关系。RoPE 不同维度使用不同频率,有些维度在原训练窗口内已经经历较多旋转,有些只变化很小。把位置简单扩大几十倍,意味着不同频率进入新区域的程度并不一致。于是“所有维度统一按一个公式继续转”虽然简单,却未必是最稳妥的外推方式。

Position Interpolation 为什么要把长位置压回训练范围

Position Interpolation 的直觉是:既然模型熟悉的是原始位置区间,就不要直接要求它处理远超训练范围的新坐标,而是把更长的实际位置映射回原有区间。假设训练窗口是 4096,现在希望处理 16384,可以把实际位置按比例缩小,使末端位置在 RoPE 看来仍接近 4096。

这种方法的好处是明显的。旋转相位不会立刻跑到原训练分布之外,模型面对的是被压缩的位置坐标,而不是完全陌生的远端位置。但代价也同样明确:所有位置间距都同时缩小。原来相隔 4 个 Token 的局部关系和相隔 4000 个 Token 的长距离关系,都会发生比例变化。

这会带来一个工程判断:上下文扩展不能只看“最大位置是否映射回来了”,还要看原模型最依赖的局部位置结构有没有被压得过头。语言模型大量模式来自相邻或近邻 Token。如果为了覆盖极长距离,让局部位置分辨率明显变化,模型可能在长文档之外先损伤短文本能力。

因此,更成熟的 RoPE scaling 方法通常不会把所有频率视为完全相同的资源。不同频率对局部和长距离模式承担的作用不同,缩放策略也需要考虑这种差异。这正是 YaRN 等方法继续改进的位置。

YaRN 为什么要区分不同频率的缩放

RoPE 可以被看成一组不同频率的旋转器。高频维度在位置稍微变化时,相位就改变较多;低频维度变化更慢,可以跨越更长范围。统一缩放虽然实现简单,却可能让原本负责局部区分的高频部分和负责长距离关系的低频部分一起被扭曲。

YaRN 的思路不是把所有频率机械地乘上同一个比例,而是针对不同频率区域采用更细致的调整,并在原始行为与长距离扩展之间进行平滑过渡。用户提供的 YaRN 论文材料报告了其在长上下文扩展训练效率上的优势,并展示了扩展到 32K、64K 和 128K 的实验。这些数字属于论文具体训练和模型配置,不能直接推导出任意 RoPE 模型都能用同样成本获得相同窗口。

从系统视角看,YaRN 改变的仍然是“位置如何映射到旋转角度”。它不会自动减少 Attention 的计算量,也不会让 KV Cache 变小。模型从 4K 扩展到 128K 后,即便位置编码行为已经稳定,128K Token 仍会产生更大的 KV Cache、更长的 prefill 和更高的 Attention 成本。

因此,在一套长上下文服务中,RoPE scaling 与推理优化属于不同层。前者负责让模型的位置坐标在更长区间里尽量可用;后者负责让这么多 Token 真正能在硬件上被处理。两者缺一不可,但也不能互相替代。

LongRoPE 把问题推进到逐维度搜索

随着目标窗口继续扩大,统一规则越来越难覆盖所有频率和位置区域。LongRoPE 进一步把缩放策略细化,搜索不同维度更合适的位置缩放关系,并结合位置插值等方法扩展上下文。用户提供的 LongRoPE 材料报告了在 LLaMA2、Mistral 等模型上的超长上下文实验,其中包含超过两百万 Token 的配置探索。

这类结果首先说明一个事实:RoPE 的位置映射存在很大的可调空间,原训练窗口并不是一个由公式强制限定的绝对硬上限。但它也暴露出另一个现实:上下文扩得越远,越不能只靠一个缩放倍数判断质量。搜索策略、微调数据、训练阶段、原模型频率设置以及评估任务都会影响最终结果。

LongRoPE2 等后续工作继续关注扩展后短上下文能力的保持问题。这一点非常重要。一个模型如果把最大上下文从 32K 扩到更长,却让常用的 2K、8K 请求质量下降,那么生产价值未必增加。线上流量通常是长短请求混合分布,而不是每个请求都接近最大窗口。

从工程验证角度看,RoPE 扩展至少要建立两组基准。一组覆盖原训练长度以内的常规任务,用来检查回归;另一组覆盖目标长窗口中的检索、复制、跨段组合与推理任务,用来确认扩展确实带来有效能力。只测试“能够成功 prefill 到目标长度”远远不够。

支持百万 Token 与有效利用百万 Token 是两件事

“最大上下文长度”是一个很容易传播的模型指标,因为它是单个数字。但这个数字实际混合了模型配置、位置编码、训练数据、推理引擎和硬件容量等多个条件。一个系统允许输入一百万 Token,只说明这条执行链可以接收这么长的序列;它没有单独证明模型能稳定使用任意位置的信息。

可以把长上下文能力拆成三个不同问题:

层面核心问题典型失败表现需要的验证
输入容量模型和推理系统能否接收目标长度OOM、位置越界、kernel 不支持最大长度压力测试、显存与 prefill 监控
信息定位远距离关键信息能否被准确找到信息存在但回答忽略、位置越远命中越差Needle 类测试、不同深度检索测试
组合推理分散在长文档中的证据能否联合使用单点检索成功但多跳推理失败LongBench、InfiniteBench 类任务和业务集
短上下文保持扩展后原有能力是否退化常规短请求质量下降原始长度回归集、A/B 评估

这也是为什么模型宣传的窗口长度不能替代业务评估。对于代码库问答,真正重要的不是一次能塞进多少文件,而是问题涉及两个相距很远的模块时,模型能否找到两处实现并正确建立依赖关系。对于企业文档问答,则需要检查证据埋在不同位置时的召回和引用质量。

用户材料中提到了 Llama 3.1 等公开模型的长上下文配置,也提到了更激进的百万级窗口探索。使用这些数字时,应始终把“该模型或论文报告的支持长度”与“在某个业务任务上的可靠有效长度”分开描述。前者可以从模型或论文材料确认,后者必须由真实评估得出。

RoPE 扩展后,系统瓶颈会转移到 KV Cache 和 prefill

假设位置编码已经能够稳定覆盖 128K,推理系统仍然需要为 128K Token 计算并保存 Attention 状态。上下文变长后,KV Cache 会随着历史 Token 数增加;prefill 需要处理更长输入;标准 Attention 的计算规模也随序列增长迅速增加。RoPE scaling 没有改变这些资源规律。

这正好解释了 GQA、KV Cache Quantization、PagedAttention 和 FlashAttention 为什么经常和长上下文一起出现。它们各自处理不同层面的瓶颈:

技术所在层主要解决的问题不能替代什么
RoPE Scaling / YaRN位置表示让更远位置的旋转关系尽量保持可用不降低 KV Cache 或 Attention 成本
GQA模型结构减少 KV Head 数量,从源头降低缓存规模不管理缓存碎片,也不扩展位置编码
KV Cache Quantization数据表示降低每个缓存元素的存储成本不解决位置外推质量
PagedAttention内存管理更灵活地分配和回收 KV Cache不减少 Attention 算术量
FlashAttentionKernel / IO减少 Attention 中间数据搬运不保证长距离信息能被模型利用

对于一个长文档在线服务,可以把完整请求看成连续状态变化。请求进入后先执行 tokenizer 和位置编号;RoPE scaling 决定这些位置如何映射到 Q/K 旋转;prefill 计算整个 Prompt 并建立 KV Cache;随后 decode 每生成一个 Token 都继续追加缓存。如果位置编码已经扩展成功但 TTFT 很高,问题可能在 prefill;如果显存先耗尽,应该检查 GQA、量化和缓存管理;如果输入能跑完但远端证据持续丢失,则应回到模型长上下文能力和 RoPE 扩展质量本身。

生产环境因此需要把指标拆开。位置扩展相关的质量指标包括不同位置深度的检索准确率、短上下文回归、长距离组合任务成功率;系统指标则包括 Prompt 长度分布、TTFT、prefill tokens/s、KV Cache 使用量、显存峰值、decode 延迟和并发容量。只有两组指标同时成立,最大上下文窗口才真正具有业务价值。

什么时候不应该把问题归因于 RoPE

长上下文回答失败时,RoPE 很容易成为第一个被怀疑的组件,但失败原因可能完全在其他层。比如模型能准确复述位于 100K 位置的一段文本,却无法根据这段文本完成复杂推理,此时位置表示未必是主要问题,更可能是训练数据或推理能力不足。

另一类情况是服务根本无法完成目标长度 prefill。如果 GPU 在 80K Token 就 OOM,或者 Attention backend 不支持当前形状,那么即使 RoPE 数学上已经扩展到更远位置也没有意义。这属于容量和执行路径问题,应先处理缓存、并行和 kernel。

还可能出现 tokenizer 或数据预处理导致的误判。业务方认为文档只有几十万“字”,实际转换成 Token 后长度完全不同。不同语言、代码和结构化文本的 tokenization 密度差异明显。评估长上下文系统时应使用真实 Token 数,而不是文件大小、字符数或页数来推断模型窗口占用。

位置扩展本身也存在回退条件。如果一种 scaling 策略在目标长窗口上提高了远距离能力,却让原始短上下文任务出现明显回归,生产系统未必应该直接全量启用。更合理的做法可能是重新训练、调整缩放参数或为不同模型版本维护清晰的上下文能力边界,而不是把一个更大的 max_position_embeddings 当成完成标志。

如何验证一次 RoPE 长上下文扩展是否真的成功

最小可行验证可以从“位置、质量、系统”三条线同时展开。位置线确认目标长度没有越界或数值异常,并覆盖窗口不同深度;质量线比较扩展前后的短文本能力和扩展后的远距离检索、跨段推理;系统线观察显存、prefill、Attention kernel 和 KV Cache 是否在目标负载下可接受。

测试样本不要只把答案固定放在文档末尾。更可靠的方法是把关键信息分别埋在开头、四分之一、中央、四分之三和末端,观察距离变化是否形成系统性退化。对于多跳问题,还需要把两个必要证据放在相距较远的位置,避免模型只靠单点检索通过测试。

长短请求混合测试也不可省略。多数生产服务不会连续收到接近最大窗口的请求。如果扩展后的模型为了处理极长序列改变了 kernel、缓存布局或调度策略,可能影响普通请求延迟。应在真实流量分布下同时测量短、中、长 Prompt,而不是只跑最大长度压力测试。

最后才是最大窗口数字。一个可靠的结论应该类似“在目标硬件上可以稳定处理某长度,短上下文回归在接受范围内,远距离检索和组合任务通过业务阈值”,而不是单独说“支持某个 Token 数”。RoPE 解决了 Transformer 的位置坐标问题,RoPE scaling 扩大了这套坐标系;真正把坐标系变成可用长上下文能力,还需要训练、Attention、KV Cache、调度和评估一起成立。

资料来源

  1. RoFormer: Enhanced Transformer with Rotary Position Embedding
  2. YaRN: Efficient Context Window Extension of Large Language Models
  3. LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens
  4. LongRoPE2: Near-Lossless LLM Context Window Scaling
  5. Introducing Llama 3.1
  6. MrRoPE: Mixed-radix Rotary Position Embedding
  7. Rotary Position Embedding for Vision Transformer