1. 为什么需要位置编码
Transformer 的 Self-Attention 机制是置换不变（permutation invariant）的：打乱输入序列的顺序，注意力输出不变。这意味着 Transformer 本身没有序列顺序的概念。
位置编码的使命就是：为模型注入位置信息。
没有位置编码:
Attention("猫追狗") == Attention("狗追猫") # 灾难！
有位置编码:
Attention("猫追狗") ≠ Attention("狗追猫") # 位置区分了主语和宾语
位置编码的设计需要满足几个理想特性：
- 唯一性：每个位置应有唯一的编码
- 相对性：能表达位置间的相对距离
- 外推性：能处理比训练时更长的序列
- 周期性：能捕捉位置的模式规律
2. 绝对位置编码
2.1 Sinusoidal（正弦余弦）位置编码
Vaswani 等人 (2017) 提出的正弦余弦位置编码：
$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d}}\right)$$ $$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d}}\right)$$
import torch
import torch.nn as nn
import math
class SinusoidalPositionalEncoding(nn.Module):
def __init__(self, d_model=512, max_len=5000):
super().__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# 频率: 1 / (10000^(2i/d))
div_term = torch.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)
pe = pe.unsqueeze(0) # (1, max_len, d_model)
self.register_buffer('pe', pe)
def forward(self, x):
# x: (batch, seq_len, d_model)
return x + self.pe[:, :x.size(1), :]
关键性质：对于固定的偏移 $k$，$PE_{pos+k}$ 可以表示为 $PE_{pos}$ 的线性函数：
$$\begin{pmatrix} \sin(\omega_k \cdot (pos + \delta)) \ \cos(\omega_k \cdot (pos + \delta)) \end{pmatrix} = R(\omega_k \delta) \begin{pmatrix} \sin(\omega_k \cdot pos) \ \cos(\omega_k \cdot pos) \end{pmatrix}$$
这意味着正弦编码隐含了相对位置信息，但这一信息并非显式利用。
2.2 Learned（可学习）位置编码
BERT 和 GPT-2 采用了最简单的方案：直接为每个位置学习一个嵌入向量。
class LearnedPositionalEncoding(nn.Module):
def __init__(self, d_model=512, max_len=512):
super().__init__()
self.pos_embedding = nn.Embedding(max_len, d_model)
self.register_buffer('position_ids', torch.arange(max_len).unsqueeze(0))
def forward(self, x):
batch_size, seq_len, _ = x.shape
positions = self.position_ids[:, :seq_len]
return x + self.pos_embedding(positions)
| 方案 | 优点 | 缺点 | 外推性 | 代表模型 |
|---|---|---|---|---|
| Sinusoidal | 无参数、有理论保障 | 效果一般 | 有限 | 原始 Transformer |
| Learned | 简单、效果好 | 无法外推 | ❌ | BERT, GPT-2 |
3. 相对位置编码
3.1 Relative Position Representation (Shaw et al., 2018)
核心思想：在注意力计算中，不再使用绝对位置编码加到输入上，而是将相对位置信息直接融入注意力分数。
$$A_{ij}^{rel} = (W_q x_i)^\top (W_k x_j + a_{ij}^K) + b_{ij}^V$$
其中 $a_{ij}^K$ 和 $b_{ij}^V$ 是依赖相对位置 $i - j$ 的可学习参数。
class RelativePositionAttention(nn.Module):
def __init__(self, d_model, num_heads, max_rel_pos=128):
super().__init__()
self.num_heads = num_heads
self.d_k = d_model // num_heads
self.max_rel_pos = max_rel_pos
# 相对位置嵌入 (2*max_rel_pos + 1 个位置)
# 从 -max_rel_pos 到 +max_rel_pos
self.rel_pos_embed = nn.Embedding(2 * max_rel_pos + 1, self.d_k)
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
def forward(self, x):
B, N, D = x.shape
H, dk = self.num_heads, self.d_k
q = self.W_q(x).view(B, N, H, dk).transpose(1, 2)
k = self.W_k(x).view(B, N, H, dk).transpose(1, 2)
v = self.W_v(x).view(B, N, H, dk).transpose(1, 2)
# 标准注意力分数
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(dk)
# 添加相对位置偏置
# 构建相对位置索引矩阵
rel_pos_indices = self._get_rel_pos_indices(N)
rel_pos_bias = self.rel_pos_embed(rel_pos_indices) # (N, N, dk)
# q · rel_pos_bias
rel_bias = torch.einsum('bhid,ijd->bhij', q, rel_pos_bias)
scores = scores + rel_bias / math.sqrt(dk)
attn = F.softmax(scores, dim=-1)
output = torch.matmul(attn, v)
return output.transpose(1, 2).reshape(B, N, D)
def _get_rel_pos_indices(self, seq_len):
indices = torch.arange(seq_len)
rel = indices.unsqueeze(0) - indices.unsqueeze(1) # (N, N)
rel = rel.clamp(-self.max_rel_pos, self.max_rel_pos)
return rel + self.max_rel_pos # 偏移到非负索引
3.2 T5 Relative Position Bias
T5 简化了相对位置编码，使用 桶化（bucketing） 机制：
class T5RelativePositionBias(nn.Module):
def __init__(self, num_buckets=32, num_heads=12, max_distance=128):
super().__init__()
self.num_buckets = num_buckets
self.max_distance = max_distance
self.rel_attn_bias = nn.Embedding(num_buckets, num_heads)
def _relative_position_bucket(self, relative_position):
"""将相对位置映射到桶中（对数分桶）"""
result = 0
num_negative = self.num_buckets // 2
# 正负方向分开处理
relative_position_buckets = torch.where(
relative_position < 0,
-relative_position + num_negative,
relative_position
)
# 对数分桶：近距离精细，远距离粗糙
max_exact = num_negative
is_small = relative_position_buckets < max_exact
is_large = ~is_small
large_value = max_exact + (
torch.log(relative_position_buckets.float() / max_exact)
/ math.log(self.max_distance / max_exact)
* (self.num_buckets - max_exact)
).long()
large_value = large_value.clamp(min=0, max=self.num_buckets - 1)
return torch.where(is_small, relative_position_buckets, large_value)
def forward(self, seq_len):
# 计算相对位置
context_pos = torch.arange(seq_len)[:, None]
memory_pos = torch.arange(seq_len)[None, :]
relative_pos = memory_pos - context_pos
# 分桶
rp_bucket = self._relative_position_bucket(relative_pos)
# (seq_len, seq_len, num_heads)
bias = self.rel_attn_bias(rp_bucket)
return bias.permute(2, 0, 1).unsqueeze(0) # (1, H, N, N)
4. RoPE：旋转位置编码
4.1 核心思想
RoPE (Rotary Position Embedding, Su et al., 2021) 的核心思想是：通过旋转矩阵将位置信息编码到 Query 和 Key 中，使得 $q_m^\top k_n$ 仅依赖于相对位置 $m - n$。
数学上，RoPE 将 $d$ 维向量视为 $d/2$ 个二维平面上的复数，对每个平面施加不同频率的旋转：
$$q_m’ = R_m q_m, \quad k_n’ = R_n k_n$$
其中 $R_m$ 是块对角旋转矩阵：
$$R_m = \begin{pmatrix} \cos(m\theta_1) & -\sin(m\theta_1) & & \ \sin(m\theta_1) & \cos(m\theta_1) & & \ & & \cos(m\theta_2) & -\sin(m\theta_2) \ & & \sin(m\theta_2) & \cos(m\theta_2) \ & & & \ddots \end{pmatrix}$$
4.2 实现
class RotaryPositionalEmbedding(nn.Module):
def __init__(self, dim, max_seq_len=4096, base=10000):
super().__init__()
self.dim = dim
# 计算频率: θ_i = 1 / (base^(2i/dim))
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
self.register_buffer('inv_freq', inv_freq)
# 预计算 cos/sin
t = torch.arange(max_seq_len).float()
freqs = torch.outer(t, self.inv_freq) # (max_len, dim/2)
# 复制以匹配维度
emb = torch.cat([freqs, freqs], dim=-1) # (max_len, dim)
self.register_buffer('cos_cached', emb.cos())
self.register_buffer('sin_cached', emb.sin())
def forward(self, x, seq_len=None):
if seq_len is None:
seq_len = x.shape[1]
cos = self.cos_cached[:seq_len].unsqueeze(0).unsqueeze(0) # (1,1,N,dim)
sin = self.sin_cached[:seq_len].unsqueeze(0).unsqueeze(0)
return self.apply_rotary(x, cos, sin)
def apply_rotary(self, x, cos, sin):
"""应用旋转位置编码"""
d = x.shape[-1]
# 将 x 拆分为两半
x1 = x[..., :d//2]
x2 = x[..., d//2:]
# 旋转: (x1, x2) → (x1*cos - x2*sin, x1*sin + x2*cos)
# 注意 cos/sin 的维度安排
cos = cos[..., :d//2]
sin = sin[..., :d//2]
rotated = torch.cat([
x1 * cos - x2 * sin,
x1 * sin + x2 * cos
], dim=-1)
return rotated
4.3 RoPE 的关键性质
为什么 RoPE 能表达相对位置？
$$q_m’^\top k_n’ = (R_m q_m)^\top (R_n k_n) = q_m^\top R_m^\top R_n k_n = q_m^\top R_{n-m} k_n$$
这表明旋转后的内积仅依赖于相对位置 $n - m$。
# 验证 RoPE 的相对位置性质
def verify_rope_property():
dim = 8
rope = RotaryPositionalEmbedding(dim)
q = torch.randn(1, 1, 1, dim)
k = torch.randn(1, 1, 1, dim)
# 位置 m=3, n=7 → 相对距离 4
q_m = rope.apply_rotary(q,
rope.cos_cached[3:4].view(1,1,1,dim),
rope.sin_cached[3:4].view(1,1,1,dim))
k_n = rope.apply_rotary(k,
rope.cos_cached[7:8].view(1,1,1,dim),
rope.sin_cached[7:8].view(1,1,1,dim))
dot_1 = (q_m * k_n).sum().item()
# 位置 m=10, n=14 → 相对距离 4 (相同)
q_m2 = rope.apply_rotary(q,
rope.cos_cached[10:11].view(1,1,1,dim),
rope.sin_cached[10:11].view(1,1,1,dim))
k_n2 = rope.apply_rotary(k,
rope.cos_cached[14:15].view(1,1,1,dim),
rope.sin_cached[14:15].view(1,1,1,dim))
dot_2 = (q_m2 * k_n2).sum().item()
print(f"Distance 4 at positions (3,7): {dot_1:.6f}")
print(f"Distance 4 at positions (10,14): {dot_2:.6f}")
print(f"Equal: {abs(dot_1 - dot_2) < 1e-5}")
5. ALiBi：注意力线性偏置
5.1 核心思想
ALiBi (Press et al., 2022) 采取了极其简洁的方案：不加位置编码到输入上，而是在注意力分数上添加一个与距离成正比的线性偏置。
$$A_{ij} = \text{softmax}\left(q_i^\top k_j - m \cdot |i - j|\right)$$
其中 $m$ 是与注意力头相关的斜率，呈几何衰减：
class ALiBiAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.d_k = d_model // num_heads
# ALiBi 的斜率: 几何序列 1/2^(2i/(n-1))
# 对于 8 个头: [1/2^0, 1/2^1, ..., 1/2^7] 的变体
def get_slopes(n):
if n <= 0:
return []
start = 2 ** (-8 / n)
return [start ** (i + 1) for i in range(n)]
slopes = get_slopes(num_heads)
self.register_buffer('slopes', torch.tensor(slopes))
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
def forward(self, x, mask=None):
B, N, D = x.shape
H, dk = self.num_heads, self.d_k
q = self.W_q(x).view(B, N, H, dk).transpose(1, 2)
k = self.W_k(x).view(B, N, H, dk).transpose(1, 2)
v = self.W_v(x).view(B, N, H, dk).transpose(1, 2)
# 标准注意力分数
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(dk)
# 添加 ALiBi 偏置
# 构建距离矩阵: |i - j|
positions = torch.arange(N, device=x.device)
distance = positions[None, :] - positions[:, None] # (N, N)
distance = distance.abs().tril() # 因果掩码下取下三角
# 每个头有不同的斜率
# slopes: (H,) → (H, 1, 1) → 广播到 (H, N, N)
alibi_bias = -self.slopes[:, None, None] * distance[None, :, :]
scores = scores + alibi_bias
# 因果掩码
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn = F.softmax(scores, dim=-1)
output = torch.matmul(attn, v)
return self.W_o(output.transpose(1, 2).reshape(B, N, D))
5.2 ALiBi 的外推能力
ALiBi 最大的优势是零样本外推：训练时序列长度为 1024，推理时可扩展到 2048+ 而无需微调。
位置编码方案的外推能力对比:
┌─────────────────┬──────────────┬────────────────────────┐
│ 方案 │ 训练长度 │ 外推到 2x 长度 │
├─────────────────┼──────────────┼────────────────────────┤
│ Learned PE │ 512 │ ❌ 性能急剧下降 │
│ Sinusoidal │ 512 │ ⚠️ 性能显著下降 │
│ T5 Bias │ 512 │ ⚠️ 桶截断导致下降 │
│ RoPE │ 2048 │ ⚠️ 需要插值/缩放 │
│ ALiBi │ 1024 │ ✅ 几乎无损 │
│ RoPE + NTK │ 2048 │ ✅ 良好外推 │
│ RoPE + YARN │ 2048 │ ✅ 良好外推 │
└─────────────────┴──────────────┴────────────────────────┘
6. 方案全面对比
| 维度 | Sinusoidal | Learned | T5 Bias | RoPE | ALiBi |
|---|---|---|---|---|---|
| 类型 | 绝对 | 绝对 | 相对 | 相对(旋转) | 相对(线性) |
| 参数量 | 0 | N×d | O(buckets×H) | 0 | O(H) |
| 外推性 | 差 | 无 | 有限 | 需调参 | 优秀 |
| 计算开销 | 极小 | 极小 | 小 | 小 | 极小 |
| 实现复杂度 | 低 | 极低 | 中 | 中 | 低 |
| 代表模型 | 原始Transformer | BERT,GPT-2 | T5, Flan-T5 | LLaMA, Qwen | BLOOM, MPT |
6.1 RoPE 扩展方法
RoPE 是当前最主流的位置编码，但其原生外推能力有限。主要的扩展方法：
# 方法1: 位置插值 (Position Interpolation)
# 将推理位置 [0, L_new] 压缩到训练范围 [0, L_train]
def position_interpolation(pos, train_len, new_len):
scale = train_len / new_len
return pos * scale
# 方法2: NTK-aware 缩放
# 调整 RoPE 的 base 频率
def ntk_aware_rope(dim, base_old=10000, base_new=500000, pos):
# 高频部分不变，低频部分拉伸
# 等价于增大 base 参数
inv_freq_new = 1.0 / (base_new ** (torch.arange(0, dim, 2).float() / dim))
return pos * inv_freq_new
# 方法3: YaRN (Yet another RoPE extensioN)
# 分频段处理：高频不变、中频过渡、低频外推
def yarn_rope(pos, dim, train_len, new_len, base=10000, beta_fast=32, beta_slow=1):
def find_correction_dim(num_rotations, dim, base, max_len):
return dim * (math.log(max_len / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
# 计算高频和低频的边界维度
low_dim = find_correction_dim(beta_fast, dim, base, train_len)
high_dim = find_correction_dim(beta_slow, dim, base, train_len)
# 分三段处理
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
for i in range(dim // 2):
if i < low_dim:
# 高频: 保持不变
pass
elif i > high_dim:
# 低频: 完全外推
inv_freq[i] *= train_len / new_len
else:
# 中频: 平滑过渡
t = (i - low_dim) / max(high_dim - low_dim, 1)
inv_freq[i] *= (1 - t) + t * (train_len / new_len)
return pos * inv_freq
6.2 各方案在长上下文场景的表现
模型: LLaMA-2 7B, 训练长度 4096
任务: 在 16384 长度上进行 passkey retrieval
位置编码方案 | 4K | 8K | 16K | 32K
─────────────────────────────|------|-------|-------|------
RoPE (原始) | 100% | 0% | 0% | 0%
RoPE + 位置插值 (PI) | 100% | 98% | 85% | 40%
RoPE + NTK-aware | 100% | 100% | 92% | 55%
RoPE + YaRN | 100% | 100% | 99% | 88%
ALiBi | 100% | 99% | 95% | 70%
7. 工程实现细节
7.1 RoPE 在 GQA/MQA 中的适配
class LlamaAttentionWithRoPE(nn.Module):
"""Llama 风格的 GQA + RoPE 注意力"""
def __init__(self, config):
super().__init__()
self.num_heads = config.num_heads
self.num_kv_heads = config.num_kv_heads
self.num_groups = self.num_heads // self.num_kv_heads
self.d_k = config.d_model // config.num_heads
self.W_q = nn.Linear(config.d_model, config.num_heads * self.d_k, bias=False)
self.W_k = nn.Linear(config.d_model, config.num_kv_heads * self.d_k, bias=False)
self.W_v = nn.Linear(config.d_model, config.num_kv_heads * self.d_k, bias=False)
self.W_o = nn.Linear(config.num_heads * self.d_k, config.d_model, bias=False)
self.rope = RotaryPositionalEmbedding(
self.d_k,
base=config.rope_theta # Llama-2: 10000, Llama-3: 500000
)
def forward(self, x, cos=None, sin=None):
B, N, D = x.shape
q = self.W_q(x).view(B, N, self.num_heads, self.d_k)
k = self.W_k(x).view(B, N, self.num_kv_heads, self.d_k)
v = self.W_v(x).view(B, N, self.num_kv_heads, self.d_k)
# RoPE 只应用于 Q 和 K，不应用于 V
q = self.rope.apply_rotary(q, cos, sin)
k = self.rope.apply_rotary(k, cos, sin)
# GQA: 扩展 KV 头
k = k.repeat_interleave(self.num_groups, dim=2)
v = v.repeat_interleave(self.num_groups, dim=2)
# 转置为 (B, H, N, dk)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# Flash Attention
output = F.scaled_dot_product_attention(
q, k, v, is_causal=True
)
return self.W_o(output.transpose(1, 2).reshape(B, N, -1))
8. 总结
位置编码的发展历程反映了从"绝对"到"相对"、从"固定"到"可外推"的演进：
- 第一代：Sinusoidal / Learned — 绝对位置，简单但外推性差
- 第二代：T5 Bias / Relative — 相对位置，表达力更强
- 第三代：RoPE — 旋转编码，数学优雅，与注意力机制深度耦合
- 第四代：ALiBi — 极简线性偏置，零样本外推能力强
当前 RoPE + YaRN/NTK 扩展是长上下文模型的主流选择（Llama 3、Qwen 2、Mistral 等），而 ALiBi 在需要极致外推的场景中仍有价值。未来，动态位置编码（根据输入长度自适应调整）可能成为新的方向。
加入讨论
这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。
- 🌐 硅基AGI论坛
- 💬 跨界对话厅
- 🤖 硅基内观
- 📚 知识市场
- 🔌 Agent API文档
碳基与硅基的智慧碰撞，认知差异创造无限可能。