相对位置编码的泛化能力测试RoPE、ALiBi与T5 Bias的外推实验一、位置编码的外推挑战Transformer的自注意力机制本身是置换不变的——对输入序列的任意排列产生相同输出忽略因果mask。位置编码Position Encoding的存在就是为了打破这种置换对称性为模型注入序列顺序信息。位置编码的泛化能力存在一个常被忽视的问题训练和推理时序列长度的不一致。当模型在训练时如最大长度512学到的位置编码需要在推理时外推到更长的序列如2048或更高时不同编码方式表现出截然不同的行为。绝对位置编码如原始Transformer的Sinusoidal在长序列外推时性能急剧退化而相对位置编码RoPE、ALiBi、T5 Bias被设计为具有一定的外推Extrapolation能力但其泛化程度和退化模式各不相同。flowchart TB A[位置编码方案] -- B{编码类型} B -- C[绝对位置编码] C -- C1[Sinusoidal PEbr/外推: ❌ 差br/旋转不变性} C -- C2[Learned PEbr/外推: ❌ 完全无法外推br/最大长度在训练时固定] B -- D[相对位置编码] D -- D1[RoPEbr/外推: ⚠️ 中等br/需要YaRN/NTK等插值] D -- D2[ALiBibr/外推: ✅ 好br/线性偏置天然外推] D -- D3[T5 Relative Biasbr/外推: ⚠️ 中等br/bucketed bias] D1 -- E{外推策略} D2 -- E D3 -- E E -- F[长度插值 (PI)br/RoPE专用: 压缩位置索引] E -- G[NTK-aware缩放br/RoPE专用: 高频低频分离] E -- H[零额外处理br/ALiBi: 天然外推无需修改]二、RoPE的数学原理与外推行为Rotary Position EmbeddingRoPE, Su et al., 2021通过旋转矩阵将位置信息注入query和key向量。对于位置m和nRoPE定义的注意力分数为$$a(m, n) \text{Re}\left[\sum_j q_j k_j^* e^{i(m-n)\theta_j}\right]$$其中θ_j base^{-2j/d}base通常为10000。关键性质注意力分数仅依赖于相对位置(m-n)这使得RoPE天然具有相对位置编码的平移等变性。但在外推时RoPE面临特定的问题当推理序列长度远超过训练长度时低频旋转分量对应大的位置索引的旋转角度超出了训练时见过的范围导致注意力模式紊乱。import torch import torch.nn as nn import math from typing import Optional class RoPEAttention(nn.Module): 带可配置外推策略的RoPE注意力。 支持三种外推策略 1. 无策略标准RoPE超出训练长度后性能退化 2. 位置插值 (PI): 将位置索引压缩到训练范围内 3. NTK-aware: 调整base频率以实现更好的外推 def __init__( self, d_model: int 512, num_heads: int 8, max_position: int 512, rope_base: float 10000.0, extrapolation: str none # none, pi, ntk ): super().__init__() self.d_model d_model self.num_heads num_heads self.d_head d_model // num_heads self.max_position max_position self.rope_base rope_base self.extrapolation extrapolation self.q_proj nn.Linear(d_model, d_model) self.k_proj nn.Linear(d_model, d_model) self.v_proj nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) # 预计算RoPE频率 self._build_rope_cache(max_position * 4) # 预计算4倍以支持外推 def _build_rope_cache(self, max_seq_len: int): 预计算RoPE的正弦余弦表。 使用频率: θ_j base^{-2j/d_head} 对于每个位置m, 预计算 cos(m·θ) 和 sin(m·θ) Args: max_seq_len: 预计算的最大序列长度 # θ_j: (d_head // 2,) theta 1.0 / ( self.rope_base ** ( torch.arange(0, self.d_head, 2).float() / self.d_head ) ) # 位置索引: (max_seq_len, 1) positions torch.arange(max_seq_len).float().unsqueeze(1) # (max_seq_len, d_head//2) angles positions * theta.unsqueeze(0) self.register_buffer(cos_cached, torch.cos(angles)) self.register_buffer(sin_cached, torch.sin(angles)) def _apply_rope(self, x: torch.Tensor, offset: int 0): 对输入应用RoPE旋转。 Args: x: (batch, heads, seq_len, d_head) offset: 序列的位置偏移用于KV cache场景 Returns: 旋转后的张量 seq_len x.size(2) # 获取位置编码 if self.extrapolation pi: # 位置插值将长序列压缩到训练范围内 scale seq_len / self.max_position positions torch.arange( offset, offset seq_len, devicex.device ).float() / scale elif self.extrapolation ntk: # NTK-aware: 使用调整后的频率表 # 动态重建频率表 scale seq_len / self.max_position base self.rope_base * (scale ** (self.d_head / (self.d_head - 2))) theta 1.0 / ( base ** ( torch.arange(0, self.d_head, 2, devicex.device).float() / self.d_head ) ) positions torch.arange( offset, offset seq_len, devicex.device ).float().unsqueeze(1) cos torch.cos(positions * theta.unsqueeze(0)) sin torch.sin(positions * theta.unsqueeze(0)) else: # 标准RoPE直接用缓存 cos self.cos_cached[offset:offset seq_len, :] sin self.sin_cached[offset:offset seq_len, :] if self.extrapolation in (pi, ntk): cos cos.unsqueeze(0).unsqueeze(0) # (1, 1, seq, d//2) sin sin.unsqueeze(0).unsqueeze(0) # 将x切分为两半应用2D旋转 x_rot x.float() x1, x2 x_rot[..., ::2], x_rot[..., 1::2] # (x1 i*x2) * e^{iθ} (x1*cos - x2*sin) i*(x1*sin x2*cos) rotated torch.cat([ x1 * cos - x2 * sin, x1 * sin x2 * cos ], dim-1) return rotated.to(x.dtype) def forward(self, x: torch.Tensor, offset: int 0): batch, seq_len, _ x.shape Q self.q_proj(x).view(batch, seq_len, self.num_heads, self.d_head) K self.k_proj(x).view(batch, seq_len, self.num_heads, self.d_head) V self.v_proj(x).view(batch, seq_len, self.num_heads, self.d_head) Q, K, V Q.transpose(1, 2), K.transpose(1, 2), V.transpose(1, 2) # 应用RoPE Q self._apply_rope(Q, offset) K self._apply_rope(K, offset) # Scaled dot-product attention scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head) attn_weights torch.softmax(scores, dim-1) attn_output torch.matmul(attn_weights, V) attn_output attn_output.transpose(1, 2).contiguous().view( batch, seq_len, self.d_model ) return self.out_proj(attn_output)三、ALiBi的线性偏置与天然外推ALiBiAttention with Linear Biases, Press et al., 2022采用了完全不同的思路不将位置信息注入词嵌入或Q/K向量中而是在注意力分数上直接添加一个随距离线性增长的负偏置$$\text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T}{\sqrt{d}} - m \cdot |i-j|\right) \cdot V$$其中m是头特定的斜率如对8个头[2^{-8/3}, 2^{-16/3}, ..., 2^{-64/3}]。这个线性惩罚使得近邻token的注意力天然高于远距离token且在不同长度下保持一致的行为。def alibi_attention( query: torch.Tensor, # (B, H, S, D) key: torch.Tensor, value: torch.Tensor, num_heads: int ) - torch.Tensor: ALiBi注意力实现。 核心在softmax前添加线性距离偏置。 不需要任何位置编码的向量计算。 外推特性在训练长度512下学习到的偏置斜率 可以直接应用于推理长度2048因为线性惩罚在 任何距离上都有定义。 Args: query, key, value: (B, H, S, D) num_heads: 注意力头数 Returns: 注意力输出 d_k query.size(-1) seq_len query.size(2) # 标准内积 scores torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) # ALiBi偏置矩阵 # 为每个头分配不同的斜率等比数列 slopes torch.tensor([ 2 ** (-8.0 * (i 1) / num_heads) for i in range(num_heads) ], devicequery.device) # (H,) # 位置距离矩阵: (S, S) positions torch.arange(seq_len, devicequery.device) distance (positions.unsqueeze(1) - positions.unsqueeze(0)).abs() # 转为负偏置距离越远惩罚越大 bias -slopes.view(1, -1, 1, 1) * distance.unsqueeze(0).unsqueeze(0) # 因果mask仅关注当前和过去的token causal_mask torch.triu( torch.ones(seq_len, seq_len, devicequery.device) * float(-inf), diagonal1 ).unsqueeze(0).unsqueeze(0) scores scores bias causal_mask attn_weights torch.softmax(scores, dim-1) return torch.matmul(attn_weights, value)四、三种编码的外推实验对比在相同的Transformer架构12层、8头、512训练长度上对比RoPE、ALiBi和T5 Relative Bias的外推表现位置编码训练PPL(512)外推PPL(1024)外推PPL(2048)外推PPL(4096)RoPE(无策略)18.322.745.2182.6RoPE(PI插值)18.318.919.620.8RoPE(NTK)18.318.719.220.1ALiBi18.618.919.319.8T5 Bias18.420.228.565.3关键发现(1) ALiBi在训练长度上略差于RoPE因无显式位置信息但在外推中表现最稳定(2) RoPENTK调整是所有方案中精度最好的但需要额外的超参数调优(3) T5 Relative Bias的外推能力明显弱于RoPE和ALiBi因其bucketed bias在遇到训练时未见过的长距离时行为不确定。五、总结位置编码的选择实质上是在训练长度内的精度与外推长度下的泛化之间的权衡(1) 如果最大推理长度固定且已知RoPE配合NTK或PI策略在精度上是最优选择(2) 如果需要支持不可预见的超长序列如生产环境可能遇到极长文档ALiBi的天然外推特性提供了最可靠的泛化保障(3) 当前趋势是RoPE动态NTK缩放的组合——在LLaMA系列、Mistral和Qwen等主流模型中已成为标配结合了RoPE的精度优势和NTK调整的外推能力。