位置编码是 Transformer 架构中最容易被忽视却至关重要的组件之一。2017 年原始论文引入的正弦位置编码已衍生出数十种变体,但这些方法之间是否存在深层联系长期缺乏系统性理论解释。Jane Street 研究员的最新工作从群论角度给出了令人惊讶的答案:在少数合理假设下,所有可行的位置编码方案都源自同一数学结构 —— 单参数群(one-parameter group)。这一发现不仅统一了 RoPE、指数衰减、ALiBi 等主流方法,更为工程实践提供了明确的参数选择边界。
形式化位置编码的数学约束
在无位置编码的原始注意力机制中,查询与键向量的点积计算不包含任何位置信息。为将位置信息注入注意力分数,研究者通常对查询和键向量进行与时间相关的变换。形式上,设查询向量和键向量分别为时间 t 的函数 F (t) 和 G (t),注意力分数为 F (t)ᵀG (s)。我们期望通过某种变换使该分数依赖于 t 和 s 的相对位置。
为推导出合理的位置编码形式,Jane Street 的分析引入了三个关键假设。首先是线性假设:位置编码在时间维度上必须是线性的,即 F (t)=F・A (t) 和 G (t)=G・B (t),其中 A (t) 和 B (t) 是方阵。这一假设虽然难以完全证明,但在神经网络框架下是自然的,因为嵌入空间本身是向量空间。其次是位移不变性:注意力分数仅依赖于相对位置 t-s,而不依赖于绝对时间。这意味着任意时间平移不应改变相对位置编码的性质。最后是连续性:A (t) 作为 t 的函数必须连续,以避免数值实现中的病态行为。
这三个假设联合推出了一个强约束:矩阵函数 A (t) 必须构成一个单参数群。更精确地说,存在固定生成元矩阵 G 使得 A (t)=exp (tG)。这正是群论介入的临界点 —— 单参数矩阵群的结构已被数学家彻底研究清楚,我们得以借此完整枚举所有可能的位置编码方案。
生成元矩阵的分类与编码方法对应
单参数群的分类取决于生成元矩阵 G 的谱特性。当 G 可对角化时,通过时间无关的基变换,我们可以在复数域将其化为对角形式。实特征值对应一维子空间上的作用,而共轭复特征值对则产生二维子空间。
实特征值 λ=0 的情况最为简单,此时 exp (tG) 是恒等矩阵,位置编码不产生任何效果,即 NoPE(无位置编码)方法。这种情况在某些需要位置不变性的任务中反而有益,例如某些图像分类任务或对称性本就明显的数据。
当 λ<0 时,exp (tG) 产生指数衰减因子。这正是线性注意力(Linear Attention)变体中常见的模式,其中注意力权重随相对距离增大而指数衰减。值得注意的是,在门控模型(如 Mamba 系列)中习得的与数据相关的衰减因子,从位置编码的视角看,可以理解为学习了一个非标准的时间推进速率,而非改变衰减率本身。当 λ 为纯虚数时,exp (tG) 对应旋转操作,这就是 RoPE(旋转位置编码)的数学本质。通过将查询和键向量在不同维度对上执行不同频率的旋转,RoPE 实现了相对位置编码的等变性。
更一般的情况是复数特征值带有负实部,这产生了旋转与指数衰减的组合,即阻尼 RoPE。RetNet 和 Mamba-3 等模型采用的正是这种位置编码:旋转提供周期性的位置信息,指数衰减则确保远距离 token 的影响被适当削弱。这种组合在处理超长上下文时尤为重要,因为纯粹的旋转编码在长距离上可能产生混淆。
缺陷矩阵与 ALiBi 的数学渊源
当生成元矩阵不可对角化(即为缺陷矩阵)时,情况变得更为复杂。数学上,这意味着 G 的 Jordan 标准形包含非平凡的 Jordan 块。此时 exp (tG) 不仅包含指数和三角函数,还会出现多项式因子。最简单的例子是 G=[[0,1],[0,0]],满足 G²=0,其矩阵指数为 I+tG,产生了线性增长项。
这一看似 pathological 的数学构造实际上与 ALiBi(注意力线性偏置)存在深层联系。ALiBi 通过对点积施加线性惩罚来实现位置偏置,其形式为 query・key - m・|i-j|。Jane Street 的研究指出,通过精心设计查询和键向量的维度扩展,可以使用缺陷生成元矩阵实现这一效果。具体而言,选择适当的二维向量使得它们的点积正比于时间差,即可在标准注意力框架内复现 ALiBi 的行为。这揭示了一个有趣的事实:ALiBi 虽然表面上不符合线性位置编码的形式,但实际上仍可纳入统一框架,只是需要扩展向量维度。
工程实践参数建议
基于上述理论分析,我们可以为不同场景下的位置编码选择提供具体指导。对于标准的自回归语言模型,RoPE 仍是最稳妥的选择,其基础频率参数 rope_theta 通常设置在 10000 至 100000 范围内,更大的模型或需要更长上下文时倾向使用更大值。例如 LLaMA 系列采用 10000,而某些长上下文变体会将 theta 提升至 500000 甚至更高。
当处理超长上下文(超过 8192 tokens)时,建议启用位置插值策略或采用阻尼 RoPE。位置插值通过缩放频率来适应更长的上下文长度,而阻尼 RoPE 通过引入指数衰减来自然地限制远距离注意力。在训练更大上下文模型时,将 max_position_embeddings 设置为目标长度的 1.5 至 2 倍可以为推理时的长度外推提供缓冲。
对于某些特定任务,如时间序列预测或需要明确建模位置距离衰减的场景,指数衰减位置编码值得考虑。这种情况下,衰减系数的选择应与数据的典型依赖范围匹配:对于金融时间序列,可能需要较快的衰减(衰减因子在 0.8-0.95 之间);而对于代码补全等任务,可以采用较慢的衰减以保留更长的依赖信息。
实现层面需注意,RoPE 要求隐藏维度为偶数以便进行维度配对旋转。若使用滑动窗口注意力,建议为窗口层和全局层设置不同的 rope_theta,因为二者需要建模的位置范围差异显著。全局层可以使用更大的 theta 以获得更精细的旋转粒度,而滑动窗口层则使用较小值以避免信息混淆。
统一视角的实践意义
群论框架的价值不仅在于理论统一,更在于它为位置编码的选择提供了决策边界。工程师无需在数十种方法中盲目试错,而可以依据任务特性在数学上等价的有限选项中做出有理有据的选择。位置编码的本质是告诉注意力机制如何理解 token 之间的距离,而群论揭示了这个「距离」必须满足的基本公理。一旦接受这一视角,RoPE、指数衰减、ALiBi 不再是彼此独立的设计发明,而是同一数学实在的不同投影。
资料来源:本文数学框架主要基于 Jane Street 博客《Using group theory to explore the space of positional encodings for attention》的理论分析,工程参数参考了主流开源实现的默认配置。