注意力机制的"位置盲区"
在上一章中,我们学习了注意力机制如何通过QKV矩阵计算Token之间的相关性。但这里有一个严重的问题:
注意力机制天生是"位置不敏感"的!
问题演示
考虑以下两个句子:
- "猫 吃 鱼"
- "鱼 吃 猫"
对于注意力机制来说,如果我们交换Token的顺序,计算过程是这样的:
句子1的注意力分数矩阵:句子2(交换位置后):Scores1=Q1⋅K1TScores2=Q2⋅K2T
由于 Q、K、V 都是通过相同的权重矩阵 WQ、WK、WV 从Embedding计算得到的,如果我们只是交换了Token的顺序,而不告诉模型"位置信息",那么注意力机制会认为这两个句子是等价的!
具体来说,注意力计算公式:
Attention(Q,K,V)=softmax(dkQ⋅KT)⋅V
这个公式中,没有任何地方体现Token的位置信息。
为什么位置很重要?
自然语言中,位置决定语义:
- "我 不 喜欢 你" vs "我 喜欢 你 不?"(语义完全不同)
- "小明 打 了 小红" vs "小红 打 了 小明"(主宾关系颠倒)
- "因为下雨,所以取消" vs "取消,所以因为下雨"(因果关系混乱)
更技术性的原因:
- 语法结构:主语在前、谓语在中、宾语在后
- 时间顺序:事件发生的先后顺序
- 依赖关系:前面的Token被后面的Token引用(代词指代)
- 自回归生成:生成第n+1个Token时,只能看前n个Token,不能看"未来"
因此,我们必须给模型注入位置信息,这就是位置编码的作用。
位置编码的核心思想
位置编码的目标很简单:
为每个位置生成一个唯一的向量,并将其加到Token的Embedding上,让模型知道"这个Token在序列中的第几个位置"
数学表达:
输入带位置信息的表示Xwith_pos[i]=Token Embedding+Positional Encoding=X[i]+PE[i]
其中:
- X[i] 是第i个Token的原始Embedding(dmodel维)
- PE[i] 是第i个位置的位置编码向量(同样是dmodel维)
- Xwith_pos[i] 是最终输入到注意力层的表示
原始位置编码(Sinusoidal Positional Encoding)
Transformer原始论文(Vaswani et al., 2017)提出了一种基于正弦和余弦函数的位置编码方案。
公式
对于位置 pos(第几个Token,从0开始)和维度 i(向量的第几维,从0开始):
PE(pos,2i)PE(pos,2i+1)=sin(100002i/dmodelpos)=cos(100002i/dmodelpos)
参数解释:
- pos:Token的位置索引(0, 1, 2, 3, ...)
- i:位置编码向量的维度索引(0,1,2,…,dmodel/2−1)
- 2i:偶数维度使用sin函数
- 2i+1:奇数维度使用cos函数
- 10000:基数,控制频率的衰减速度
- dmodel:位置编码向量的维度(与Token Embedding维度相同)
直观理解
这个公式的核心思想是:使用不同频率的正弦波来编码位置
想象一下时钟:
- 秒针:转得很快,频率高(对应高维度,2i很大)
- 分针:转得中等,频率中等
- 时针:转得很慢,频率低(对应低维度,2i很小)
不同的时刻,秒针、分针、时针的组合是唯一的,这就能唯一标识一个时间点。
类似地:
- 低维度(2i小):使用低频正弦波,变化慢,能区分远距离的位置
- 高维度(2i大):使用高频正弦波,变化快,能区分近距离的位置
具体例子
假设 dmodel=4(简化),我们计算前3个位置的位置编码:
位置 pos=0:
PE(0,0)PE(0,1)PE(0,2)PE(0,3)PE[0]=sin(0/100000/4)=sin(0)=0=cos(0/100000/4)=cos(0)=1=sin(0/100002/4)=sin(0)=0=cos(0/100002/4)=cos(0)=1=[0,1,0,1]
位置 pos=1:
PE(1,0)PE(1,1)PE(1,2)PE(1,3)PE[1]=sin(1/100000/4)=sin(1)≈0.841=cos(1/100000/4)=cos(1)≈0.540=sin(1/100002/4)=sin(0.01)≈0.01=cos(1/100002/4)=cos(0.01)≈1.0=[0.841,0.540,0.01,1.0]
位置 pos=2:
PE[2]=[0.909,−0.416,0.02,0.9998]
可以看到,每个位置都有一个唯一的向量表示。
Sinusoidal编码的优势
- 确定性:不需要学习,直接用公式计算
- 外推性:即使训练时只见过长度100的序列,也能为长度200的序列生成位置编码
- 相对位置:通过三角函数的性质,模型可以学习相对位置关系
- 数学性质:PE(pos+k) 可以表示为 PE(pos) 的线性变换
Sinusoidal编码的劣势
- 外推性有限:虽然理论上可以外推,但实际效果在超长序列上会下降
- 位置信息弱:只是简单的加法,位置信息容易被Embedding淹没
- 无法区分绝对位置重要性:远距离和近距离的位置用相同的编码方式
可学习的绝对位置编码(Learned Positional Encoding)
另一种简单的方案是:把位置编码当作模型参数,在训练中学习
实现方式
创建一个可学习的Embedding矩阵:
PElearned∈Rmax_seq_len×dmodel
对于位置 pos:
PE[pos]=PElearned[pos](直接查表)
参数解释:
- max_seq_len:支持的最大序列长度(比如512、2048等)
- PElearned 是一个可训练的参数矩阵
- 训练时,这个矩阵会通过反向传播更新
优势与劣势
优势:
- 灵活性高,模型可以自己学习最优的位置表示
- 实现简单
劣势:
- 无法外推:如果训练时最大长度是512,那么无法处理长度超过512的序列
- 参数量增加:需要额外存储 max_seq_len×dmodel 个参数
这种方案在BERT、GPT等早期模型中使用,但现代大模型更倾向于使用RoPE等相对位置编码。
RoPE:旋转位置编码(Rotary Position Embedding)
RoPE(Su et al., 2021)是目前最流行的位置编码方案之一,被LLaMA、GPT-NeoX、PaLM等主流大模型采用。
核心思想:直接作用于注意力计算
RoPE与传统位置编码的最大区别在于:它不是在输入阶段添加位置信息,而是直接作用在注意力机制的计算过程中。
传统方法(Sinusoidal、Learned PE):
步骤1:在输入阶段加入位置信息Xwith_pos=X+PE步骤2:计算Q、K、VQ=Xwith_pos⋅WQK=Xwith_pos⋅WKV=Xwith_pos⋅WV步骤3:计算注意力分数Score=Q⋅KT
RoPE方法:
步骤1:先计算Q、K(不含位置信息)Q=X⋅WQK=X⋅WKV=X⋅WV步骤2:对Q、K应用旋转矩阵(注入位置信息)Qwith_pos[m]=RΘ(m)⋅Q[m](位置m的Query向量)Kwith_pos[n]=RΘ(n)⋅K[n](位置n的Key向量)步骤3:计算注意力分数Score(m,n)=Qwith_pos[m]⋅Kwith_pos[n]T
关键区别:
- 传统方法:位置信息通过加法混入Embedding,然后一起参与Q、K、V的线性变换
- RoPE:位置信息通过旋转变换直接作用在已计算好的Q、K上,在注意力分数计算时才引入位置信息
为什么这样更好?
- 传统加法:位置信息和内容信息在Embedding层混合,不同的线性变换会破坏位置关系
- RoPE旋转:位置信息通过几何旋转方式注入,保持了相对位置的数学性质,使得 Score(m,n) 自然地只依赖相对位置 (m−n)
RoPE的核心:通过旋转矩阵,将位置信息直接融入注意力计算的核心步骤,而不是简单地加在输入上
为什么叫"旋转"?
在二维平面上,旋转一个向量 θ 角度,可以用旋转矩阵表示:
[x′y′]=[cosθsinθ−sinθcosθ][xy]
RoPE就是将这个思想推广到高维空间:每对维度作为一个平面,进行不同角度的旋转
RoPE的数学公式
对于位置 m 的Query向量和位置 n 的Key向量,RoPE将它们分别旋转:
qmkn=RΘ(m)⋅WQ⋅xm=RΘ(n)⋅WK⋅xn
其中,旋转矩阵 RΘ(pos) 是一个分块对角矩阵:
RΘ(pos)=cos(pos⋅θ0)sin(pos⋅θ0)00⋮−sin(pos⋅θ0)cos(pos⋅θ0)00⋮00cos(pos⋅θ1)sin(pos⋅θ1)⋮00−sin(pos⋅θ1)cos(pos⋅θ1)⋮⋯⋯⋯⋯⋱
参数解释:
- pos:Token的位置(0, 1, 2, ...)
- θi:第i对维度的旋转频率,计算方式:θi=10000−2i/dmodel
- 每两个维度一组,共 dmodel/2 组
- 每组使用不同的旋转频率 θi
简化表示(逐元素形式)
为了更直观,我们可以用逐元素的方式表示RoPE:
对于Query向量的第 2i 和 2i+1 维(一对维度):
qm[2i]qm[2i+1]=q[2i]⋅cos(m⋅θi)−q[2i+1]⋅sin(m⋅θi)=q[2i]⋅sin(m⋅θi)+q[2i+1]⋅cos(m⋅θi)
对于Key向量同理:
kn[2i]kn[2i+1]=k[2i]⋅cos(n⋅θi)−k[2i+1]⋅sin(n⋅θi)=k[2i]⋅sin(n⋅θi)+k[2i+1]⋅cos(n⋅θi)
其中:
θi=10000−2i/dmodel
RoPE融合进注意力计算的完整流程
让我们详细看看RoPE是如何一步步融合进注意力机制的计算过程的。
假设我们有一个序列:["我", "喜欢", "猫"],共3个Token,位置分别为0、1、2。
步骤1:获取Token Embedding(不含位置信息)
X=x0x1x2∈R3×dmodel
其中 x0、x1、x2 分别是"我"、"喜欢"、"猫"的Embedding向量。
步骤2:计算原始的Q、K、V(仍不含位置信息)
QKV=X⋅WQ=q0q1q2∈R3×dk=X⋅WK=k0k1k2∈R3×dk=X⋅WV=v0v1v2∈R3×dv
注意:到这一步为止,Q、K、V都还没有任何位置信息!
步骤3:对Q、K应用RoPE旋转(注入位置信息)
这是RoPE的核心步骤!对每个位置的Q和K向量应用旋转:
位置0的Token:"我"Qrot[0]=RΘ(0)⋅q0(旋转0度,保持不变)Krot[0]=RΘ(0)⋅k0位置1的Token:"喜欢"Qrot[1]=RΘ(1)⋅q1(旋转θ角度)Krot[1]=RΘ(1)⋅k1位置2的Token:"猫"Qrot[2]=RΘ(2)⋅q2(旋转2θ角度)Krot[2]=RΘ(2)⋅k2
V向量不旋转,因为V包含的是"内容信息",只有Q和K需要位置信息来计算相关性。
步骤4:计算注意力分数矩阵(位置信息已融合)
现在计算所有位置对之间的注意力分数:
Scores=Qrot⋅KrotT=Qrot[0]⋅Krot[0]TQrot[1]⋅Krot[0]TQrot[2]⋅Krot[0]TQrot[0]⋅Krot[1]TQrot[1]⋅Krot[1]TQrot[2]⋅Krot[1]TQrot[0]⋅Krot[2]TQrot[1]⋅Krot[2]TQrot[2]⋅Krot[2]T
关键:每个分数 Qrot[m]⋅Krot[n]T 自动包含了位置m和位置n之间的相对位置信息 (m−n)!
步骤5:应用Softmax和加权求和(标准流程)
Attention_WeightsOutput=softmax(dkScores)=Attention_Weights⋅V
对比总结:RoPE vs 传统方法
| 步骤 | 传统位置编码 | RoPE |
|---|
| 1. 输入 | X+PE | X(纯内容) |
| 2. 计算QKV | Q=(X+PE)⋅WQ | Q=X⋅WQ |
| 3. 位置注入 | ❌(已在步骤1完成) | ✅ Qrot=RΘ(pos)⋅Q |
| 4. 注意力分数 | Q⋅KT(位置信息已稀释) | Qrot⋅KrotT(位置信息精确) |
| 结果 | 位置信息间接、可能被削弱 | 位置信息直接、保留相对关系 |
核心优势:RoPE在注意力分数计算的关键时刻才引入位置信息,通过旋转的几何性质,保证了注意力分数只依赖相对位置差,而不是绝对位置。
RoPE的关键性质
性质1:注意力分数自动包含相对位置信息
当计算位置m和位置n之间的注意力分数时:
Attention_Score(m,n)=qmT⋅kn=(展开后,对于第i对维度)=[q[2i]⋅k[2i]+q[2i+1]⋅k[2i+1]]×cos((m−n)⋅θi)
核心发现:注意力分数只依赖于 (m−n),即相对位置差,而不是绝对位置m或n!
这意味着:
- 位置0的Token看位置1的Token,与位置5的Token看位置6的Token,注意力模式相同(相对位置都是+1)
- 模型自然地学习到相对位置关系
性质2:远距离衰减
由于使用了不同频率的旋转:
- 低频分量(θi 小):捕捉长距离依赖
- 高频分量(θi 大):捕捉短距离依赖
相对位置距离越远,高频分量的点积越接近0,注意力自然衰减。
具体例子
假设 dmodel=4,我们计算位置0和位置1的Query向量:
步骤1:计算旋转频率
θ0θ1=10000−0/4=1=10000−2/4=0.01
步骤2:对位置m=0,旋转角度为0
q0[0]q0[1]q0[2]q0[3]=q[0]⋅cos(0⋅1)−q[1]⋅sin(0⋅1)=q[0]=q[0]⋅sin(0⋅1)+q[1]⋅cos(0⋅1)=q[1]=q[2]⋅cos(0⋅0.01)−q[3]⋅sin(0⋅0.01)=q[2]=q[2]⋅sin(0⋅0.01)+q[3]⋅cos(0⋅0.01)=q[3]
位置0不旋转,保持原样。
步骤3:对位置m=1,旋转角度为θ
q1[0]q1[1]q1[2]q1[3]=q[0]⋅cos(1⋅1)−q[1]⋅sin(1⋅1)≈0.540⋅q[0]−0.841⋅q[1]=q[0]⋅sin(1⋅1)+q[1]⋅cos(1⋅1)≈0.841⋅q[0]+0.540⋅q[1]=q[2]⋅cos(1⋅0.01)−q[3]⋅sin(1⋅0.01)≈0.99995⋅q[2]−0.01⋅q[3]=q[2]⋅sin(1⋅0.01)+q[3]⋅cos(1⋅0.01)≈0.01⋅q[2]+0.99995⋅q[3]
可以看到:
- 第0-1维(高频):旋转了约54度,变化明显
- 第2-3维(低频):只旋转了约0.57度,几乎不变
RoPE的优势
- 相对位置编码:注意力分数自动包含相对位置信息,不依赖绝对位置
- 外推性好:理论上可以处理任意长度的序列
- 长距离衰减:远距离Token的注意力自然衰减,符合语言学规律
- 无额外参数:不增加模型参数量
- 高效实现:可以预计算旋转矩阵,推理时直接查表
RoPE的实现细节
在实际代码中,RoPE通常这样实现。下面展示完整的带RoPE的注意力计算流程:
import torch
import torch.nn.functional as F
def precompute_rope_cache(d_model, max_seq_len=2048):
"""
预计算RoPE需要的cos和sin值
"""
theta = 10000 ** (-2 * torch.arange(d_model // 2) / d_model)
pos = torch.arange(max_seq_len)
freqs = torch.outer(pos, theta)
cos_cache = freqs.cos()
sin_cache = freqs.sin()
return cos_cache, sin_cache
def apply_rope(x, cos, sin):
"""
对Q或K向量应用RoPE旋转
Args:
x: shape (batch, seq_len, d_model) - Q或K矩阵
cos: shape (seq_len, d_model/2) - 预计算的cos值
sin: shape (seq_len, d_model/2) - 预计算的sin值
Returns:
旋转后的向量,shape (batch, seq_len, d_model)
"""
x1 = x[..., 0::2]
x2 = x[..., 1::2]
x_rotated = torch.stack([
x1 * cos - x2 * sin,
x1 * sin + x2 * cos
], dim=-1).flatten(-2)
return x_rotated
def attention_with_rope(X, W_Q, W_K, W_V, cos_cache, sin_cache):
"""
完整的注意力计算流程,展示RoPE如何融合进来
Args:
X: shape (batch, seq_len, d_model) - 输入的Token Embeddings(不含位置信息)
W_Q, W_K, W_V: 权重矩阵
cos_cache, sin_cache: 预计算的RoPE缓存
"""
batch, seq_len, d_model = X.shape
Q = torch.matmul(X, W_Q)
K = torch.matmul(X, W_K)
V = torch.matmul(X, W_V)
print("步骤1完成:计算Q、K、V(纯内容,无位置信息)")
cos = cos_cache[:seq_len]
sin = sin_cache[:seq_len]
Q_rot = apply_rope(Q, cos, sin)
K_rot = apply_rope(K, cos, sin)
print("步骤2完成:对Q、K应用旋转(位置信息已注入)")
d_k = Q_rot.shape[-1]
scores = torch.matmul(Q_rot, K_rot.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))
print("步骤3完成:计算注意力分数(自动包含相对位置信息)")
attn_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attn_weights, V)
print("步骤4完成:加权求和得到输出")
return output, attn_weights
d_model = 512
max_seq_len = 2048
cos_cache, sin_cache = precompute_rope_cache(d_model, max_seq_len)
batch_size = 2
seq_len = 10
X = torch.randn(batch_size, seq_len, d_model)
W_Q = torch.randn(d_model, d_model)
W_K = torch.randn(d_model, d_model)
W_V = torch.randn(d_model, d_model)
output, attn_weights = attention_with_rope(X, W_Q, W_K, W_V, cos_cache, sin_cache)
print(f"\n最终输出 shape: {output.shape}")
代码关键点解释:
-
预计算阶段(precompute_rope_cache):
- 只在模型初始化时执行一次
- 计算所有可能位置的 cos(pos⋅θi) 和 sin(pos⋅θi)
- 存储在缓存中,推理时直接查表
-
RoPE应用阶段(apply_rope):
- 在计算完Q、K之后立即应用
- 将向量的每对维度 (2i,2i+1) 作为一个平面进行旋转
- V向量不旋转,因为V存储的是"内容",不需要位置信息
-
融合进注意力计算(attention_with_rope):
- 步骤1:Q=X⋅WQ(不含位置)
- 步骤2:Qrot=RΘ(pos)⋅Q(RoPE在这里注入位置)
- 步骤3:Scores=Qrot⋅KrotT(位置信息自动体现在分数中)
- 步骤4:标准的softmax和加权求和
与传统方法的对比:
| 时间点 | 传统位置编码 | RoPE |
|---|
| 输入阶段 | X = X + PE(位置信息混入) | X(纯内容) |
| 计算QKV | Q = X · W_Q(位置已混入) | Q = X · W_Q(纯内容) |
| 位置注入 | ❌(已完成) | ✅ Q_rot = apply_rope(Q)(在这里!) |
| 计算分数 | Q · K^T | Q_rot · K_rot^T |
RoPE的优势在于:位置信息在注意力分数计算的关键时刻才引入,通过旋转的几何性质精确地编码了相对位置关系。
RoPE的长度扩展技术
虽然RoPE理论上可以外推,但在实际应用中,当序列长度远超训练时的长度时,性能会下降。为此,研究者提出了多种长度扩展技术。
问题:为什么需要长度扩展?
假设模型在训练时只见过长度≤2048的序列,当推理时输入长度4096的序列:
- 位置编码会遇到"未见过"的旋转角度
- 高频分量的旋转角度过大,导致注意力模式混乱
- 模型性能显著下降
方法1:位置插值(Position Interpolation, PI)
核心思想:将长序列的位置"压缩"到训练时的范围内
原始位置:压缩后:pos∈[0,Lnew]pos′=pos⋅LnewLtrain
参数解释:
- Ltrain:训练时的最大序列长度(如2048)
- Lnew:推理时的目标序列长度(如8192)
- pos′:压缩后的位置,范围在 [0,Ltrain] 内
举例:
- 训练长度2048,推理长度8192
- 推理时位置4096 → 压缩为 4096×(2048/8192)=1024
- 推理时位置8192 → 压缩为 8192×(2048/8192)=2048
优势:
- 简单有效,只需修改位置索引
- 所有位置都在训练范围内,模型"见过"
劣势:
- 改变了相对位置的含义(相邻Token的距离变小了)
- 需要少量微调来适应
方法2:NTK-Aware插值
核心思想:不是简单压缩位置,而是调整旋转频率的基数(将10000改为更大的值)
原始频率:NTK频率:θi=10000−2i/dmodelθi′=(10000⋅scale)−2i/dmodel
其中:
scale=LtrainLnew
参数解释:
- scale:长度扩展倍数
- 基数从10000增大到 10000×scale
- 旋转频率整体降低,适应更长的序列
举例:
- 训练长度2048,推理长度8192,scale=4
- 基数从10000变为40000
- 旋转速度降低4倍,适配4倍长的序列
优势:
- 保持了相对位置的语义
- 不需要微调,零样本外推效果好
劣势:
方法3:YaRN(Yet another RoPE extensioN)
核心思想:对不同频率分量采用不同的插值策略
- 低频分量(θ小):捕捉长距离依赖,使用NTK插值
- 高频分量(θ大):捕捉短距离依赖,保持不变或轻微插值
- 中频分量:渐进式插值
这种方法在LLaMA-2等模型中取得了很好的效果,可以将上下文长度扩展到32k甚至更长。
长度扩展对比
| 方法 | 是否需要微调 | 外推效果 | 计算开销 |
|---|
| 位置插值(PI) | 需要少量微调 | 好 | 无额外开销 |
| NTK-Aware | 零样本 | 较好 | 无额外开销 |
| YaRN | 零样本或少量微调 | 很好 | 无额外开销 |
小结
-
位置编码的必要性:注意力机制天生无法感知位置,必须显式注入位置信息
-
传统位置编码:
- Sinusoidal编码:使用sin/cos函数,确定性、可外推,但效果一般
- 可学习编码:灵活但无法外推
-
RoPE(旋转位置编码):
- 通过旋转矩阵将位置信息融入Q、K
- 注意力分数自动包含相对位置信息
- 外推性好、无额外参数、长距离自然衰减
- 成为现代大模型的主流选择
-
长度扩展技术:
- 通过位置插值、频率调整等方法,让模型适应超长序列
- 核心是平衡"训练时的位置模式"和"推理时的长度需求"
位置编码看似简单,但对大模型的性能至关重要。RoPE的成功说明,好的位置编码应该捕捉相对位置关系,而不是绝对位置,这样才能具备良好的泛化能力。