在上一章中,我们学习了注意力机制如何通过QKV矩阵计算Token之间的相关性。但这里有一个严重的问题:
注意力机制天生是"位置不敏感"的!
考虑以下两个句子:
对于注意力机制来说,如果我们交换Token的顺序,计算过程是这样的:
句子1的注意力分数矩阵:句子2(交换位置后):句子1的注意力分数矩阵:句子2(交换位置后):Scores1=Q1⋅K1TScores2=Q2⋅K2T由于 QQQ、KKK、VVV 都是通过相同的权重矩阵 WQW_QWQ、WKW_KWK、WVW_VWV 从Embedding计算得到的,如果我们只是交换了Token的顺序,而不告诉模型"位置信息",那么注意力机制会认为这两个句子是等价的!
具体来说,注意力计算公式:
Attention(Q,K,V)=softmax(Q⋅KTdk)⋅Vtext{Attention}(Q, K, V) = text{softmax}left(frac{Q cdot K^T}{sqrt{d_k}}right) cdot VAttention(Q,K,V)=softmax(dkQ⋅KT)⋅V这个公式中,没有任何地方体现Token的位置信息。
自然语言中,位置决定语义:
更技术性的原因:
因此,我们必须给模型注入位置信息,这就是位置编码的作用。
位置编码的目标很简单:
数学表达:
输入带位置信息的表示=TokenEmbedding+PositionalEncodingXwith_pos[i]=X[i]+PE[i]begin{aligned}text{输入带位置信息的表示} &= text{Token Embedding} + text{Positional Encoding} X_{text{with_pos}}[i] &= X[i] + text{PE}[i]end{aligned}输入带位置信息的表示Xwith_pos[i]=TokenEmbedding+PositionalEncoding=X[i]+PE[i]其中:
Transformer原始论文(Vaswani et al., 2017)提出了一种基于正弦和余弦函数的位置编码方案。
对于位置 postext{pos}pos(第几个Token,从0开始)和维度 iii(向量的第几维,从0开始):
PE(pos,2i)=sin(pos100002i/dmodel)PE(pos,2i+1)=cos(pos100002i/dmodel)begin{aligned}text{PE}(text{pos}, 2i) &= sinleft(frac{text{pos}}{10000^{2i/d_{text{model}}}}right) text{PE}(text{pos}, 2i+1) &= cosleft(frac{text{pos}}{10000^{2i/d_{text{model}}}}right)end{aligned}PE(pos,2i)PE(pos,2i+1)=sin(100002i/dmodelpos)=cos(100002i/dmodelpos)参数解释:
这个公式的核心思想是:使用不同频率的正弦波来编码位置
想象一下时钟:
不同的时刻,秒针、分针、时针的组合是唯一的,这就能唯一标识一个时间点。
类似地:
假设 dmodel=4d_{text{model}} = 4dmodel=4(简化),我们计算前3个位置的位置编码:
位置 pos=0:
PE(0,0)=sin(0/100000/4)=sin(0)=0PE(0,1)=cos(0/100000/4)=cos(0)=1PE(0,2)=sin(0/100002/4)=sin(0)=0PE(0,3)=cos(0/100002/4)=cos(0)=1PE[0]=[0,1,0,1]begin{aligned}text{PE}(0, 0) &= sin(0 / 10000^{0/4}) = sin(0) = 0 text{PE}(0, 1) &= cos(0 / 10000^{0/4}) = cos(0) = 1 text{PE}(0, 2) &= sin(0 / 10000^{2/4}) = sin(0) = 0 text{PE}(0, 3) &= cos(0 / 10000^{2/4}) = cos(0) = 1 \text{PE}[0] &= [0, 1, 0, 1]end{aligned}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)=sin(1/100000/4)=sin(1)≈0.841PE(1,1)=cos(1/100000/4)=cos(1)≈0.540PE(1,2)=sin(1/100002/4)=sin(0.01)≈0.01PE(1,3)=cos(1/100002/4)=cos(0.01)≈1.0PE[1]=[0.841,0.540,0.01,1.0]begin{aligned}text{PE}(1, 0) &= sin(1 / 10000^{0/4}) = sin(1) approx 0.841 text{PE}(1, 1) &= cos(1 / 10000^{0/4}) = cos(1) approx 0.540 text{PE}(1, 2) &= sin(1 / 10000^{2/4}) = sin(0.01) approx 0.01 text{PE}(1, 3) &= cos(1 / 10000^{2/4}) = cos(0.01) approx 1.0 \text{PE}[1] &= [0.841, 0.540, 0.01, 1.0]end{aligned}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]text{PE}[2] = [0.909, -0.416, 0.02, 0.9998]PE[2]=[0.909,−0.416,0.02,0.9998]可以看到,每个位置都有一个唯一的向量表示。
另一种简单的方案是:把位置编码当作模型参数,在训练中学习
创建一个可学习的Embedding矩阵:
PElearned∈Rmax_seq_len×dmodeltext{PE}_{text{learned}} in mathbb{R}^{text{max_seq_len} times d_{text{model}}}PElearned∈Rmax_seq_len×dmodel对于位置 postext{pos}pos:
PE[pos]=PElearned[pos]PE[pos]=PElearned[pos](直接查表)参数解释:
优势:
劣势:
这种方案在BERT、GPT等早期模型中使用,但现代大模型更倾向于使用RoPE等相对位置编码。
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⋅KTbegin{aligned}&text{步骤1:在输入阶段加入位置信息} &X_{text{with_pos}} = X + text{PE} \&text{步骤2:计算Q、K、V} &Q = X_{text{with_pos}} cdot W_Q &K = X_{text{with_pos}} cdot W_K &V = X_{text{with_pos}} cdot W_V \&text{步骤3:计算注意力分数} &text{Score} = Q cdot K^Tend{aligned}步骤1:在输入阶段加入位置信息Xwith_pos=X+PE步骤2:计算Q、K、VQ=Xwith_pos⋅WQK=Xwith_pos⋅WKV=Xwith_pos⋅WV步骤3:计算注意力分数Score=Q⋅KTRoPE方法:
步骤1:先计算Q、K(不含位置信息)Q=X⋅WQK=X⋅WKV=X⋅WV步骤2:对Q、K应用旋转矩阵(注入位置信息)Qwith_pos[m]=RΘ(m)⋅Q[m]Kwith_pos[n]=RΘ(n)⋅K[n]步骤3:计算注意力分数Score(m,n)=Qwith_pos[m]⋅Kwith_pos[n]Tbegin{aligned}&text{步骤1:先计算Q、K(不含位置信息)} &Q = X cdot W_Q &K = X cdot W_K &V = X cdot W_V \&text{步骤2:对Q、K应用旋转矩阵(注入位置信息)} &Q_{text{with_pos}}[m] = R_Theta(m) cdot Q[m] quad text{(位置m的Query向量)} &K_{text{with_pos}}[n] = R_Theta(n) cdot K[n] quad text{(位置n的Key向量)} \&text{步骤3:计算注意力分数} &text{Score}(m,n) = Q_{text{with_pos}}[m] cdot K_{text{with_pos}}[n]^Tend{aligned}步骤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关键区别:
为什么这样更好?
在二维平面上,旋转一个向量 θthetaθ 角度,可以用旋转矩阵表示:
[x′y′]=[cosθ−sinθsinθcosθ][xy]begin{bmatrix} x' y' end{bmatrix} = begin{bmatrix} costheta & -sintheta sintheta & costheta end{bmatrix} begin{bmatrix} x y end{bmatrix}[x′y′]=[cosθsinθ−sinθcosθ][xy]RoPE就是将这个思想推广到高维空间:每对维度作为一个平面,进行不同角度的旋转
对于位置 mmm 的Query向量和位置 nnn 的Key向量,RoPE将它们分别旋转:
qm=RΘ(m)⋅WQ⋅xmkn=RΘ(n)⋅WK⋅xnbegin{aligned}q_m &= R_Theta(m) cdot W_Q cdot x_m k_n &= R_Theta(n) cdot W_K cdot x_nend{aligned}qmkn=RΘ(m)⋅WQ⋅xm=RΘ(n)⋅WK⋅xn其中,旋转矩阵 RΘ(pos)R_Theta(text{pos})RΘ(pos) 是一个分块对角矩阵:
RΘ(pos)=[cos(pos⋅θ0)−sin(pos⋅θ0)00⋯sin(pos⋅θ0)cos(pos⋅θ0)00⋯00cos(pos⋅θ1)−sin(pos⋅θ1)⋯00sin(pos⋅θ1)cos(pos⋅θ1)⋯⋮⋮⋮⋮⋱]R_Theta(text{pos}) = begin{bmatrix}cos(text{pos} cdot theta_0) & -sin(text{pos} cdot theta_0) & 0 & 0 & cdots sin(text{pos} cdot theta_0) & cos(text{pos} cdot theta_0) & 0 & 0 & cdots & 0 & cos(text{pos} cdot theta_1) & -sin(text{pos} cdot theta_1) & cdots & 0 & sin(text{pos} cdot theta_1) & cos(text{pos} cdot theta_1) & cdots vdots & vdots & vdots & vdots & ddotsend{bmatrix}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)⋮⋯⋯⋯⋯⋱参数解释:
为了更直观,我们可以用逐元素的方式表示RoPE:
对于Query向量的第 2i2i2i 和 2i+12i+12i+1 维(一对维度):
qm[2i]=q[2i]⋅cos(m⋅θi)−q[2i+1]⋅sin(m⋅θi)qm[2i+1]=q[2i]⋅sin(m⋅θi)+q[2i+1]⋅cos(m⋅θi)begin{aligned}q_m[2i] &= q[2i] cdot cos(m cdot theta_i) - q[2i+1] cdot sin(m cdot theta_i) q_m[2i+1] &= q[2i] cdot sin(m cdot theta_i) + q[2i+1] cdot cos(m cdot theta_i)end{aligned}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]=k[2i]⋅cos(n⋅θi)−k[2i+1]⋅sin(n⋅θi)kn[2i+1]=k[2i]⋅sin(n⋅θi)+k[2i+1]⋅cos(n⋅θi)begin{aligned}k_n[2i] &= k[2i] cdot cos(n cdot theta_i) - k[2i+1] cdot sin(n cdot theta_i) k_n[2i+1] &= k[2i] cdot sin(n cdot theta_i) + k[2i+1] cdot cos(n cdot theta_i)end{aligned}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/dmodeltheta_i = 10000^{-2i/d_{text{model}}}θi=10000−2i/dmodel让我们详细看看RoPE是如何一步步融合进注意力机制的计算过程的。
假设我们有一个序列:["我", "喜欢", "猫"],共3个Token,位置分别为0、1、2。
其中 x0x_0x0、x1x_1x1、x2x_2x2 分别是"我"、"喜欢"、"猫"的Embedding向量。
注意:到这一步为止,Q、K、V都还没有任何位置信息!
这是RoPE的核心步骤!对每个位置的Q和K向量应用旋转:
位置0的Token:"我"位置1的Token:"喜欢"位置2的Token:"猫"位置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)⋅k2V向量不旋转,因为V包含的是"内容信息",只有Q和K需要位置信息来计算相关性。
现在计算所有位置对之间的注意力分数:
Scores=Qrot⋅KrotT=[Qrot[0]⋅Krot[0]TQrot[0]⋅Krot[1]TQrot[0]⋅Krot[2]TQrot[1]⋅Krot[0]TQrot[1]⋅Krot[1]TQrot[1]⋅Krot[2]TQrot[2]⋅Krot[0]TQrot[2]⋅Krot[1]TQrot[2]⋅Krot[2]T]text{Scores} = Q_{text{rot}} cdot K_{text{rot}}^T = begin{bmatrix}Q_{text{rot}}[0] cdot K_{text{rot}}[0]^T & Q_{text{rot}}[0] cdot K_{text{rot}}[1]^T & Q_{text{rot}}[0] cdot K_{text{rot}}[2]^T Q_{text{rot}}[1] cdot K_{text{rot}}[0]^T & Q_{text{rot}}[1] cdot K_{text{rot}}[1]^T & Q_{text{rot}}[1] cdot K_{text{rot}}[2]^T Q_{text{rot}}[2] cdot K_{text{rot}}[0]^T & Q_{text{rot}}[2] cdot K_{text{rot}}[1]^T & Q_{text{rot}}[2] cdot K_{text{rot}}[2]^Tend{bmatrix}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]TQ_{text{rot}}[m] cdot K_{text{rot}}[n]^TQrot[m]⋅Krot[n]T 自动包含了位置m和位置n之间的相对位置信息 (m−n)(m-n)(m−n)!
| 步骤 | 传统位置编码 | RoPE |
|---|---|---|
| 1. 输入 | X+PEX + text{PE}X+PE | XXX(纯内容) |
| 2. 计算QKV | Q=(X+PE)⋅WQQ = (X + text{PE}) cdot W_QQ=(X+PE)⋅WQ | Q=X⋅WQQ = X cdot W_QQ=X⋅WQ |
| 3. 位置注入 | ❌(已在步骤1完成) | ✅ Qrot=RΘ(pos)⋅QQ_{text{rot}} = R_Theta(text{pos}) cdot QQrot=RΘ(pos)⋅Q |
| 4. 注意力分数 | Q⋅KTQ cdot K^TQ⋅KT(位置信息已稀释) | Qrot⋅KrotTQ_{text{rot}} cdot K_{text{rot}}^TQrot⋅KrotT(位置信息精确) |
| 结果 | 位置信息间接、可能被削弱 | 位置信息直接、保留相对关系 |
核心优势: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)begin{aligned}text{Attention_Score}(m, n) &= q_m^T cdot k_n &= text{(展开后,对于第i对维度)} &= left[q[2i] cdot k[2i] + q[2i+1] cdot k[2i+1]right] times cos((m-n) cdot theta_i)end{aligned}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)(m−n),即相对位置差,而不是绝对位置m或n!
这意味着:
性质2:远距离衰减
由于使用了不同频率的旋转:
相对位置距离越远,高频分量的点积越接近0,注意力自然衰减。
假设 dmodel=4d_{text{model}} = 4dmodel=4,我们计算位置0和位置1的Query向量:
步骤1:计算旋转频率
θ0=10000−0/4=1θ1=10000−2/4=0.01begin{aligned}theta_0 &= 10000^{-0/4} = 1 theta_1 &= 10000^{-2/4} = 0.01end{aligned}θ0θ1=10000−0/4=1=10000−2/4=0.01步骤2:对位置m=0,旋转角度为0
q0[0]=q[0]⋅cos(0⋅1)−q[1]⋅sin(0⋅1)=q[0]q0[1]=q[0]⋅sin(0⋅1)+q[1]⋅cos(0⋅1)=q[1]q0[2]=q[2]⋅cos(0⋅0.01)−q[3]⋅sin(0⋅0.01)=q[2]q0[3]=q[2]⋅sin(0⋅0.01)+q[3]⋅cos(0⋅0.01)=q[3]begin{aligned}q_0[0] &= q[0] cdot cos(0 cdot 1) - q[1] cdot sin(0 cdot 1) = q[0] q_0[1] &= q[0] cdot sin(0 cdot 1) + q[1] cdot cos(0 cdot 1) = q[1] q_0[2] &= q[2] cdot cos(0 cdot 0.01) - q[3] cdot sin(0 cdot 0.01) = q[2] q_0[3] &= q[2] cdot sin(0 cdot 0.01) + q[3] cdot cos(0 cdot 0.01) = q[3]end{aligned}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]=q[0]⋅cos(1⋅1)−q[1]⋅sin(1⋅1)≈0.540⋅q[0]−0.841⋅q[1]q1[1]=q[0]⋅sin(1⋅1)+q[1]⋅cos(1⋅1)≈0.841⋅q[0]+0.540⋅q[1]q1[2]=q[2]⋅cos(1⋅0.01)−q[3]⋅sin(1⋅0.01)≈0.99995⋅q[2]−0.01⋅q[3]q1[3]=q[2]⋅sin(1⋅0.01)+q[3]⋅cos(1⋅0.01)≈0.01⋅q[2]+0.99995⋅q[3]begin{aligned}q_1[0] &= q[0] cdot cos(1 cdot 1) - q[1] cdot sin(1 cdot 1) approx 0.540 cdot q[0] - 0.841 cdot q[1] q_1[1] &= q[0] cdot sin(1 cdot 1) + q[1] cdot cos(1 cdot 1) approx 0.841 cdot q[0] + 0.540 cdot q[1] q_1[2] &= q[2] cdot cos(1 cdot 0.01) - q[3] cdot sin(1 cdot 0.01) approx 0.99995 cdot q[2] - 0.01 cdot q[3] q_1[3] &= q[2] cdot sin(1 cdot 0.01) + q[3] cdot cos(1 cdot 0.01) approx 0.01 cdot q[2] + 0.99995 cdot q[3]end{aligned}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]可以看到:
在实际代码中,RoPE通常这样实现。下面展示完整的带RoPE的注意力计算流程:
import torchimport torch.nn.functional as F# ============ 第一步:预计算RoPE的旋转矩阵(初始化时执行一次) ============def precompute_rope_cache(d_model, max_seq_len=2048):"""预计算RoPE需要的cos和sin值"""# 计算旋转频率 θ_i = 10000^(-2i/d_model)theta = 10000 ** (-2 * torch.arange(d_model // 2) / d_model)# theta shape: (d_model/2,)# 生成位置索引 [0, 1, 2, ..., max_seq_len-1]pos = torch.arange(max_seq_len)# pos shape: (max_seq_len,)# 计算所有位置和所有频率的组合:pos * θ_ifreqs = torch.outer(pos, theta)# shape: (max_seq_len, d_model/2)# 预计算cos和sin值,推理时直接查表cos_cache = freqs.cos()# shape: (max_seq_len, d_model/2)sin_cache = freqs.sin()# shape: (max_seq_len, d_model/2)return cos_cache, sin_cache# ============ 第二步:应用RoPE旋转(在每次forward时执行) ============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)"""# 将x分为偶数维和奇数维x1 = x[..., 0::2]# shape: (batch, seq_len, d_model/2) - 第0,2,4,...维x2 = x[..., 1::2]# shape: (batch, seq_len, d_model/2) - 第1,3,5,...维# 应用旋转公式:# x_rot[2i] = x[2i] * cos(pos*θ_i) - x[2i+1] * sin(pos*θ_i)# x_rot[2i+1] = x[2i] * sin(pos*θ_i) + x[2i+1] * cos(pos*θ_i)x_rotated = torch.stack([x1 * cos - x2 * sin,# 偶数维x1 * sin + x2 * cos # 奇数维], dim=-1).flatten(-2)# 交错拼接回 (batch, seq_len, d_model)return x_rotated# ============ 第三步:完整的带RoPE的注意力计算 ============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# 步骤1:计算原始的Q、K、V(不含位置信息)Q = torch.matmul(X, W_Q)# shape: (batch, seq_len, d_k)K = torch.matmul(X, W_K)# shape: (batch, seq_len, d_k)V = torch.matmul(X, W_V)# shape: (batch, seq_len, d_v)print("步骤1完成:计算Q、K、V(纯内容,无位置信息)")# 步骤2:对Q、K应用RoPE旋转(注入位置信息)# 这是RoPE的核心!位置信息在这里融入cos = cos_cache[:seq_len]# 截取当前序列长度sin = sin_cache[:seq_len]Q_rot = apply_rope(Q, cos, sin)# shape: (batch, seq_len, d_k)K_rot = apply_rope(K, cos, sin)# shape: (batch, seq_len, d_k)# 注意:V不旋转!V只包含内容信息print("步骤2完成:对Q、K应用旋转(位置信息已注入)")# 步骤3:计算注意力分数(位置信息已在Q_rot和K_rot中)d_k = Q_rot.shape[-1]scores = torch.matmul(Q_rot, K_rot.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))# scores shape: (batch, seq_len, seq_len)print("步骤3完成:计算注意力分数(自动包含相对位置信息)")# 步骤4:Softmax + 加权求和(标准流程)attn_weights = F.softmax(scores, dim=-1)output = torch.matmul(attn_weights, V)# shape: (batch, seq_len, d_v)print("步骤4完成:加权求和得到输出")return output, attn_weights# ============ 使用示例 ============# 初始化(只需一次)d_model = 512max_seq_len = 2048cos_cache, sin_cache = precompute_rope_cache(d_model, max_seq_len)# 前向传播(每次推理)batch_size = 2seq_len = 10X = torch.randn(batch_size, seq_len, d_model)# 输入Token Embeddings# 假设已初始化权重矩阵W_Q = torch.randn(d_model, d_model)W_K = torch.randn(d_model, d_model)W_V = torch.randn(d_model, d_model)# 执行注意力计算(RoPE在步骤2自动应用)output, attn_weights = attention_with_rope(X, W_Q, W_K, W_V, cos_cache, sin_cache)print(f"n最终输出 shape: {output.shape}")# (batch_size, seq_len, d_model)
代码关键点解释:
预计算阶段(precompute_rope_cache):
RoPE应用阶段(apply_rope):
融合进注意力计算(attention_with_rope):
与传统方法的对比:
| 时间点 | 传统位置编码 | 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理论上可以外推,但在实际应用中,当序列长度远超训练时的长度时,性能会下降。为此,研究者提出了多种长度扩展技术。
假设模型在训练时只见过长度≤2048的序列,当推理时输入长度4096的序列:
核心思想:将长序列的位置"压缩"到训练时的范围内
原始位置:压缩后:′=pos⋅LtrainLnewbegin{aligned}text{原始位置:} & quad text{pos} in [0, L_{text{new}}] text{压缩后:} & quad text{pos}' = text{pos} cdot frac{L_{text{train}}}{L_{text{new}}}end{aligned}原始位置:压缩后:pos∈[0,Lnew]pos′=pos⋅LnewLtrain参数解释:
举例:
优势:
劣势:
核心思想:不是简单压缩位置,而是调整旋转频率的基数(将10000改为更大的值)
原始频率:NTK频率:′=(10000⋅scale)−2i/dmodelbegin{aligned}text{原始频率:} & quad theta_i = 10000^{-2i/d_{text{model}}} text{NTK频率:} & quad theta_i' = (10000 cdot text{scale})^{-2i/d_{text{model}}}end{aligned}原始频率:NTK频率:θi=10000−2i/dmodelθi′=(10000⋅scale)−2i/dmodel其中:
scale=LnewLtraintext{scale} = frac{L_{text{new}}}{L_{text{train}}}scale=LtrainLnew参数解释:
举例:
优势:
劣势:
核心思想:对不同频率分量采用不同的插值策略
这种方法在LLaMA-2等模型中取得了很好的效果,可以将上下文长度扩展到32k甚至更长。
| 方法 | 是否需要微调 | 外推效果 | 计算开销 |
|---|---|---|---|
| 位置插值(PI) | 需要少量微调 | 好 | 无额外开销 |
| NTK-Aware | 零样本 | 较好 | 无额外开销 |
| YaRN | 零样本或少量微调 | 很好 | 无额外开销 |
位置编码的必要性:注意力机制天生无法感知位置,必须显式注入位置信息
传统位置编码:
RoPE(旋转位置编码):
长度扩展技术:
位置编码看似简单,但对大模型的性能至关重要。RoPE的成功说明,好的位置编码应该捕捉相对位置关系,而不是绝对位置,这样才能具备良好的泛化能力。