Llama改进之——RoPE旋转位置编码

引言

旋转位置编码(Rotary Position Embedding, RoPE)将绝对相对位置依赖纳入自注意力机制中，以增强Transformer架构的性能。目前很火的大模型LLaMA、QWen等都应用了旋转位置编码。

之前在[论文笔记]ROFORMER中对旋转位置编码的原始论文进行了解析，重点推导了旋转位置编码的公式，本文侧重实现，同时尽量简化数学上的推理，详细内容可见最后的参考文章。

复数与极坐标

复数由两个部分组成：实部(real part)和虚部(imaginary part)。实部就是一个普通的数字，可以是零、正数或负数。虚部是另一个实数与$i$相乘。比如$2+3i$是一个复数，其中$2$是实部；$3i$是虚部。下面这些数字都是复数：
$$2,2+2i,1-3i,-4i,17i$$
可以看到复数是实数的扩展，包含了实数，比如$2$可以看成是虚部为$0$。

通常实数放前面，然后是$i$。但当$i$与三角函数($\sin ,\cos$)在一起通常把$i$放在前面：$i\sin \theta ,i\cos \theta$​​。

$i$我们可以理解为就是一个简单的数学对象，满足$i^2=-1$。

极坐标系是一个二维坐标系统。该坐标系统中任意位置可由一个夹角和一段相对原点——极点的距离来表示。如上图(来自百度百科)所示。

给定极坐标系内的任意一个复数$x+yi$(对应二维向量$[x,y]$)，要将其(逆时针)旋转$\theta$度，只需要乘上旋转子：
$${R\text{R}R}_{\theta }=\cos \theta +i\sin \theta ({\sin }^2\theta +{\cos }^2\theta =1)\text{\hspace{1em}(1)}$$
可以相乘再展开，然后利用$i^2=-1$可得：
$$\begin{array}{rl}{x}^{\prime }+y^{\prime }i& =(\cos \theta +i\sin \theta )(x+yi)\\ & =(x\cos \theta -y\sin \theta )+(x\sin \theta +y\cos \theta )i\end{array}$$
对应二维平面中点$[x,y]$关于原点的逆时针旋转：
$$\left[\begin{array}{c}{x}^{\prime }\\ y^{\prime }\end{array}\right]=\left[\begin{array}{cc}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array}\right]\left[\begin{array}{c}x\\ y\end{array}\right]$$
其中包含$\theta$的矩阵是一个旋转矩阵。

旋转位置编码

${x\text{x}x}_i\in {\mathbb{R}}^d$是无位置信息的标记$w_i$的$d$维词嵌入向量。自注意力首先将位置信息与单词嵌入相结合，并将其转化为query、key和value的表示形式。
$$\begin{array}{rl}{q\text{q}q}_m& =f_q({x\text{x}x}_m,m)\\ {k\text{k}k}_n& =f_k({x\text{x}x}_n,n)\\ {v\text{v}v}_n& =f_v({x\text{x}x}_n,n)\end{array}\text{\hspace{1em}(2)}$$
其中${q\text{q}q}_m,{k\text{k}k}_n$和${v\text{v}v}_n$分别通过$f_q,f_k$和$f_v$整合了第m和第n个位置信息。query和key然后用于计算注意力权重，而输出为value的加权和。
$$
\begin{aligned}
a_{m,n} &= \frac{\exp(\frac{\pmb q^T_m \pmb k_n}{\sqrt d})}{\sum_{j=1}^N \exp \frac{\pmb q^T_m \pmb k_j}{\sqrt d}} \
\pmb o_m &= \sum_{n=1}^N a_{m,n}\pmb v_n \

\end{aligned} \tag 3
$$

Transformer通过自注意机制利用各个标记的位置信息，如等式(3)中所见，${q\text{q}q}_m^T{k\text{k}k}_n$通常可以在不同位置的标记之间传递知识。为了融入相对位置信息，我们需要将查询${q\text{q}q}_m$和键${k\text{k}k}_n$的内积公式转化为一个函数$g$，该函数只接受词嵌入${x\text{x}x}_m,{x\text{x}x}_n$以及它们的相对位置$m-n$​作为输入变量。换句话说，我们希望内积只以相对形式编码位置信息：

$$⟨f_q({x\text{x}x}_m,m),f_k({x\text{x}x}_n,n)⟩=g({x\text{x}x}_m,{x\text{x}x}_n,m-n)\text{\hspace{1em}(4)}$$
最终目标是找到一个等价的编码方式来求解函数$f_q({x\text{x}x}_m,m)$和$f_k({x\text{x}x}_n,n)$​，以符合上等式。

从简单的维度$d=2$的情况开始，这样可以利用二维平面上向量的几何特性及其复数形式来证明公式(4)的一个解是：
$$\begin{array}{rl}{f}_q({x\text{x}x}_m,m)& =({W\text{W}W}_q{x\text{x}x}_m)e^{im\theta }\\ f_k({x\text{x}x}_n,n)& =({W\text{W}W}_k{x\text{x}x}_n)e^{in\theta }\\ g({x\text{x}x}_m,{x\text{x}x}_n,m-n)& =\text{Re}[({W\text{W}W}_q{x\text{x}x}_m)({W\text{W}W}_k{x\text{x}x}_n{)}^{\ast }{e}^{i(m-n)\theta }]\end{array}\text{\hspace{1em}(5)}$$
这里$\text{Re}[\cdot ]$表示复数的实部；$({W\text{W}W}_k{x\text{x}x}_n{)}^{\ast }$表示$({W\text{W}W}_k{x\text{x}x}_n)$的共轭复数；$\theta \in \mathbb{R}$表示一个非零常数。

可以进一步将 f { q , k } f_{\{q,k\}} f{q,k}​写成矩阵乘法形式：
f { q , k } ( x m , m ) = ( cos ⁡ m θ − sin ⁡ m θ sin ⁡ m θ cos ⁡ m θ ) ( W { q , k } ( 11 ) W { q , k } ( 12 ) W { q , k } ( 21 ) W { q , k } ( 22 ) ) ( x m ( 1 ) x m ( 2 ) ) (6) f_{\{q,k\}} (\pmb x_m,m) =\begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix}\begin{pmatrix} W_{\{q,k\}}^{(11)} & W_{\{q,k\}}^{(12)} \\ W_{\{q,k\}}^{(21)} & W_{\{q,k\}}^{(22)} \end{pmatrix} \begin{pmatrix} x_m^{(1)} \\ x_m^{(2)} \end{pmatrix} \tag{6} f{q,k}​(xxxm​,m)=(cosmθsinmθ​−sinmθcosmθ​)(W{q,k}(11)​W{q,k}(21)​​W{q,k}(12)​W{q,k}(22)​​)(xm(1)​xm(2)​​)(6)
这里的 { q , k } \{q,k\} {q,k}表示$q$和$k$的集合，比如上式对$f_q$和$f_k$​都成立；包含$\sin m\theta$或$\cos m\theta$的矩阵是上面介绍的旋转矩阵。

其中$ (x^{(1)}_m, x^{(2)}_m) $为$x_m$ 在二维坐标中的表示。类似地，$g$ 可以被视为一个矩阵，从而能够在二维情况下求解等式$(4)$。具体来说，结合相对位置嵌入是很直接的：只需将仿射变换后的词嵌入向量旋转一定角度乘位置索引(旋转$m\theta$​)，从而解释了旋转位置嵌入背后的直觉。

我们进行直观理解，假设两个向量$q\text{q}q$和$k\text{k}k$它们的夹角为$\theta$，根据向量夹角的余弦我们知道$q\text{q}q\cdot k\text{k}k=|q\text{q}q||k\text{k}k|\cos \theta$​。

当$q\text{q}q$(逆时针)旋转$\alpha$角度后，与$k\text{k}k$的夹角变成了$\theta +\alpha$：

当$k\text{k}k$旋转$\beta$角度后，与$q\text{q}q$的夹角变成了$\theta -\beta$：

当两个向量同时旋转后，它们的夹角变成了$\theta +\alpha -\beta$。内积表达式为：
$$q\text{q}q\cdot k\text{k}k=|q\text{q}q||k\text{k}k|\cos (\theta +\alpha -\beta )$$
特殊地，当$\alpha -\beta =0$​​时，即两个向量旋转的角度相同，它们的内积不变。通过这两个向量的夹角来影响内积的值。通过这种直觉，公式(4)是成立的。

为了将我们在二维空间中的结果推广到任意 ${x\text{x}x}_i\in {\mathbb{R}}^d$，其中 $d$ 是偶数。我们可以将 $d$ 维空间划分为 $d/2 $个子空间(分块矩阵)，并结合内积的线性特性进行组合，将 f { q , k } f_{\{q,k\}} f{q,k}​​ 转化为：
f { q , k } = ( x m , m ) = R Θ , m d W { q , k } x m (7) f_{\{q,k\}} = (\pmb x_m,m) = \pmb R_{\Theta,m}^d \pmb W_{\{q,k\}} \pmb x_m \tag{7} f{q,k}​=(xxxm​,m)=RRRΘ,md​WWW{q,k}​xxxm​(7)

这里说的特性是指线性叠加性：

1. 定义：内积的定义是两个向量对应分量相乘后再相加。假设有两个向量 $\vec{v}=(v_1,v_2,...,v_n)$ 和 $\vec{w}=(w_1,w_2,...,w_n)$，它们的内积可以表示为 $\vec{v}\cdot \vec{w}=v_1{w}_1+v_2{w}_2+...+v_n{w}_n$。

2. 线性性质：内积满足线性叠加性，即对于任意标量 $a$ 和向量 $\vec{v},\vec{w},\vec{u}$，有以下性质：

   • 可加性：$\vec{v}\cdot (\vec{w}+\vec{u})=\vec{v}\cdot \vec{w}+\vec{v}\cdot \vec{u}$
   • 齐次性：$(a\vec{v})\cdot \vec{w}=a(\vec{v}\cdot \vec{w})$

其中
$${R\text{R}R}_{\Theta ,m}^d=\left(\begin{array}{ccccccc}\cos m{\theta }_1& -\sin m{\theta }_1& 0& 0& \cdots & 0& 0\\ \sin m{\theta }_1& \cos m{\theta }_1& 0& 0& \cdots & 0& 0\\ 0& 0& \cos m{\theta }_2& -\sin m{\theta }_2& \cdots & 0& 0\\ 0& 0& \sin m{\theta }_2& \cos m{\theta }_2& \cdots & 0& 0\\ ⋮& ⋮& ⋮& ⋮& \ddots & ⋮& ⋮\\ 0& 0& 0& 0& \cdots & \cos m{\theta }_{d/2}& -\sin m{\theta }_{d/2}\\ 0& 0& 0& 0& \cdots & \sin m{\theta }_{d/2}& \cos m{\theta }_{d/2}\end{array}\right)\text{\hspace{1em}(8)}$$
是一个带有预定义参数 Θ = { θ i = 1000 0 − 2 ( i − 1 ) / d , i ∈ [ 1 , 2 , . . . , d / 2 ] } Θ = \{θ_i = 10000^{−2(i−1)/d}, i ∈ [1, 2, ..., d/2]\} Θ={θi​=10000−2(i−1)/d,i∈[1,2,...,d/2]}​ 的旋转矩阵。RoPE的图示如原论文中的图(1)所示。将RoPE应用于等式(3)中的自注意力机制，我们可以得到：
$${q\text{q}q}_m^{\top }{k\text{k}k}_n=({R\text{R}R}_{\Theta ,m}^d{W\text{W}W}_q{x\text{x}x}_m{)}^{\top }({R\text{R}R}_{\Theta ,n}^d{W\text{W}W}_k{x\text{x}x}_n)={x\text{x}x}_m^{\top }{W\text{W}W}_q{R\text{R}R}_{\Theta ,n-m}^d{W\text{W}W}_k{x\text{x}x}_n\text{\hspace{1em}(9)}$$
其中${R\text{R}R}_{\Theta ,n-m}^d=({R\text{R}R}_{\Theta ,m}^d{)}^{\top }{R\text{R}R}_{\Theta ,n}^d$。值得指出的是，${R\text{R}R}_{\Theta }$​是一个正交矩阵，它不会改变向量的模长，因此通常来说它不会改变原模型的稳定性。

我们可以增大$\theta$的base以支持更长的上下文，这里是10000。

上图所说的是一个长度为6的序列，在进行自注意力计算时，Query和Key向量经过旋转位置编码变换的过程。首先对于位置1来说，记为$m$。然后仅考虑第一个二维子空间，即$(x_1,x_2)$向量，旋转$m{\theta }_1$后得到的增强表示。

由于公式(8)中${R\text{R}R}_{\Theta ,m}^d$的稀疏性，可以通过下述等价方式来实现${R\text{R}R}_{\Theta ,m}^d$和$x\text{x}x\in {\mathbb{R}}^d$的乘法：
KaTeX parse error: No such environment: equation at position 37: …\pmb x = \begin{̲e̲q̲u̲a̲t̲i̲o̲n̲}̲\begin{pmatrix}…
其中$\otimes$​是逐位对应相乘。

为什么可以简化成这样子，把乘$x\text{x}x$带入公式(8)得到：
$${R\text{R}R}_{\Theta ,m}^{d}x\text{x}x=\left(\begin{array}{c}\begin{array}{ccccccc}\cos m{\theta }_1& -\sin m{\theta }_1& 0& 0& \cdots & 0& 0\\ \sin m{\theta }_1& \cos m{\theta }_1& 0& 0& \cdots & 0& 0\\ 0& 0& \cos m{\theta }_2& -\sin m{\theta }_2& \cdots & 0& 0\\ 0& 0& \sin m{\theta }_2& \cos m{\theta }_2& \cdots & 0& 0\\ ⋮& ⋮& ⋮& ⋮& \ddots & ⋮& ⋮\\ 0& 0& 0& 0& \cdots & \cos m{\theta }_{d/2}& -\sin m{\theta }_{d/2}\\ 0& 0& 0& 0& \cdots & \sin m{\theta }_{d/2}& \cos m{\theta }_{d/2}\end{array}\end{array}\right)\left(\begin{array}{c}{x}_1\\ x_2\\ x_3\\ x_4\\ ⋮\\ x_{d-1}\\ x_d\end{array}\right)$$
根据分块矩阵的乘法，我们仅考虑左右两边矩阵的第一块，其得到(10)中向量的第1和第2个元素：
$$\left(\begin{array}{cc}\cos m{\theta }_1& -\sin m{\theta }_1\\ \sin m{\theta }_1& \cos m{\theta }_1\end{array}\right)\left(\begin{array}{c}{x}_1\\ x_2\end{array}\right)=\left(\begin{array}{c}{x}_1\cos m{\theta }_1-x_2\sin m{\theta }_1\\ x_1\sin m{\theta }_1+x_2\cos m{\theta }_1\end{array}\right)$$
因此这是成立的。

代码实现

本节参考LLaMA源码来实现旋转位置编码，同时底层实现逻辑进行一个解释。

首先定义一个函数生成旋转矩阵：

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):"""给定维度预计算频率(\theta) Tensor的复指数(complex exponentials,cis)Args:dim (int): dimension of the frequency tensorend (int): end index for precomputing frequenciestheta (float, optional): scaling factor for frequency computation. Defaults to 10000.0.Returns:torch.Tensor: Precomputed frequency tensor with complex exponentials."""# freqs (dim/2, )# theta_i = 10000 ** (-2(i-1)/dim) for i = [1,2,...,dim / 2]# theta_i# we start from 0 dont need to do i-1freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))# generate token sequence m = [0, 1, ..., seq_len - 1]# m (end, )m = torch.arange(end, device=freqs.device)# compute m * \theta# freqs (end, dim / 2)freqs = torch.outer(m, freqs).float()# freqs_cis (end, dim / 2)freqs_cis = torch.polar(torch.ones_like(freqs), freqs)return freqs_cis

这个函数用于生成公式(8)中的旋转矩阵。

首先计算预定义参数 Θ = { θ i = 1000 0 − 2 ( i − 1 ) / d , i ∈ [ 1 , 2 , . . . , d / 2 ] } Θ = \{θ_i = 10000^{−2(i−1)/d}, i ∈ [1, 2, ..., d/2]\} Θ={θi​=10000−2(i−1)/d,i∈[1,2,...,d/2]} ，我们的$i$从$0$开始因此不需要$i-1$，对应上面的Line 17。

然后考虑所有的位置，生成一个m = (seq_len, )形状的向量，Line 20。

计算m和Line 17计算出来的freqs的外积，即m中的每个位置$m_i$都会乘上$\Theta$的每个元素，得到一个(seq_len, dim / 2)形状的矩阵。假设序列的长度

假设$m=[m_1,m_2,\cdots ,m_T]=[1,2,\cdots ,N]$​，这里$N$表示序列长度。

它们的乘积是一个矩阵：
$$\left(\begin{array}{cccc}{m}_1{\theta }_1& m_1{\theta }_2& \cdots & m_1{\theta }_{d/2}\\ m_2{\theta }_1& m_2{\theta }_2& \cdots & m_2{\theta }_{d/2}\\ ⋮& ⋮& \ddots & ⋮\\ m_N{\theta }_1& m_N{\theta }_2& \cdots & m_N{\theta }_{d/2}\end{array}\right)$$
最后在Line 25通过torch.polar将它们转换为复数形式：
$$\left(\begin{array}{cccc}\cos (m_1{\theta }_1)+i\cdot \sin (m_1{\theta }_1)& \cos (m_1{\theta }_2)+i\cdot \sin (m_1{\theta }_2)& \cdots & \cos (m_1{\theta }_{d/2})+i\cdot \sin (m_1{\theta }_{d/2})\\ \cos (m_2{\theta }_1)+i\cdot \sin (m_2{\theta }_1)& \cos (m_2{\theta }_2)+i\cdot \sin (m_2{\theta }_2)& \cdots & \cos (m_2{\theta }_{d/2})+i\cdot \sin (m_2{\theta }_{d/2})\\ ⋮& ⋮& \ddots & ⋮\\ \cos (m_N{\theta }_1)+i\cdot \sin (m_N{\theta }_1)& \cos (m_N{\theta }_2)+i\cdot \sin (m_N{\theta }_2)& \cdots & \cos (m_N{\theta }_{d/2})+i\cdot \sin (m_N{\theta }_{d/2})\end{array}\right)$$
torch.polar(abs, angle)基于abs和angle计算出一个极坐标系中的复数表示：

那如何达到公式(10)的结果呢，为了简单，这里只展示$d=4$​的情况，考虑某个Token $x\text{x}x$：
$$x\text{x}x=\left[\begin{array}{cccc}{x}_1& x_2& x_3& x_4\end{array}\right]$$
第一步把$x\text{x}x$的元素两两分组：
$$x\text{x}x=\left[\begin{array}{cc}[x_1,x_2]& [x_3,x_4]\end{array}\right]$$
也不考虑批次维度，形状由(1,4)变成(1,2,2)。然后把新的$x\text{x}x$转换成复数的形式，形状变成了(1, 2)：
$$x\text{x}x=\left[\begin{array}{cc}{x}_1+i\cdot x_2& x_3+i\cdot x_4\end{array}\right]$$
即每个二维向量变成了一个复数。然后我们把这个向量矩阵和freqs_cis对应的向量对应位置相乘(分别旋转$m{\theta }_1,m{\theta }_2$角度：$d/2=4/2=2$)，这里假设当前位置为$m$​，然后有：
$$\begin{array}{rl}x\text{x}x& =\left[\begin{array}{cc}{x}_1+i\cdot x_2& x_3+i\cdot x_4\end{array}\right]\otimes \left[\begin{array}{cc}\cos (m{\theta }_1)+i\cdot \sin (m{\theta }_1)& \cos (m{\theta }_2)+i\cdot \sin (m{\theta }_2)\end{array}\right]\\ & =\left[\begin{array}{cc}(x_1+i\cdot x_2)[\cos (m{\theta }_1)+i\cdot \sin (m{\theta }_1)]& (x_3+i\cdot x_4)[\cos (m{\theta }_2)+i\cdot \sin (m{\theta }_2)]\end{array}\right]\\ & =\left[\begin{array}{cc}{x}_1\cos m{\theta }_1+i\cdot x_1\sin m{\theta }_1+i\cdot x_2\cos m{\theta }_1-x_2\sin m{\theta }_1& x_3\cos m{\theta }_2+i\cdot x_3\sin m{\theta }_2+i\cdot x_4\cos m{\theta }_2-x_4\sin m{\theta }_2\end{array}\right]\\ & =\left[\begin{array}{cc}{x}_1\cos m{\theta }_1-x_2\sin m{\theta }_1+i(x_1\sin m{\theta }_1+x_2\cos m{\theta }_1)& x_3\cos m{\theta }_2-x_4\sin m{\theta }_2+i(x_3\sin m{\theta }_2+x_4\cos m{\theta }_2)\end{array}\right]\end{array}$$

得到一个形状为(1,2)的复数项链。

然后我们把里面的复数变为二维向量：
$$x\text{x}x=\left[\begin{array}{cc}\left[\begin{array}{c}{x}_1\cos m_1{\theta }_1-x_2\sin m_1{\theta }_1\\ x_1\sin m_1{\theta }_1+x_2\cos m_1{\theta }_1\end{array}\right]& \left[\begin{array}{c}{x}_3\cos m_1{\theta }_2-x_4\sin m_1{\theta }_2\\ x_3\sin m_1{\theta }_2+x_4\cos m_1{\theta }_2\end{array}\right]\end{array}\right]$$
最后拉平其中的二维向量：
$$x\text{x}x=\left[\begin{array}{cccc}{x}_1\cos m{\theta }_1-x_2\sin m{\theta }_1& x_1\sin m{\theta }_1+x_2\cos m{\theta }_1& x_3\cos m{\theta }_2-x_4\sin m{\theta }_2& x_3\sin m{\theta }_2+x_4\cos m_1{\theta }_2\end{array}\right]$$
比较公式(10)中前4行的结果，可以发现是一样的，只不过列向量变成了行向量。

基于上面的过程我们就不难理解下面的代码：

def apply_rotary_emb(xq: Tensor, xk: Tensor, freq_cis: Tensor):"""使用给定的频率Tensor将旋转嵌入应用到输入张量中。该函数使用提供的频率使用给定的频率Tensor将旋转嵌入应用到输入张量中。freqs_cis将旋转嵌入应用到给定的查询xq和键xk张量上。输入张量被重塑为复数，并且频率张量被重塑以匹配广播兼容性。生成的张量包含旋转嵌入，并作为实张量返回。Args:xq (torch.Tensor): Query tensor to apply rotary embeddings.xk (torch.Tensor): Key tensor to apply rotary embeddings.freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials.Returns:Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings."""# xq (batch_size, seq_len, n_head, head_dim)# xq_ (batch_size, seq_len, n_head, head_dim // 2, 2)xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)# turn to complex# xq_ (batch_size, seq_len, n_head, head_dim // 2)xq_ = torch.view_as_complex(xq_)xk_ = torch.view_as_complex(xk_)# 应用旋转操作，然后将结果转回实数# xq_out (batch_size, seq_len, n_head, head_dim)xq_out = torch.view_as_real(xq_ * freq_cis).flatten(2)xk_out = torch.view_as_real(xk_ * freq_cis).flatten(2)return xq_out.type_as(xq), xk_out.type_as(xk)

下篇文章我们会探讨如何应用旋转位置编码到自注意力上。

参考

1. [论文笔记]ROFORMER
2. 复数与二维空间旋转