一文看懂 LLaMA 中的旋转式位置编码（Rotary Position Embedding） - 链载Ai

旋转式位置编码（RoPE）最早是论文[1]提出的一种能够将相对位置信息依赖集成到 self-attention 中并提升 transformer 架构性能的位置编码方式。而目前很火的 LLaMA 模型也是采用该位置编码方式。

接下来结合代码和论文来解读一下 RoPE。

基本概念

首先论文中定义一个长度为 N 的输入序列为:

其中 wi 表示输入序列中第 i 个 token，而输入序列 SN 对应的 embedding 表示为：

其中 xi 表示第 i 个 token wi 对应的 d 维词嵌入向量。

接着在做 self-attention 之前，会用词嵌入向量计算 q, k, v 向量同时加入位置信息，函数公式表达如下：

其中 qm 表示第 m 个 token 对应的词向量 xm 集成位置信息 m 之后的 query 向量。而 kn 和 vn 则表示第 n 个 token 对应的词向量 xn 集成位置信息 n 之后的 key 和 value 向量。

而基于 transformer 的位置编码方法都是着重于构造一个合适的 f{q,k,v} 函数形式。

而计算第 m 个词嵌入向量 xm 对应的 self-attention 输出结果，就是 qm 和其他 kn 都计算一个 attention score ，然后再将 attention score 乘以对应的 vn 再求和得到输出向量 om：

绝对位置编码

对于位置编码，常规的做法是在计算 query, key 和 value 向量之前，会计算一个位置编码向量 pi 加到词嵌入 xi 上，位置编码向量 pi 同样也是 d 维向量，然后再乘以对应的变换矩阵 W{q,k,v}：

其中 p_{i,2t} 表示位置 d 维度向量 pi 中的第 2t 个元素也就是偶数索引位置的计算公式，而 p_{i,2t+1} 就对应奇数索引位置的计算公式。

python 代码如下：

#position就对应token序列中的位置索引i
#hidden_dim就对应词嵌入维度大小d
#seq_len表示token序列长度
defget_position_angle_vec(position):
return[position/np.power(10000,2*(hid_j//2)/hidden_dim)forhid_jinrange(hidden_dim)]

#position_angle_vecs.shape=[seq_len,hidden_dim]
position_angle_vecs=np.array([get_position_angle_vec(pos_i)forpos_iinrange(seq_len)])

#分别计算奇偶索引位置对应的sin和cos值
position_angle_vecs[:,0::2]=np.sin(position_angle_vecs[:,0::2])#dim2t
position_angle_vecs[:,1::2]=np.cos(position_angle_vecs[:,1::2])#dim2t+1

#positional_embeddings.shape=[1,seq_len,hidden_dim]
positional_embeddings=torch.FloatTensor(position_angle_vecs).unsqueeze(0)

旋转式位置编码

接着论文中提出为了能利用上 token 之间的相对位置信息，假定 query 向量 qm 和 key 向量 kn 之间的内积操作可以被一个函数 g 表示，该函数 g 的输入是词嵌入向量 xm ， xn 和它们之间的相对位置 m - n：