相对位置编码的泛化能力测试：RoPE、ALiBi与T5 Bias的外推实验

相对位置编码的泛化能力测试：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 PE<br/>外推: ❌ 差<br/>旋转不变性"} C --> C2["Learned PE<br/>外推: ❌ 完全无法外推<br/>最大长度在训练时固定"] B --> D[相对位置编码] D --> D1["RoPE<br/>外推: ⚠️ 中等<br/>需要YaRN/NTK等插值"] D --> D2["ALiBi<br/>外推: ✅ 好<br/>线性偏置天然外推"] D --> D3["T5 Relative Bias<br/>外推: ⚠️ 中等<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 Embedding（RoPE, Su et al., 2021）通过旋转矩阵将位置信息注入query和key向量。对于位置m和n，RoPE定义的注意力分数为：

$$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, device=x.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, device=x.device).float() / self.d_head ) ) positions = torch.arange( offset, offset + seq_len, device=x.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的线性偏置与天然外推

ALiBi（Attention 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) ], device=query.device) # (H,) # 位置距离矩阵: (S, S) positions = torch.arange(seq_len, device=query.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, device=query.device) * float('-inf'), diagonal=1 ).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的外推表现：

关键发现：(1) ALiBi在训练长度上略差于RoPE（因无显式位置信息），但在外推中表现最稳定；(2) RoPE+NTK调整是所有方案中精度最好的，但需要额外的超参数调优；(3) T5 Relative Bias的外推能力明显弱于RoPE和ALiBi，因其bucketed bias在遇到训练时未见过的长距离时行为不确定。

五、总结

位置编码的选择实质上是在"训练长度内的精度"与"外推长度下的泛化"之间的权衡：(1) 如果最大推理长度固定且已知，RoPE（配合NTK或PI策略）在精度上是最优选择；(2) 如果需要支持不可预见的超长序列（如生产环境可能遇到极长文档），ALiBi的天然外推特性提供了最可靠的泛化保障；(3) 当前趋势是RoPE+动态NTK缩放的组合——在LLaMA系列、Mistral和Qwen等主流模型中已成为标配，结合了RoPE的精度优势和NTK调整的外推能力。