在 Transformer 架构中,注意力机制通过计算 Query 与 Key 的点积来确定 token 之间的关联强度。然而,原始的注意力计算完全不包含位置信息 —— 点积结果与 token 在序列中的相对或绝对位置无关。由于自然语言的高度位置相关性,位置编码(Positional Encoding)成为 Transformer 实现序列建模的关键组件。Jane Street 的最新研究从群论角度对位置编码的设计空间进行了系统性分析,发现看似多样的编码方法实际上都隶属于一个统一的数学结构:一位参数群(One-Parameter Group)。这一发现不仅深化了对现有编码方案的理解,更为工程实践提供了明确的选型依据。
位置编码的数学形式化
在探讨位置编码的理论基础之前,首先需要对其数学表达进行形式化定义。设给定时间序列的 Query 向量为 $q (t)$,Key 向量为 $k (t)$,其中 $t$ 表示时间索引,可为整数或连续值。原始注意力计算为 $q (t)^\top k (s)$,即两个向量的内积,它无法区分 $t$ 与 $s$ 之间的相对关系。位置编码的目标正是要对这一内积运算进行修改,使其依赖于 $t$ 和 $s$ 的相对位置。
Jane Street 的分析提出了三个关键约束条件来筛选 “合理” 的位置编码函数。第一是线性性(Linearity):位置编码对 Query 和 Key 的变换必须是线性的,即存在矩阵函数 $F (t)$ 和 $G (t)$ 使得 $q_{enc}(t) = F (t) q (t)$、$k_{enc}(s) = G (s) k (s)$。这一假设虽非严格不可缺,但在向量空间框架下是自然的简化。第二是平移不变性(Translation Invariance):编码效果只应依赖于相对位置 $d = t - s$,而非绝对索引本身。这一性质对模型泛化至更长序列至关重要 —— 若编码暴露了绝对索引,模型在训练时见过最长序列为 $N$,当推理时遇到超过 $N$ 的序列将面临分布外问题。第三是连续性(Continuity):编码函数 $A (d)$ 必须是连续的,这是保证模型训练稳定性的工程实践共识。
在这三个约束下,Attention 计算可改写为 $q (t)^\top A (t-s) k (s)$,其中 $A (d)$ 是描述相对位置 $d$ 如何修正内积的矩阵函数。关键洞察在于:$A (d)$ 必须满足 $A (d_1 + d_2) = A (d_1) A (d_2)$,这正是群论中一位参数群的定义条件。因此,所有满足上述约束的位置编码都唯一地对应于某个矩阵指数形式 $A (d) = \exp (d \cdot M)$,其中 $M$ 为生成元矩阵(Generator Matrix)。
对角化生成元的分类讨论
对生成元矩阵 $M$ 的结构进行分析,可将位置编码方案划分为两大类别。第一类是对角化可实现的情形,即 $M$ 可在某个正交基下分解为对角块。对角化后,一维子空间上的作用退化为标量乘法 $e^{\lambda d}$,其中 $\lambda$ 为实数特征值;二阶子空间(对应共轭特征值对)则表现为旋转与缩放的组合,可表示为 $e^{\alpha d} \cdot R (\beta d)$,其中 $R (\theta)$ 为二维旋转矩阵。
针对一维子空间的分析揭示了三种可能情况。当 $\lambda > 0$ 时,$e^{\lambda d}$ 随时序距离增加呈指数增长,这会导致注意力分数随距离增加而放大,显然不符合建模直觉,应被排除。当 $\lambda = 0$ 时,$e^{0 \cdot d}$ 恒等于单位矩阵,此时相当于不使用任何位置编码(NoPE),注意力完全退化为内容导向。当 $\lambda < 0$ 时,$e^{\lambda d}$ 实现指数衰减,这一机制正是线性注意力(Linear Attention)变体的核心数学基础,在 Gated Linear Attention 等架构中被显式利用。
二阶子空间产生的旋转结构尤其值得关注。通过基变换可将共轭特征值对转化为标准的二维旋转矩阵形式,其数学表达为 $R (\omega d) = \begin {pmatrix} \cos (\omega d) & -\sin (\omega d) \ \sin (\omega d) & \cos (\omega d) \end {pmatrix}$。这正是旋转位置编码(RoPE,Rotary Position Embedding)的核心机制 —— 对 Query 和 Key 的特定维度进行旋转,旋转角度与位置成线性关系。若同时引入指数衰减因子 $e^{-\alpha d}$,则得到带阻尼的 RoPE 形式,这一变体在 RetNet 和 Mamba-3 中被实际采用。参数 $\omega$ 控制旋转频率(对应位置编码的 “周期”),$\alpha$ 控制衰减速度,两者均为可学习的超参数。
缺陷矩阵与多项式编码
第二类别涉及生成元 $M$ 不可对角化(即亏损矩阵,Defective Matrix)的情形。当 $M$ 的 Jordan 标准形出现重复特征值对应的 Jordan 块时,指数作用将产生多项式因子而非纯粹的指数或三角函数。最简单的例子是 $M = \begin {pmatrix} 0 & 1 \ 0 & 0 \end {pmatrix}$,此时 $A (d) = \exp (dM) = \begin {pmatrix} 1 & d \ 0 & 1 \end {pmatrix}$,其中出现了线性增长的多项式项 $d$。
这一发现具有理论意义但工程价值有限。多项式位置编码的行为难以直观解释 —— 旋转编码对应 “时钟” 直觉,指数衰减对应 “记忆衰退” 直觉,而多项式修正缺乏明确的物理或语义对应。然而,Jane Street 的分析指出,ALiBi(Attention with Linear Biases)实际上可视为亏损矩阵的意外实现:ALiBi 对点积结果施加 $2^{-|t-s|}$ 的指数衰减,这一形式可通过特定的亏损生成元精确构造。这一联系揭示了 ALiBi 在群论框架下的数学本质,同时也说明看似 “非线性” 的 ALiBi 仍可纳入统一的线性编码体系(只需对 Query 和 Key 维度进行适当扩充)。
工程实践参数建议
基于上述理论分析,可为工程实践提供以下具体参数指导。在编码方案选型上,RoPE 是当前综合表现最优的方案,其数学形式(旋转)具有良好的理论基础,且与注意力机制兼容自然。实现时需关注基础频率 $\omega$ 的设置,典型值为 10000 至 100000,较大的值支持更细粒度的位置区分,但可能导致训练不稳定。RoPE 的最大序列长度由 $\omega$ 和向量维度共同决定:设维度为 $d$,第 $i$ 维(从 0 开始)的旋转周期为 $2\pi / (\omega \cdot 2^{i/d})$,实际部署时需确保推理长度不超过训练长度的 1.5 至 2 倍,否则需重新训练。
若场景对长程依赖有强需求(如文档摘要、代码生成),可考虑带衰减的 RoPE 变体,即在旋转基础上引入指数衰减因子 $e^{-\alpha d}$。衰减系数 $\alpha$ 的选取需在短程建模精度与长程信息保留之间取得平衡,建议范围为 0.01 至 0.1,可通过验证集 perplexity 进行网格搜索优化。对于超长上下文(超过 100k token)场景,建议采用 ALiBi 或其变体,因其显式的衰减机制在理论上有更好的长度外推性质,但需注意 ALiBi 无法通过标准因果掩码实现完整的相对位置建模,需在注意力掩码中额外引入偏置项。
统一视角的理论价值
群论框架的核心价值在于为纷繁复杂的位置编码方案提供了统一的分析语言。在此视角下,RoPE、指数衰减线性注意力、带阻尼的旋转编码不再是彼此独立的设计,而是同一位参数群在不同基下的表现形态。这一认识对模型架构搜索具有重要指导意义:既然满足约束的编码方案空间已被严格限定,工程师无需再耗费精力探索 “全新” 编码 —— 现有方案已覆盖所有数学上合法的可能性。未来的改进方向更应聚焦于编码参数的优化学习(如数据依赖的位置偏移)而非架构层面的根本性创新。
资料来源:本文核心事实与数学推导基于 Jane Street 博客文章《Using group theory to explore the space of positional encodings for attention》(2024),该文由 Jane Street ML 研究员 Alok 撰写。