预训练模型中的位置编码：绝对位置、相对位置与旋转位置编码

点击“AladdinEdu你的AI学习实践工作坊”注册即送-H卡级别算力沉浸式云原生集成开发环境80G大显存多卡并行按量弹性计费教育用户更享超低价。1. 引言当自注意力忘记了顺序Transformer架构的核心是自注意力机制。其计算公式为[\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) \text{softmax}\left( \frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d_k}} \right) \mathbf{V}]仔细观察这一公式你会发现一个容易被忽略的事实它对输入序列的排列是等变的。如果将输入序列中的两个词交换位置自注意力输出的对应位置也会相应交换但每个位置的输出值本身不会因位置改变而变化。换句话说纯自注意力不包含任何关于token顺序的信息。对于语言而言顺序至关重要。“狗咬人”与“人咬狗”含义截然相反“我喜欢你”与“你喜欢我”情感关系完全不同。因此必须设法将位置信息注入到模型中。这一注入方式即为位置编码。从Transformer原论文的正弦位置编码到BERT的可学习绝对位置嵌入再到相对位置编码、旋转位置编码RoPE以及最新的ALiBi位置编码的设计经历了深刻的演进。不同的位置编码方案不仅影响模型对序列长度的泛化能力长度外推也关乎训练稳定性与计算效率。在LLaMA、GPT-NeoX、PaLM等大模型中RoPE已成为事实标准而在某些高效架构中ALiBi以其简洁有效崭露头角。本文将系统性地解析位置编码的技术脉络。我们将从绝对位置编码出发剖析其“为每个位置赋予唯一向量”的思想及局限性继而深入相对位置编码探讨如何在注意力分数中注入相对距离偏置最后重点阐释旋转位置编码的数学之美——如何用复数旋转将相对位置信息隐式编码进内积。文章将结合数学推导、代码实现与模型案例为你揭开位置编码的神秘面纱。2. 绝对位置编码为每个位置赋予身份绝对位置编码的核心思想是为序列中的每一个绝对位置索引 (pos) 生成一个唯一的向量表示 (\mathbf{p}_{pos})并将其与词嵌入相加作为模型的输入。2.1 正弦位置编码Transformer原论文采用了正弦位置编码。对于位置 (pos) 和维度索引 (i)(0 \le i d_{\text{model}})[PE_{(pos, 2i)} \sin\left( \frac{pos}{10000^{2i / d_{\text{model}}}} \right)][PE_{(pos, 2i1)} \cos\left( \frac{pos}{10000^{2i / d_{\text{model}}}} \right)]设计直觉不同频率的正余弦函数沿位置轴变化使每个位置产生唯一编码。波长范围从 (2\pi) 到约 (10000 \cdot 2\pi)覆盖了短程和长程位置变化。对于固定偏移量 (k)(PE_{posk}) 可以表示为 (PE_{pos}) 的线性函数因正余弦的和角公式这一性质理论上有利于模型学习相对位置信息。优缺点优点无需训练参数可处理任意长度序列理论上可外推。缺点外推至远超训练长度的序列时性能下降明显位置信息与词嵌入简单相加交互方式固定。2.2 可学习位置嵌入BERT、GPT等模型则采用可学习的位置嵌入定义一个嵌入矩阵 (\mathbf{P} \in \mathbb{R}^{L_{\max} \times d})其中 (L_{\max}) 为最大序列长度。位置 (pos) 的嵌入为 (\mathbf{P}_{pos})与词嵌入相加。优势位置表示完全由数据驱动学习灵活性高在训练长度范围内通常效果优于正弦编码。劣势无法处理超过 (L_{\max}) 的序列遇到更长输入需截断或插值参数随序列长度线性增长。2.3 绝对位置编码的共性问题无论是正弦还是可学习绝对位置编码都存在两个根本缺陷绝对位置依赖模型学习的是“第5个位置”的表示而非“当前词与上文相距多远”。当测试序列长度与训练分布不一致时绝对位置编码的外推能力有限。加法注入的刚性位置信息与词义信息以简单相加混合模型无法根据上下文动态调整位置与内容的交互权重。这些局限性催生了相对位置编码的设计。3. 相对位置编码关注距离而非绝对坐标相对位置编码的核心思想是不直接表示绝对位置而是在计算注意力分数时注入两个token之间的相对距离信息。3.1 原始Transformer中的相对位置偏置在原始的Transformer中位置信息已经通过正弦编码加到了输入中注意力计算时位置信息通过非线性变换间接交互。而显式的相对位置编码则直接将相对距离作为一个可学习的偏置项加入注意力分数。基本形式设 (i) 和 (j) 为两个位置它们的相对距离为 (i-j)。定义一个可学习的嵌入表 (\mathbf{B} \in \mathbb{R}^{(2L_{\max}-1) \times \text{num_heads}})或每个头独立在计算注意力分数时添加[\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) \text{softmax}\left( \frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d_k}} \mathbf{B}_{i-j} \right) \mathbf{V}]这里 (\mathbf{B}_{i-j}) 是一个标量取决于相对距离。为控制参数量通常设定一个最大相对距离 (K)超过该距离的偏置共享同一个值裁剪。3.2 Transformer-XL与T5的相对位置编码Transformer-XL提出了一种更精细的相对位置编码方案。它将注意力分数分解为内容-内容、内容-位置、位置-内容、位置-位置四项其中位置项使用相对位置嵌入。具体地[\mathbf{A}{ij}^{\text{rel}} \underbrace{\mathbf{E}{x_i}^\top \mathbf{W}q^\top \mathbf{W}k \mathbf{E}{x_j}}{(a)} \underbrace{\mathbf{E}{x_i}^\top \mathbf{W}q^\top \mathbf{W}k \mathbf{R}{i-j}}{(b)} \underbrace{\mathbf{u}^\top \mathbf{W}k \mathbf{E}{x_j}}{©} \underbrace{\mathbf{v}^\top \mathbf{W}k \mathbf{R}{i-j}}_{(d)}]其中 (\mathbf{R}_{i-j}) 是可学习的相对位置嵌入(\mathbf{u}, \mathbf{v}) 为可训练向量。T5采用了一种简化的相对位置偏置每个注意力头有一个可学习的标量偏置 (b_{\text{rel}})仅依赖于相对距离且进行分桶处理以限制参数。3.3 ALiBi线性偏置的简洁之美ALiBiAttention with Linear Biases是一种极简的相对位置编码不添加任何位置嵌入仅在注意力分数中加上一个与相对距离成正比的负偏置。[\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) \text{softmax}\left( \frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d_k}} - m \cdot |i-j| \right) \mathbf{V}]其中 (m) 是头特定的斜率如第 (h) 个头的斜率为 (2^{-8h/H})。这个简单的线性偏置强迫模型更关注邻近的token对远处token施加更大惩罚。ALiBi的优势零参数无需学习任何位置嵌入。极强外推能力由于偏置随距离线性增长模型在训练时未见过的长序列上也能合理分配注意力外推性能显著优于RoPE和绝对位置编码。ALiBi已被用于BLOOM等模型展示了无位置嵌入架构的潜力。3.4 相对位置编码的小结相对位置编码通过直接建模相对距离部分解决了绝对位置编码的长度外推问题。但许多实现如T5偏置、Transformer-XL仍受限于最大相对距离的裁剪外推时对超长距离的泛化依然受限。ALiBi的线性偏置在理论上可无限外推但牺牲了对复杂位置模式的细粒度建模能力。4. 旋转位置编码用复数旋转编码相对位置旋转位置编码Rotary Position Embedding, RoPE由Su等人在2021年提出现已成为LLaMA、GPT-NeoX、PaLM等主流大模型的标准配置。其核心思想极其优雅通过复数域中的旋转变换将token的绝对位置信息转化为query和key的内积中的相对位置信息。4.1 数学动机内积只依赖相对位置设词嵌入经线性投影得到查询向量 (\mathbf{q}_m)位置 (m)和键向量 (\mathbf{k}_n)位置 (n)。RoPE的目标是找到一个变换 (f)使得[\langle f(\mathbf{q}_m, m), f(\mathbf{k}_n, n) \rangle g(\mathbf{q}_m, \mathbf{k}_n, m-n)]即变换后的内积仅依赖于向量本身和相对位置 (m-n)而不依赖于绝对位置 (m) 和 (n)。4.2 二维情形下的旋转首先考虑二维向量 (\mathbf{q} \in \mathbb{R}^2)。将其视为复平面上的点 (q e^{i\theta})。RoPE将位置 (m) 编码为一个旋转角度 (m\theta)通过复数乘法旋转矩阵施加于向量[f(\mathbf{q}, m) \begin{pmatrix} \cos(m\theta) -\sin(m\theta) \ \sin(m\theta) \cos(m\theta) \end{pmatrix} \begin{pmatrix} q_0 \ q_1 \end{pmatrix}]此时两个旋转后向量的内积为[f(\mathbf{q}, m)^\top f(\mathbf{k}, n) \mathbf{q}^\top \mathbf{R}_{m-n} \mathbf{k}]其中 (\mathbf{R}_{m-n}) 是角度为 ((m-n)\theta) 的旋转矩阵。内积仅依赖于相对位置 (m-n)。4.3 推广至高维对于 (d) 维向量RoPE将维度分成 (d/2) 对每对作为一个二维子空间独立旋转。不同子空间采用不同的旋转频率[\theta_i 10000^{-2i/d}, \quad i 0, 1, \dots, d/2-1]对于查询向量 (\mathbf{q})位置 (m) 的RoPE变换为[\begin{pmatrix} q_{2i} \ q_{2i1} \end{pmatrix} \leftarrow\begin{pmatrix} \cos(m\theta_i) -\sin(m\theta_i) \ \sin(m\theta_i) \cos(m\theta_i) \end{pmatrix}\begin{pmatrix} q_{2i} \ q_{2i1} \end{pmatrix}]键向量 (\mathbf{k}) 同理。4.4 RoPE的实现与计算效率在实际代码中RoPE不显式构造旋转矩阵而是利用复数运算高效实现defrotate_half(x):x1,x2x.chunk(2,dim-1)returntorch.cat([-x2,x1],dim-1)defapply_rotary_pos_emb(q,k,cos,sin):# cos, sin: (seq_len, dim) 预计算好的频率q_embed(q*cos)(rotate_half(q)*sin)k_embed(k*cos)(rotate_half(k)*sin)returnq_embed,k_embed其中cos和sin根据位置和频率预计算形状为(max_seq_len, dim)。4.5 RoPE的优势相对位置的内化内积自然依赖于相对距离模型对相对位置的感知更加直接。良好的外推能力通过调整推理时的频率缩放如NTK-aware插值RoPE可有效外推到训练长度数倍的序列。与自注意力无缝融合RoPE仅修改了(\mathbf{Q})和(\mathbf{K})的计算方式不改变Transformer其他部分易于集成。保持向量模长旋转是正交变换不改变向量范数有助于训练稳定性。4.6 RoPE的变体与扩展线性RoPE调整频率基数为更大值如500,000使高频分量衰减更慢改善长距离外推。NTK-aware插值在外推时对高频维度进行非线性频率缩放避免直接插值导致的性能下降。YaRN结合NTK插值与温度缩放进一步提升极长上下文下的外推能力。5. 位置编码的对比与选型5.1 核心特性对比表特性绝对位置正弦绝对位置可学习相对位置偏置RoPEALiBi参数学习无有有无无长度外推能力较弱差硬截断中等受裁剪限制良好配合插值优秀对相对距离建模间接间接直接直接内积中直接线性偏置计算复杂度低低中中极低代表模型Transformer原论文BERT, GPT, GPT-2T5, Transformer-XLLLaMA, GPT-NeoX, PaLMBLOOM5.2 选型建议NLU任务、训练长度固定BERT式的可学习绝对位置嵌入足够简单有效。需要生成长文本、重视外推RoPE是当前大模型首选配合插值技术可外推至32k甚至100k上下文。极致外推需求、算力受限ALiBi几乎零成本且外推能力最强适合长文档建模。追求可解释性与理论优雅RoPE通过复数旋转将相对位置编码内嵌数学形式优美。6. 代码实践实现与对比6.1 正弦绝对位置编码importtorchimportmathdefsinusoidal_position_encoding(seq_len,d_model):petorch.zeros(seq_len,d_model)positiontorch.arange(0,seq_len).unsqueeze(1)div_termtorch.exp(torch.arange(0,d_model,2)*-(math.log(10000.0)/d_model))pe[:,0::2]torch.sin(position*div_term)pe[:,1::2]torch.cos(position*div_term)returnpe6.2 RoPE实现简化版classRotaryPositionEmbedding(torch.nn.Module):def__init__(self,dim,max_seq_len2048,base10000):super().__init__()inv_freq1.0/(base**(torch.arange(0,dim,2).float()/dim))ttorch.arange(max_seq_len).float()freqstorch.einsum(i,j-ij,t,inv_freq)embtorch.cat((freqs,freqs),dim-1)self.register_buffer(cos,emb.cos())self.register_buffer(sin,emb.sin())defforward(self,q,k,seq_len):cosself.cos[:seq_len,:].unsqueeze(0).unsqueeze(0)# (1, 1, seq_len, dim)sinself.sin[:seq_len,:].unsqueeze(0).unsqueeze(0)q_embed(q*cos)(rotate_half(q)*sin)k_embed(k*cos)(rotate_half(k)*sin)returnq_embed,k_embed6.3 ALiBi偏置生成defget_alibi_slopes(num_heads):# 生成指数递减的斜率returntorch.tensor([2**(-8*i/num_heads)foriinrange(1,num_heads1)])defapply_alibi(attn_scores,alibi_slopes,seq_len):# attn_scores: (batch, heads, seq_len, seq_len)distancetorch.arange(seq_len).unsqueeze(0)-torch.arange(seq_len).unsqueeze(1)distancedistance.abs().unsqueeze(0).unsqueeze(0)# (1, 1, seq_len, seq_len)alibi_bias-alibi_slopes.view(-1,1,1)*distancereturnattn_scoresalibi_bias7. 前沿探索无位置编码与混合方案7.1 无显式位置编码的可能性近期研究表明通过特殊的注意力掩码或递归架构模型可隐式习得位置信息。例如Hyena和Mamba等状态空间模型不使用显式位置编码在长序列任务上表现出色。这为未来彻底摆脱位置编码提供了新思路。7.2 长度外推的持续挑战尽管RoPE配合插值技术已将上下文窗口扩展至百万token级别但真正有效的长程依赖建模仍面临注意力稀疏化、内存瓶颈等问题。位置编码是必要条件但远非充分条件。8. 结语位置序列的灵魂位置编码的设计演进是Transformer架构中最能体现数学之美的篇章之一。从绝对位置的直白赋予到相对位置的偏置注入再到旋转编码的复数优雅每一步都凝结着研究者对“如何让注意力感知顺序”这一根本问题的深刻洞察。在LLaMA、GPT-4等巨人身躯中RoPE以其卓越的外推性能和理论完备性稳坐王座ALiBi则以极致简洁在某些场景下异军突起。理解这些位置编码的原理与差异不仅是读懂当今大模型技术栈的必要基础更是未来设计新一代序列模型时的灵感源泉。当你下次调用model.generate时或许可以多想一想那些流畅生成的文本背后是无数复数旋转在默默编织着位置的丝线让词语在时间的长河中找到属于自己的坐标。点击“AladdinEdu你的AI学习实践工作坊”注册即送-H卡级别算力沉浸式云原生集成开发环境80G大显存多卡并行按量弹性计费教育用户更享超低价。