Transformer 位置编码深入解析从正弦编码到 RoPE、ALiBi1. 引言Transformer 的自注意力机制是置换不变的——打乱输入 token 的顺序输出不变。位置编码Positional Encoding是让模型知道谁在前、谁在后的关键。本文将从最初的正弦编码到现代的 RoPE、ALiBi系统讲解位置编码的演进。位置编码分类类型代表特点绝对位置编码正弦编码、可学习编码加到输入嵌入上相对位置编码ALiBi、Relative PE编码 token 间距离旋转位置编码RoPE注意力分数中注入位置无显式编码CNN 隐式位置依赖结构归纳偏置2. 正弦位置编码Sinusoidal2.1 原理PE(pos, 2i) sin(pos / 10000^(2i/d)) PE(pos, 2i1) cos(pos / 10000^(2i/d)) 其中 pos token 在序列中的位置 i 嵌入维度的索引 d 嵌入维度2.2 实现importtorchimportmathclassSinusoidalPositionalEncoding(torch.nn.Module):def__init__(self,d_model,max_len5000):super().__init__()petorch.zeros(max_len,d_model)positiontorch.arange(0,max_len,dtypetorch.float).unsqueeze(1)div_termtorch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))pe[:,0::2]torch.sin(position*div_term)pe[:,1::2]torch.cos(position*div_term)pepe.unsqueeze(0)# (1, max_len, d_model)self.register_buffer(pe,pe)defforward(self,x):# x: (batch, seq_len, d_model)returnxself.pe[:,:x.size(1)]2.3 特点优点无需训练可泛化到任意长度缺点绝对位置编码与注意力计算分离关键性质PE(posk) 可以表示为 PE(pos) 的线性变换3. 可学习位置编码classLearnedPositionalEncoding(torch.nn.Module):def__init__(self,d_model,max_len512):super().__init__()self.embeddingtorch.nn.Embedding(max_len,d_model)defforward(self,x):positionstorch.arange(x.size(1),devicex.device)returnxself.embedding(positions)ViT、GPT-2 使用此方案。简单有效但无法泛化到训练时未见过的长度。4. RoPERotary Position Embedding4.1 核心思想RoPE 不直接修改输入嵌入而是在注意力计算时注入位置信息q_m R(m) · q // 查询向量旋转 m 个位置 k_n R(n) · k // 键向量旋转 n 个位置 q_m · k_n (R(m)·q) · (R(n)·k) q · R(m-n)·k 结果只依赖相对位置 (m-n)4.2 二维旋转对于二维向量 [x, y]旋转角度 θ R(θ) [cos θ -sin θ] [sin θ cos θ] R(θ)·[x, y] [x·cos θ - y·sin θ, x·sin θ y·cos θ]4.3 高维 RoPE 实现classRotaryPositionalEmbedding:RoPE 实现def__init__(self,dim,base10000):# 计算频率inv_freq1.0/(base**(torch.arange(0,dim,2).float()/dim))self.register_buffer(inv_freq,inv_freq)def_build_cache(self,seq_len,device):ttorch.arange(seq_len,devicedevice).float()freqstorch.outer(t,self.inv_freq.to(device))# cos 和 sinself._cos_cachedfreqs.cos()self._sin_cachedfreqs.sin()def_apply_rotary(self,x,cos,sin):应用旋转# x: (batch, heads, seq_len, head_dim)dx.size(-1)//2x1,x2x[...,:d],x[...,d:]# 扩展 cos/sin 维度coscos[:x.size(2),:].unsqueeze(0).unsqueeze(0)sinsin[:x.size(2),:].unsqueeze(0).unsqueeze(0)# 旋转out1x1*cos-x2*sin out2x1*sinx2*cosreturntorch.cat([out1,out2],dim-1)defforward(self,q,k,seq_len):self._build_cache(seq_len,q.device)q_rotself._apply_rotary(q,self._cos_cached,self._sin_cached)k_rotself._apply_rotary(k,self._cos_cached,self._sin_cached)returnq_rot,k_rot# 简化版 RoPE实际使用defapply_rope(q,k,freqs_cis):freqs_cis: 预计算的复数频率# 转为复数q_complextorch.view_as_complex(q.float().reshape(*q.shape[:-1],-1,2))k_complextorch.view_as_complex(k.float().reshape(*k.shape[:-1],-1,2))# 乘以旋转因子q_outtorch.view_as_real(q_complex*freqs_cis).flatten(-2)k_outtorch.view_as_real(k_complex*freqs_cis).flatten(-2)returnq_out.type_as(q),k_out.type_as(k)4.4 RoPE 长度外推# NTK-aware 缩放defapply_ntk_scaling(base,dim,factor):NTK-aware RoPE 缩放支持更长上下文new_basebase*factor**(dim/(dim-2))inv_freq1.0/(new_base**(torch.arange(0,dim,2).float()/dim))returninv_freq# YaRNYet another RoPE extensioNdefyarn_rope(dim,base,factor,beta_fast32,beta_slow1):YaRN 动态缩放# 计算每个频率维度的缩放因子inv_freq1.0/(base**(torch.arange(0,dim,2).float()/dim))freqs1.0/inv_freq# 找到需要缩放的维度lowdim*math.log(beta_slow)/math.log(base)highdim*math.log(beta_fast)/math.log(base)scaletorch.ones(dim//2)foriinrange(dim//2):iffreqs[i]low:scale[i]1.0/factoreliffreqs[i]high:scale[i]1.0else:# 线性插值t(freqs[i]-low)/(high-low)scale[i](1-t)/factortreturninv_freq*scale5. ALiBiAttention with Linear Biases5.1 原理ALiBi 完全不用位置编码而是在注意力分数上加一个线性偏置Attention(Q, K) softmax(Q·K^T / √d m · bias) bias[i,j] -|i - j| // 距离越远惩罚越大 m 斜率每个头不同几何级数5.2 实现importtorchimportmathdefalibi_bias(num_heads,max_len):生成 ALiBi 偏置矩阵# 斜率几何级数 2^(-8/n), 2^(-16/n), ...slopes2**(-8*torch.arange(1,num_heads1)/num_heads)# 距离矩阵positionstorch.arange(max_len)distancepositions.unsqueeze(0)-positions.unsqueeze(1)# (max_len, max_len)distancedistance.abs().float()# 偏置 斜率 × 距离bias-slopes.unsqueeze(1).unsqueeze(2)*distance.unsqueeze(0)returnbias# (num_heads, max_len, max_len)classALiBiAttention(torch.nn.Module):def__init__(self,d_model,num_heads,max_len4096):super().__init__()self.num_headsnum_heads self.head_dimd_model//num_heads self.q_projtorch.nn.Linear(d_model,d_model)self.k_projtorch.nn.Linear(d_model,d_model)self.v_projtorch.nn.Linear(d_model,d_model)self.o_projtorch.nn.Linear(d_model,d_model)# 预计算 ALiBi 偏置self.register_buffer(alibi,alibi_bias(num_heads,max_len))defforward(self,x):B,L,_x.shape qself.q_proj(x).view(B,L,self.num_heads,self.head_dim).transpose(1,2)kself.k_proj(x).view(B,L,self.num_heads,self.head_dim).transpose(1,2)vself.v_proj(x).view(B,L,self.num_heads,self.head_dim).transpose(1,2)# 注意力 ALiBiattn(q k.transpose(-2,-1))/math.sqrt(self.head_dim)attnattnself.alibi[:,:L,:L].unsqueeze(0)attnattn.softmax(dim-1)out(attn v).transpose(1,2).reshape(B,L,-1)returnself.o_proj(out)6. 各方案对比方案外推能力训练开销推理开销长序列效果正弦编码有限无低一般可学习编码无低低差RoPE需缩放中中好ALiBi天然外推无低好YaRN RoPE强外推中中很好主流选择Llama/Qwen/DeepSeek→ RoPE配 YaRN 长度外推BLOOM/MPT→ ALiBiGPT 系列→ 可学习绝对位置编码7. 总结位置编码的核心演进正弦编码开创性工作但与注意力分离RoPE在注意力分数中注入相对位置成为主流ALiBi最简单方案线性偏置即可天然支持外推YaRNRoPE 的长度外推方案支持 128K 上下文