[论文笔记]ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING(下)

发布于:2024-04-14 ⋅ 阅读:(204) ⋅ 点赞:(0)

引言

今天带来苏神的旋转位置编码论文,ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING,论文题目翻译过来是基于旋转位置嵌入增强的Transformer:RoFormer。旋转位置编码被很多大模型架构采用。

位置编码对Transformer架构是至关重要的。它为不同位置的序列元素之间的依赖建模提供了有价值的监督。本文作者提出了一种名为Rotary Position Embedding(RoPE,旋转位置编码)的新方法,可以有效地利用位置信息。

具体而言,RoPE通过旋转矩阵编码绝对位置,并在自注意力机制中同时引入显式的相对位置依赖。值得注意的是,RoPE具有一些有价值的特性,包括序列长度的灵活性、相对距离增加时的递减的元素间依赖性以及为线性自注意力提供相对位置编码的能力。

RoFormer已经集成到Huggingface中:https://huggingface.co/docs/transformers/model_doc/roformer。

笔记比较长,因此分为上下两部分。

证明(复数)

下面沿着作者的思路1来证明一下。

首先需要为查询 q m = W q x m \pmb q_m=\pmb W_q\pmb x_m qm=Wqxm和键 k n = W k x n \pmb k_n=\pmb W_k\pmb x_n kn=Wkxn向量添加绝对位置信息,假设通过下述运算来实现:
q ~ m = f q ( x m , m ) , k ~ n = f k ( x n , n ) (p1) \pmb {\tilde q_m} = f_q(\pmb x_m,m),\quad \pmb {\tilde k_n} = f_k(\pmb x_n,n) \tag {p1} q~m=fq(xm,m),k~n=fk(xn,n)(p1)
由于Attention的核心运算就是内积,我们希望内积的结果能带有相对位置信息,因此我们可以假设存在恒等关系:
⟨ f q ( x m , m ) , f k ( x n , n ) ⟩ = g ( x m , x n , m − n ) (p2) \langle f_q(\pmb x_m,m) , f_k(\pmb x_n,n) \rangle = g(\pmb x_m,\pmb x_n, m-n) \tag{p2} fq(xm,m),fk(xn,n)⟩=g(xm,xn,mn)(p2)
注意这里的 g g g是一个函数。把内积公式转化为一个函数 g g g,该函数只接受词嵌入 x m , x n \pmb x_m,\pmb x_n xm,xn以及它们的相对位置 m − n m-n mn作为输入变量。换句话说,如果能找到这样的 f f f使得上述等式成立,我们就认为这个内积的结果是带有相对位置信息的。

同时我们设定一些初始条件,比如 f q ( x m , 0 ) = q m f_q(\pmb x_m,0)=\pmb q_m fq(xm,0)=qm f k ( x n , 0 ) = k n f_k(\pmb x_n,0)=\pmb k_n fk(xn,0)=kn​,相当与0表示还是返回当前位置,没有增加任何绝对位置信息,不过进行了查询和健向量的转换。

我们先考虑二维情形,这里利用复数来求解,在复数中对于两个向量 ⟨ q , k ⟩ = Re [ q k ∗ ] \langle \pmb q,\pmb k \rangle = \text{Re}[\pmb q \pmb k^*] q,k=Re[qk],其中 Re [ ⋅ ] \text{Re}[\cdot] Re[]代表取复数的实部,这里 k ∗ \pmb k^* k表示 k \pmb k k的共轭复数。

先来证明一下,假设 q = a + b i , k = c + d i → k ∗ = c − d i \pmb q = a + bi,\pmb k=c + di \rightarrow \pmb k^* = c - di q=a+bi,k=c+dik=cdi
q k ∗ = ( a + b i ) ( c − d i ) = a c − a d i + b c i − b d i 2 = a c + b d − ( b c − a d ) i \pmb q \pmb k^* = (a+bi)(c-di) = ac -adi + bci -bdi^2 = ac +bd - (bc -ad)i qk=(a+bi)(cdi)=acadi+bcibdi2=ac+bd(bcad)i
取其实部得 a c + b d ac + bd ac+bd。而向量内积 q ⋅ k = a c + d b \pmb q \cdot \pmb k = ac +db qk=ac+db​。这就证明了"两个二维向量的内积,等于把它们当复数看时,一个复数与另一个复数的共轭的乘积实部。"

注意这里把二维向量当成复数看,即 q = ( a , b ) \pmb q =(a,b) q=(a,b)看成是 a + b i a+bi a+bi​,然后可以利用复数的性质就像运算,是一种常用的技巧。

所以我们有:
⟨ f q ( x m , m ) , f k ( x n , n ) ⟩ = Re [ f q ( x m , m ) f k ∗ ( x n , n ) ] = g ( x m , x n , m − n ) (p3) \langle f_q(\pmb x_m,m) , f_k(\pmb x_n,n) \rangle = \text{Re}[f_q(\pmb x_m,m)f_k^*(\pmb x_n,n)] = g(\pmb x_m,\pmb x_n, m-n) \tag{p3} fq(xm,m),fk(xn,n)⟩=Re[fq(xm,m)fk(xn,n)]=g(xm,xn,mn)(p3)
利用复数来求解,把向量内积转换为复数乘积。

假设存在复数 g ( x m , x n , m − n ) \pmb g(\pmb x_m,\pmb x_n, m-n) g(xm,xn,mn),使得
f q ( x m , m ) f k ∗ ( x n , n ) = g ( x m , x n , m − n ) (p4) f_q(\pmb x_m,m)f_k^*(\pmb x_n,n) = \pmb g(\pmb x_m,\pmb x_n, m-n) \tag{p4} fq(xm,m)fk(xn,n)=g(xm,xn,mn)(p4)
其中 g \pmb g g表示复数。

我们观察(p4)式两边,左边是复数乘积,得到一个复数;右边也是一个复数。我们知道复数有实部和虚部,且复数 z = a + b i z = a+bi z=a+bi可以写成 z = r e i θ z = re^{i\theta} z=reiθ r = a 2 + b 2 r=\sqrt{a^2+b^2} r=a2+b2 是该复数的模,幅角为 θ \theta θ

我们分别通过这种形式写出(p4)式中的三个复数:
f q ( x m , m ) = R f ( x m , m ) e i Θ f ( x m , m ) f k ∗ ( x n , n ) = R f ( x n , n ) e − i Θ f ( x n , n ) g ( x m , x n , m − n ) = R g ( x m , x n , m − n ) e i Θ g ( x m , x n , m − n ) (p5) \begin{aligned} f_q(\pmb x_m,m) &= R_f(\pmb x_m,m) e^{i\Theta_f(\pmb x_m,m)} \\ f_k^*(\pmb x_n,n) &= R_f(\pmb x_n,n) e^{-i\Theta_f(\pmb x_n,n)} \\ \pmb g(\pmb x_m,\pmb x_n, m-n) &= R_g(\pmb x_m,\pmb x_n, m-n)e^{i\Theta_g(\pmb x_m,\pmb x_n, m-n)} \end{aligned} \tag{p5} fq(xm,m)fk(xn,n)g(xm,xn,mn)=Rf(xm,m)eiΘf(xm,m)=Rf(xn,n)eiΘf(xn,n)=Rg(xm,xn,mn)eiΘg(xm,xn,mn)(p5)
其中 R f , R g R_f,R_g Rf,Rg是复数的模, Θ f , Θ g \Theta_f,\Theta_g Θf,Θg​是复数的幅角。

将(p5)代入(p4)可得一个方程组:
R f ( x m , m ) R f ( x n , n ) = R g ( x m , x n , m − n ) Θ f ( x m , m ) − Θ f ( x n , n ) = Θ g ( x m , x n , m − n ) (p6) \begin{aligned} R_f(\pmb x_m,m)R_f(\pmb x_n,n) &= R_g(\pmb x_m,\pmb x_n, m-n) \\ \Theta_f(\pmb x_m,m) - \Theta_f(\pmb x_n,n) &= \Theta_g(\pmb x_m,\pmb x_n, m-n) \end{aligned} \tag{p6} Rf(xm,m)Rf(xn,n)Θf(xm,m)Θf(xn,n)=Rg(xm,xn,mn)=Θg(xm,xn,mn)(p6)
对于(p6)的第一个方程,代入 m = n m=n m=n得:
R f ( x m , m ) R f ( x n , m ) = R g ( x m , x n , 0 ) = R f ( x m , 0 ) R f ( x n , 0 ) = ∣ ∣ q m ∣ ∣ ∣ ∣ k n ∣ ∣ (p6.1) R_f(\pmb x_m,m)R_f(\pmb x_n,m) = R_g(\pmb x_m,\pmb x_n, 0) =R_f(\pmb x_m,0)R_f(\pmb x_n,0) = ||\pmb q_m|| ||\pmb k_n|| \tag{p6.1} Rf(xm,m)Rf(xn,m)=Rg(xm,xn,0)=Rf(xm,0)Rf(xn,0)=∣∣qm∣∣∣∣kn∣∣(p6.1)
分析一下,因为 m − n = 0 m-n=0 mn=0,所以得到 R g ( x m , x n , 0 ) R_g(\pmb x_m,\pmb x_n, 0) Rg(xm,xn,0),然后可以再次利用第一个方程(从等式右边往左),得到 R f ( x m , 0 ) R f ( x n , 0 ) R_f(\pmb x_m,0)R_f(\pmb x_n,0) Rf(xm,0)Rf(xn,0)。再根据条件 f q ( x m , 0 ) = q m f_q(\pmb x_m,0)=\pmb q_m fq(xm,0)=qm f k ( x n , 0 ) = k n f_k(\pmb x_n,0)=\pmb k_n fk(xn,0)=kn,最后的结果就是这两个向量的模,它们不依赖位置 m m m

对于(p6)的第二个方程,代入 m = n m=n m=n得:
Θ f ( x m , m ) − Θ f ( x n , m ) = Θ g ( x m , x n , 0 ) = Θ f ( x m , 0 ) − Θ f ( x n , 0 ) = Θ ( q m ) − Θ ( k n ) (p7) \Theta_f(\pmb x_m,m) - \Theta_f(\pmb x_n,m) = \Theta_g(\pmb x_m,\pmb x_n, 0) = \Theta_f(\pmb x_m,0) - \Theta_f(\pmb x_n,0) = \Theta (\pmb q_m) - \Theta (\pmb k_n) \tag{p7} Θf(xm,m)Θf(xn,m)=Θg(xm,xn,0)=Θf(xm,0)Θf(xn,0)=Θ(qm)Θ(kn)(p7)
这里 Θ ( q m ) , Θ ( k n ) \Theta (\pmb q_m),\Theta (\pmb k_n) Θ(qm),Θ(kn)是分别是向量 q m \pmb q_m qm k n \pmb k_n kn本身的幅角。根据(p7)有:
Θ f ( x m , m ) − Θ f ( x n , m ) = Θ ( q m ) − Θ ( k n ) ⇒ Θ f ( x m , m ) − Θ ( q m ) = Θ f ( x n , m ) − Θ ( k n ) (p8) \Theta_f(\pmb x_m,m) - \Theta_f(\pmb x_n,m) = \Theta (\pmb q_m) - \Theta (\pmb k_n) \Rightarrow \Theta_f(\pmb x_m,m) - \Theta (\pmb q_m) = \Theta_f(\pmb x_n,m) - \Theta (\pmb k_n) \tag{p8} Θf(xm,m)Θf(xn,m)=Θ(qm)Θ(kn)Θf(xm,m)Θ(qm)=Θf(xn,m)Θ(kn)(p8)
所以 Θ f ( x m , m ) − Θ ( q m ) \Theta_f(\pmb x_m,m) - \Theta (\pmb q_m) Θf(xm,m)Θ(qm)是一个只与 m m m有关,与 q m \pmb q_m qm无关的函数,记为 φ ( m ) \varphi(m) φ(m),即 Θ f ( x m , m ) = Θ ( q m ) + φ ( m ) \Theta_f(\pmb x_m,m)= \Theta (\pmb q_m) + \varphi(m) Θf(xm,m)=Θ(qm)+φ(m)

接着代入 n = m − 1 n=m-1 n=m1到(p6)的第二个等式,得到
Θ f ( x m , m ) − Θ f ( x n , m − 1 ) = Θ g ( x m , x n , 1 ) (p9) \Theta_f(\pmb x_m,m) - \Theta_f(\pmb x_n,m-1) = \Theta_g(\pmb x_m,\pmb x_n, 1) \tag{p9} Θf(xm,m)Θf(xn,m1)=Θg(xm,xn,1)(p9)
上式两边同时减去 Θ ( q m ) \Theta (\pmb q_m) Θ(qm)
Θ f ( x m , m ) − Θ ( q m ) − Θ f ( x n , m − 1 ) = Θ g ( x m , x n , 1 ) − Θ ( q m ) (p10) \Theta_f(\pmb x_m,m) - \Theta (\pmb q_m) - \Theta_f(\pmb x_n,m-1) = \Theta_g(\pmb x_m,\pmb x_n, 1) - \Theta (\pmb q_m) \tag{p10} Θf(xm,m)Θ(qm)Θf(xn,m1)=Θg(xm,xn,1)Θ(qm)(p10)
代入 Θ f ( x m , m ) − Θ ( q m ) = φ ( m ) \Theta_f(\pmb x_m,m) - \Theta (\pmb q_m)=\varphi(m) Θf(xm,m)Θ(qm)=φ(m)到上式得
φ ( m ) − Θ f ( x n , m − 1 ) = Θ g ( x m , x n , 1 ) − Θ ( q m ) (p11) \varphi(m)- \Theta_f(\pmb x_n,m-1) = \Theta_g(\pmb x_m,\pmb x_n, 1) - \Theta (\pmb q_m) \tag{p11} φ(m)Θf(xn,m1)=Θg(xm,xn,1)Θ(qm)(p11)
上式两边同时加上 Θ ( k n ) \Theta (\pmb k_n) Θ(kn)
φ ( m ) + Θ ( k n ) − Θ f ( x n , m − 1 ) = Θ g ( x m , x n , 1 ) + Θ ( k n ) − Θ ( q m ) (p12) \varphi(m)+ \Theta (\pmb k_n)- \Theta_f(\pmb x_n,m-1) = \Theta_g(\pmb x_m,\pmb x_n, 1) + \Theta (\pmb k_n)- \Theta (\pmb q_m) \tag{p12} φ(m)+Θ(kn)Θf(xn,m1)=Θg(xm,xn,1)+Θ(kn)Θ(qm)(p12)
根据(p8)两边同乘 − 1 -1 1整理得到 Θ ( k n ) − Θ f ( x n , m ) = Θ ( q m ) − Θ f ( x m , m ) \Theta (\pmb k_n) - \Theta_f(\pmb x_n,m)=\Theta (\pmb q_m) - \Theta_f(\pmb x_m,m) Θ(kn)Θf(xn,m)=Θ(qm)Θf(xm,m)代入上式得
φ ( m ) + Θ ( q m ) − Θ f ( x m , m − 1 ) = Θ g ( x m , x n , 1 ) + Θ ( k n ) − Θ ( q m ) (p13) \varphi(m)+ \Theta (\pmb q_m) - \Theta_f(\pmb x_m,m-1) = \Theta_g(\pmb x_m,\pmb x_n, 1) + \Theta (\pmb k_n)- \Theta (\pmb q_m) \tag{p13} φ(m)+Θ(qm)Θf(xm,m1)=Θg(xm,xn,1)+Θ(kn)Θ(qm)(p13)
又因为 Θ f ( x m , m − 1 ) = Θ ( q m ) + φ ( m − 1 ) ⇒ Θ ( q m ) − Θ f ( x m , m − 1 ) = − φ ( m − 1 ) \Theta_f(\pmb x_m,m-1)= \Theta (\pmb q_m) + \varphi(m-1) \Rightarrow \Theta (\pmb q_m) - \Theta_f(\pmb x_m,m-1) = -\varphi(m-1) Θf(xm,m1)=Θ(qm)+φ(m1)Θ(qm)Θf(xm,m1)=φ(m1)代入上式得
φ ( m ) − φ ( m − 1 ) = Θ g ( x m , x n , 1 ) + Θ ( k n ) − Θ ( q m ) (p14) \varphi(m)-\varphi(m-1) = \Theta_g(\pmb x_m,\pmb x_n, 1) + \Theta (\pmb k_n)- \Theta (\pmb q_m) \tag{p14} φ(m)φ(m1)=Θg(xm,xn,1)+Θ(kn)Θ(qm)(p14)
上式右边不依赖于 m m m,注意 x m , x n , q m \pmb x_m,\pmb x_n,\pmb q_m xm,xn,qm可以认为是确定的,不依赖于我们增加的绝对位置 m m m。因此左边也不依赖于 m m m,所以 { φ ( m ) } \{\varphi(m)\} {φ(m)}是等差数列。

设初始值 φ ( 0 ) = 0 \varphi(0) =0 φ(0)=0,且 φ ( 1 ) = θ \varphi(1) =\theta φ(1)=θ,我们有 φ ( m ) = m θ \varphi(m)=m\theta φ(m)=mθ。代入 Θ f ( x m , m ) − Θ ( q m ) = φ ( m ) \Theta_f(\pmb x_m,m) - \Theta (\pmb q_m)=\varphi(m) Θf(xm,m)Θ(qm)=φ(m)得到
Θ f ( x m , m ) = Θ ( q m ) + φ ( m ) = Θ ( q m ) + m θ (p15) \Theta_f(\pmb x_m,m) =\Theta (\pmb q_m) + \varphi(m) = \Theta (\pmb q_m) + m\theta \tag{p15} Θf(xm,m)=Θ(qm)+φ(m)=Θ(qm)+mθ(p15)
将其代入(p5)的第一个等式,再根据(p6.1)得
f q ( x m , m ) = R f ( x m , m ) e i Θ f ( x m , m ) = ∣ ∣ q m ∣ ∣ e i ( Θ ( q m ) + m θ ) = ∣ ∣ q m ∣ ∣ ( e i ( Θ ( q m ) ⋅ e i m θ ) = ( ∣ ∣ q m ∣ ∣ e i ( Θ ( q m ) ) e i m θ = q m e i m θ (p16) \begin{aligned} f_q(\pmb x_m,m) &= R_f(\pmb x_m,m) e^{i\Theta_f(\pmb x_m,m)} \\ &=||\pmb q_m|| e^{i(\Theta (\pmb q_m) + m\theta )} \\ &= ||\pmb q_m|| (e^{i(\Theta (\pmb q_m) }\cdot e^{im\theta }) \\ &= (||\pmb q_m|| e^{i(\Theta (\pmb q_m) } )e^{im\theta } \\ &= \pmb q_m e^{im\theta } \end{aligned} \tag{p16} fq(xm,m)=Rf(xm,m)eiΘf(xm,m)=∣∣qm∣∣ei(Θ(qm)+mθ)=∣∣qm∣∣(ei(Θ(qm)eimθ)=(∣∣qm∣∣ei(Θ(qm))eimθ=qmeimθ(p16)
这里利用了复数的知识,最后的 q m e i m θ \pmb q_m e^{im\theta } qmeimθ可以看成是复数 q m \pmb q_m qm乘上旋转子,旋转了幅角 m θ m\theta mθ

所以称为旋转位置编码,再根据旋转子 e i m θ = cos ⁡ m θ + i sin ⁡ m θ e^{im\theta} = \cos m\theta + i \sin m\theta eimθ=cosmθ+isinmθ与旋转矩阵的关系,上式还可以写成矩阵形式:
f q ( x m , m ) = q m e i m θ = [ cos ⁡ m θ − sin ⁡ m θ sin ⁡ m θ cos ⁡ m θ ] [ q m ( 1 ) q m ( 2 ) ] (p17) f_q(\pmb x_m,m) = \pmb q_m e^{im\theta }=\begin{bmatrix}\cos m\theta & -\sin m\theta\\ \sin m\theta & \cos m\theta\end{bmatrix} \begin{bmatrix}\pmb q_m^{(1)} \\ \pmb q_m^{(2)}\end{bmatrix}\tag{p17} fq(xm,m)=qmeimθ=[cosmθsinmθsinmθcosmθ][qm(1)qm(2)](p17)
至此我们验证了公式(13)。

通用形式

为了将我们在二维空间中的结果推广到任意 x i ∈ R d \pmb x_i ∈ \R^d xiRd,其中 d d d 是偶数。我们可以将 d d d 维空间划分为 d / 2 d/2 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 (14) f_{\{q,k\}} = (\pmb x_m,m) = \pmb R_{\Theta,m}^d \pmb W_{\{q,k\}} \pmb x_m \tag{14} f{q,k}=(xm,m)=RΘ,mdW{q,k}xm(14)

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

  1. 定义:内积的定义是两个向量对应分量相乘后再相加。假设有两个向量 v ⃗ = ( v 1 , v 2 , . . . , v n ) \vec{v} = (v_1, v_2, ..., v_n) v =(v1,v2,...,vn) w ⃗ = ( w 1 , w 2 , . . . , w n ) \vec{w} = (w_1, w_2, ..., w_n) w =(w1,w2,...,wn),它们的内积可以表示为 v ⃗ ⋅ w ⃗ = v 1 w 1 + v 2 w 2 + . . . + v n w n \vec{v} \cdot \vec{w} = v_1w_1 + v_2w_2 + ... + v_nw_n v w =v1w1+v2w2+...+vnwn

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

    • 可加性: v ⃗ ⋅ ( w ⃗ + u ⃗ ) = v ⃗ ⋅ w ⃗ + v ⃗ ⋅ u ⃗ \vec{v} \cdot (\vec{w} + \vec{u}) = \vec{v} \cdot \vec{w} + \vec{v} \cdot \vec{u} v (w +u )=v w +v u
    • 齐次性: ( a v ⃗ ) ⋅ w ⃗ = a ( v ⃗ ⋅ w ⃗ ) (a\vec{v}) \cdot \vec{w} = a(\vec{v} \cdot \vec{w}) (av )w =a(v w )

其中
R Θ , m d = ( cos ⁡ m θ 1 − sin ⁡ m θ 1 0 0 ⋯ 0 0 sin ⁡ m θ 1 cos ⁡ m θ 1 0 0 ⋯ 0 0 0 0 cos ⁡ m θ 2 − sin ⁡ m θ 2 ⋯ 0 0 0 0 sin ⁡ m θ 2 cos ⁡ m θ 2 ⋯ 0 0 ⋮ ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ 0 0 0 0 ⋯ cos ⁡ m θ d / 2 − sin ⁡ m θ d / 2 0 0 0 0 ⋯ sin ⁡ m θ d / 2 cos ⁡ m θ d / 2 ) (15) \pmb R_{\Theta,m}^d = \begin{pmatrix} \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 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 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{pmatrix} \tag{15} RΘ,md= cosmθ1sinmθ10000sinmθ1cosmθ1000000cosmθ2sinmθ20000sinmθ2cosmθ2000000cosmθd/2sinmθd/20000sinmθd/2cosmθd/2 (15)
是一个带有预定义参数 Θ = { θ i = 1000 0 − 2 ( i − 1 ) / d , i ∈ [ 1 , 2 , . . . , d / 2 ] } Θ = \{θ_i = 10000^{−2(i−1)/d}, i ∈ [1, 2, ..., d/2]\} Θ={θi=100002(i1)/d,i[1,2,...,d/2]}​ 的旋转矩阵。RoPE的图示如图(1)所示。将RoPE应用于方程(2)中的自注意力机制,我们可以得到:
q m ⊤ k n = ( R Θ , m d W q x m ) ⊤ ( R Θ , n d W k x n ) = x m ⊤ W q R Θ , n − m d W k x n (16) \pmb q_m^\top \pmb k_n = (\pmb R_{\Theta,m}^d \pmb W_{q}\pmb x_m)^\top (\pmb R_{\Theta,n}^d \pmb W_{k}\pmb x_n) = \pmb x_m^\top \pmb W_q \pmb R_{\Theta,n-m}^d \pmb W_k \pmb x_n \tag{16} qmkn=(RΘ,mdWqxm)(RΘ,ndWkxn)=xmWqRΘ,nmdWkxn(16)
其中 R Θ , n − m d = ( R Θ , m d ) ⊤ R Θ , n d \pmb R_{\Theta,n-m}^d=(\pmb R_{\Theta,m}^d)^\top \pmb R_{\Theta,n}^d RΘ,nmd=(RΘ,md)RΘ,nd。值得指出的是, R Θ \pmb R_{\Theta} RΘ​是一个正交矩阵,它不会改变向量的模长,因此通常来说它不会改变原模型的稳定性。

左乘一个正交矩阵不会改变向量的模长的原因在于正交矩阵的性质。一个矩阵 A A A是正交矩阵,当且仅当 A T A = I A^T A = I ATA=I,即其转置矩阵和自身的乘积等于单位矩阵。

考虑一个向量 v ⃗ \vec{v} v ,其模长为 ∥ v ⃗ ∥ \lVert \vec{v} \rVert v ,左乘一个正交矩阵 Q Q Q后得到新向量 u ⃗ = Q v ⃗ \vec{u} = Q\vec{v} u =Qv 。我们来证明 ∥ u ⃗ ∥ = ∥ v ⃗ ∥ \lVert \vec{u} \rVert = \lVert \vec{v} \rVert u =v

∥ u ⃗ ∥ 2 = u ⃗ T u ⃗ = ( v ⃗ T Q T ) ( Q v ⃗ ) = v ⃗ T ( Q T Q ) v ⃗ = v ⃗ T v ⃗ = ∥ v ⃗ ∥ 2 \begin{aligned} \lVert \vec{u} \rVert^2 &= \vec{u}^T \vec{u} \\ &= (\vec{v}^T Q^T) (Q \vec{v}) \\ &= \vec{v}^T (Q^T Q) \vec{v} \\ &= \vec{v}^T \vec{v} \\ &= \lVert \vec{v} \rVert^2 \end{aligned} u 2=u Tu =(v TQT)(Qv )=v T(QTQ)v =v Tv =v 2
因为正交矩阵 Q Q Q满足 Q T Q = I Q^T Q = I QTQ=I,所以 ∥ u ⃗ ∥ = ∥ v ⃗ ∥ \lVert \vec{u} \rVert = \lVert \vec{v} \rVert u =v

与之前的工作采用的位置嵌入方法,即(3)到(10),的加性特性相反,这里的方法是乘性的。此外,当应用于自注意力机制时,RoPE通过旋转矩阵乘积自然地融入了相对位置信息,而不是改变加性位置编码的扩展公式中的项。

image-20240413084948720

RoPE的性质

远程衰减:根据原始Transformer,作者设置 θ i = 1000 0 − 2 i / d θ_i = 10000^{−2i/d} θi=100002i/d。可以证明这种设置提供了远程衰减特性,这意味着当相对位置增加时,内积将衰减。这个特性符合这样一个直觉:相对距离较远的一对标记应该有较少的连接。

线性注意力: 注意力可以写成更通用的形式。
A t t e n t i o n ( Q , K , V ) m = ∑ n = 1 N sim ( q m , k n ) v n ∑ n = 1 N sim ( q m , k n ) (17) \begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_m = \frac{\sum_{n=1}^N \text{sim}(\boldsymbol{q}_m, \boldsymbol{k}_n)\boldsymbol{v}_n}{\sum_{n=1}^N \text{sim}(\boldsymbol{q}_m, \boldsymbol{k}_n)} \end{equation} \tag{17} Attention(Q,K,V)m=n=1Nsim(qm,kn)n=1Nsim(qm,kn)vn(17)
原始自注意力选择 sim ( q m , k n ) = exp ⁡ ( q m ⊤ k n / d ) \text{sim}(\pmb q_m,\pmb k_n)= \exp(\pmb q_m^\top \pmb k_n / \sqrt d) sim(qm,kn)=exp(qmkn/d ),注意原始自注意力会计算每对标记的query和key内积,这导致复杂度为 O ( N 2 ) O(N^2) O(N2)。根据2的工作,线性注意力重写公式(17)为:
A t t e n t i o n ( Q , K , V ) m = ∑ n = 1 N ϕ ( q m ) ⊤ φ ( k n ) v n ∑ n = 1 N ϕ ( q m ) ⊤ φ ( k n ) (18) \begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_m = \frac{\sum_{n=1}^N \phi(\boldsymbol{q}_m)^\top \varphi(\pmb k_n)\boldsymbol{v}_n}{\sum_{n=1}^N \phi(\boldsymbol{q}_m)^\top \varphi(\pmb k_n)} \end{equation} \tag{18} Attention(Q,K,V)m=n=1Nϕ(qm)φ(kn)n=1Nϕ(qm)φ(kn)vn(18)
其中 ϕ ( ⋅ ) \phi(\cdot) ϕ() φ ( ⋅ ) \varphi(\cdot) φ()通常是非负函数。Katharopoulos等人2首先提出了 ϕ ( x ) = φ ( x ) = elu ( x ) + 1 ϕ(x) = φ(x) = \text{elu}(x)+ 1 ϕ(x)=φ(x)=elu(x)+1,并使用矩阵乘法的结合性质计算了键值之间的乘积。Shen等人3使用softmax函数分别对查询和键进行归一化以进行内积运算,其等效于 ϕ ( q i ) = softmax ( q i ) ϕ(\pmb q_i) = \text{softmax}(\pmb q_i) ϕ(qi)=softmax(qi) ϕ ( k j ) = exp ⁡ ( k j ) ϕ(\pmb k_j) = \exp(\pmb k_j) ϕ(kj)=exp(kj)

我们重点讨论将RoPE与方程(18)结合起来。由于RoPE通过旋转注入位置信息,并保持隐藏表示的范数不变,我们可以通过将旋转矩阵与非负函数的输出相乘来将RoPE与线性注意力结合起来。
A t t e n t i o n ( Q , K , V ) m = ∑ n = 1 N ( R Θ , m d ϕ ( q m ) ) ⊤ ( R Θ , n d φ ( k n ) ) v n ∑ n = 1 N ϕ ( q m ) ⊤ φ ( k n ) (19) \begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_m = \frac{\sum_{n=1}^N (\pmb R_{\Theta,m}^d \phi(\boldsymbol{q}_m))^\top (\pmb R_{\Theta,n}^d \varphi(\pmb k_n))\boldsymbol{v}_n}{\sum_{n=1}^N \phi(\boldsymbol{q}_m)^\top \varphi(\pmb k_n)} \end{equation} \tag{19} Attention(Q,K,V)m=n=1Nϕ(qm)φ(kn)n=1N(RΘ,mdϕ(qm))(RΘ,ndφ(kn))vn(19)
值得注意的是,我们保持分母不变是为了避免除以零的风险,而分子中的求和可能包含负项。虽然方程(19)中每个值 v i \pmb v_i vi的权重并不严格满足概率归一化,但我们认为这种计算仍然可以对值的重要性进行建模。

理论解释

刚开始没看到作者在论文中还有推导,上面给出的证明过程也保留吧,对比着看看。

二维下RoPE的推导

d = 2 d = 2 d=2的情况下,我们考虑两个词嵌入向量 x q \pmb x_q xq x k \pmb x_k xk,分别对应查询和关键字,它们的位置分别为m和n。根据方程(1),它们的位置编码为:
q m = f q ( x q , m ) k n = f k ( x k , n ) (20) \begin{aligned} \pmb q_m &= f_q(\pmb x_q,m) \\ \pmb k_n &= f_k(\pmb x_k,n) \end{aligned} \tag{20} qmkn=fq(xq,m)=fk(xk,n)(20)
其中, q m \pmb q_m qm k n \pmb k_n kn的下标表示编码的位置信息。假设存在一个函数g,它定义了 f { q , k } f_{\{q,k\}} f{q,k}产生的向量之间的内积。
q m ⊤ k n = ⟨ f q ( x m , m ) , f k ( x n , n ) ⟩ = g ( x m , x n , n − m ) (21) \pmb q_m^\top \pmb k_n = \langle f_q(\pmb x_m,m), f_k(\pmb x_n,n)\rangle = g(\pmb x_m,\pmb x_n,n-m) \tag{21} qmkn=fq(xm,m),fk(xn,n)⟩=g(xm,xn,nm)(21)
进一步假设满足以下的初始条件:
q = f q ( x q , 0 ) k = f k ( x k , 0 ) (22) \begin{aligned} \pmb q &= f_q(\pmb x_q,0) \\ \pmb k &= f_k(\pmb x_k,0) \end{aligned} \tag{22} qk=fq(xq,0)=fk(xk,0)(22)
这可以被理解为对没有位置信息编码的向量进行编码。在这些设定下,我们试图找到 f q f_q fq f k f_k fk的解。首先,我们利用2D向量及其复数对应的几何意义,将方程(20)和(21)中的函数分解为
f q ( x q , m ) = R q ( x q , m ) e i Θ q ( x q , m ) f k ( x k , n ) = R k ( x k , n ) e i Θ k ( x k , n ) g ( x q , x k , n − m ) = R g ( x q , x k , n − m ) e i Θ g ( x q , x k , n − m ) (23) \begin{aligned} f_q(\pmb x_q,m) &= R_q(\pmb x_q,m) e^{i\Theta_q(\pmb x_q,m)} \\ f_k(\pmb x_k,n) &= R_k(\pmb x_k,n) e^{i\Theta_k(\pmb x_k,n)} \\ g(\pmb x_q,\pmb x_k, n-m) &= R_g(\pmb x_q,\pmb x_k, n-m)e^{i\Theta_g(\pmb x_q,\pmb x_k, n-m)} \end{aligned} \tag{23} fq(xq,m)fk(xk,n)g(xq,xk,nm)=Rq(xq,m)eiΘq(xq,m)=Rk(xk,n)eiΘk(xk,n)=Rg(xq,xk,nm)eiΘg(xq,xk,nm)(23)
其中 R { q , k } , R g R_{\{q,k\}},R_g R{q,k},Rg Θ { q , k } , Θ g \Theta_{\{q,k\}},\Theta_g Θ{q,k},Θg f { q , k } , g f_{\{q,k\}},g f{q,k},g对应的模和幅角。将它们带入等式(21),得到:
R q ( x q , m ) R k ( x k , n ) = R g ( x q , x k , n − m ) Θ k ( x k , n ) − Θ q ( x q , m ) = Θ g ( x q , x k , n − m ) (24) \begin{aligned} R_q(\pmb x_q,m)R_k(\pmb x_k,n) &= R_g(\pmb x_q,\pmb x_k, n-m) \\ \Theta_k(\pmb x_k,n) - \Theta_q(\pmb x_q,m) &= \Theta_g(\pmb x_q,\pmb x_k, n-m) \end{aligned} \tag{24} Rq(xq,m)Rk(xk,n)Θk(xk,n)Θq(xq,m)=Rg(xq,xk,nm)=Θg(xq,xk,nm)(24)
以及初始条件:
q = ∣ ∣ q ∣ ∣ e i θ q = R q ( x q , 0 ) e i Θ q ( x q , 0 ) k = ∣ ∣ k ∣ ∣ e i θ k = R k ( x k , 0 ) e i Θ k ( x k , 0 ) (25) \pmb q = ||\pmb q||e^{i\theta_q} = R_q(\pmb x_q, 0)e^{i\Theta_q(\pmb x_q,0)}\\ \pmb k = ||\pmb k||e^{i\theta_k} = R_k(\pmb x_k, 0)e^{i\Theta_k(\pmb x_k,0)} \tag{25} q=∣∣q∣∣eiθq=Rq(xq,0)eiΘq(xq,0)k=∣∣k∣∣eiθk=Rk(xk,0)eiΘk(xk,0)(25)
其中 ∣ ∣ q ∣ ∣ , ∣ ∣ k ∣ ∣ ||\pmb q||,||\pmb k|| ∣∣q∣∣,∣∣k∣∣ θ q , θ k \theta_q,\theta_k θq,θk分别是 q \pmb q q k \pmb k k的在二维平面上模长和幅角部分。

下面,设(24)中的 m = n m=n m=n,并考虑(25)中的初始条件:
R q ( x q , m ) R k ( x k , m ) = R g ( x q , x k , 0 ) = R q ( x q , 0 ) R k ( x k , 0 ) = ∣ ∣ q ∣ ∣ ∣ ∣ k ∣ ∣ (26a) R_q(\pmb x_q,m)R_k(\pmb x_k,m) = R_g(\pmb x_q,\pmb x_k, 0) =R_q(\pmb x_q,0)R_k(\pmb x_k,0) = ||\pmb q|| ||\pmb k|| \tag{26a} Rq(xq,m)Rk(xk,m)=Rg(xq,xk,0)=Rq(xq,0)Rk(xk,0)=∣∣q∣∣∣∣k∣∣(26a)

Θ k ( x k , m ) − Θ q ( x q , m ) = Θ g ( x q , x k , 0 ) = Θ k ( x k , 0 ) − Θ q ( x q , 0 ) = θ k − θ q (26b) \Theta_k(\pmb x_k,m) - \Theta_q(\pmb x_q,m) = \Theta_g(\pmb x_q,\pmb x_k, 0) = \Theta_k(\pmb x_k,0) - \Theta_q(\pmb x_q,0) = \theta_k -\theta_q \tag{26b} Θk(xk,m)Θq(xq,m)=Θg(xq,xk,0)=Θk(xk,0)Θq(xq,0)=θkθq(26b)

一方面,可以从(26a)式中得到 R f R_f Rf的直接解:
R q ( x q , m ) = R q ( x q , 0 ) = ∣ ∣ q ∣ ∣ R k ( x k , n ) = R k ( x k , 0 ) = ∣ ∣ k ∣ ∣ R g ( x q , x k , n − m ) = R g ( x q , x k , 0 ) = ∣ ∣ q ∣ ∣ ∣ ∣ k ∣ ∣ (27) \begin{aligned} R_q(\pmb x_q,m) &= R_q(\pmb x_q,0) = ||\pmb q|| \\ R_k(\pmb x_k,n) &= R_k(\pmb x_k,0) = ||\pmb k|| \\ R_g(\pmb x_q,\pmb x_k,n-m) &= R_g(\pmb x_q,\pmb x_k,0) = ||\pmb q||||\pmb k|| \\ \end{aligned} \tag{27} Rq(xq,m)Rk(xk,n)Rg(xq,xk,nm)=Rq(xq,0)=∣∣q∣∣=Rk(xk,0)=∣∣k∣∣=Rg(xq,xk,0)=∣∣q∣∣∣∣k∣∣(27)
这表明函数 R q R_q Rq R k R_k Rk R g R_g Rg与位置信息无关。另一方面,在(26b)中可以注意到, Θ q ( x q , m ) − θ q = Θ k ( x k , m ) − θ k Θ_q(\pmb x_q, m) − θ_q = Θ_k(\pmb x_k, m) − θ_k Θq(xq,m)θq=Θk(xk,m)θk 表明函数 Θ \Theta Θ不依赖于查询和键,我们将它们设置为 Θ f : = Θ q = Θ k Θ_f := Θ_q = Θ_k Θf:=Θq=Θk,而 Θ f ( x { q , k } , m ) − θ { q , k } Θ_f (\pmb x_{\{q,k\}}, m) − θ_{\{q,k\}} Θf(x{q,k},m)θ{q,k}是一个关于位置m的函数,与词嵌入 x { q , k } x_{\{q,k\}} x{q,k}无关,我们将其表示为$ ϕ(m)$,得到:
Θ f ( x { q , k } , m ) = ϕ ( m ) + θ { q , k } (28) \Theta_f(\pmb x_{\{q,k\}},m) = \phi(m) + \theta_{\{q,k\}} \tag{28} Θf(x{q,k},m)=ϕ(m)+θ{q,k}(28)
进一步,带入 n = m + 1 n=m+1 n=m+1到(24)并考虑上式,得到:
ϕ ( m + 1 ) − ϕ ( m ) = Θ g ( x g , x k , 1 ) + θ q − θ k (29) \phi(m+1) - \phi(m) = \Theta_g(\pmb x_g,\pmb x_k,1) + \theta_q -\theta_k \tag{29} ϕ(m+1)ϕ(m)=Θg(xg,xk,1)+θqθk(29)
由于RHS是一个与m无关的常数, ϕ ( m ) ϕ(m) ϕ(m)​在连续整数输入下产生一个等差数列:
ϕ ( m ) = m θ + γ (30) \phi(m) = m \theta+ \gamma \tag{30} ϕ(m)=mθ+γ(30)
其中 θ , γ ∈ R \theta,\gamma \in \R θ,γR实常数; θ \theta θ非零。总结(27)到(30)得:
f q ( x q , m ) = ∣ ∣ q ∣ ∣ e i θ q + m θ + γ = q e i ( m θ + γ ) f k ( x k , m ) = ∣ ∣ k ∣ ∣ e i θ k + n θ + γ = k e i ( n θ + γ ) (31) f_q(\pmb x_q,m) = ||\pmb q||e^{i\theta_q + m\theta + \gamma} = \pmb q e^{i(m\theta + \gamma)} \\ f_k(\pmb x_k,m) = ||\pmb k||e^{i\theta_k + n\theta + \gamma} = \pmb k e^{i(n\theta + \gamma)} \tag{31} fq(xq,m)=∣∣q∣∣eiθq+mθ+γ=qei(mθ+γ)fk(xk,m)=∣∣k∣∣eiθk+nθ+γ=kei(nθ+γ)(31)
需要注意的是,我们对方程(22)中的 f q f_q fq f k f_k fk没有施加任何约束,因此 f q ( x m , 0 ) f_q(\pmb x_m, 0) fq(xm,0) f k ( x n , 0 ) f_k(\pmb x_n, 0) fk(xn,0)可以自由选择。为了使我们的结果与方程(3)可比较,我们定义:
q = f q ( x m , 0 ) = W q x n k = f k ( x n , 0 ) = W k x n (32) \pmb q=f_q(\pmb x_m,0) = \pmb W_q \pmb x_n\\ \pmb k=f_k(\pmb x_n,0) = \pmb W_k \pmb x_n\\ \tag{32} q=fq(xm,0)=Wqxnk=fk(xn,0)=Wkxn(32)
然后,我们简单地设置(31)中的 γ = 0 \gamma =0 γ=0得到最终解:
f q ( x m , m ) = ( W q x m ) e i m θ , f k ( x n , n ) = ( W k x n ) e i n θ . (33) f_q(\pmb x_m, m) = (\pmb W_q \pmb x_m) e^{im\theta}, \\ f_k(\pmb x_n, n) = (\pmb W_k \pmb x_n) e^{in\theta}. \tag{33} fq(xm,m)=(Wqxm)eimθ,fk(xn,n)=(Wkxn)einθ.(33)

高效计算的旋转矩阵乘法实现

由于公式(15)中 R Θ , m d \pmb R^d_{\Theta,m} RΘ,md的稀疏性,所以直接用矩阵乘法来实现会很浪费算力,推荐通过下述方式来实现 R Θ , m d \pmb R^d_{\Theta,m} RΘ,md x ∈ R d \pmb x \in \R^d xRd的乘法:
R Θ , m d x = ( x 1 x 2 x 3 x 4 ⋮ x d − 1 x d ) ⊗ ( cos ⁡ m θ 1 cos ⁡ m θ 1 cos ⁡ m θ 2 cos ⁡ m θ 2 ⋮ cos ⁡ m θ d / 2 cos ⁡ m θ d / 2 ) + ( − x 2 x 1 − x 4 x 3 ⋮ − x d x d − 1 ) ⊗ ( sin ⁡ m θ 1 sin ⁡ m θ 1 sin ⁡ m θ 2 sin ⁡ m θ 2 ⋮ sin ⁡ m θ d / 2 sin ⁡ m θ d / 2 ) (34) \pmb R^d_{\Theta,m} \pmb x = \begin{equation}\begin{pmatrix}x_1 \\ x_2 \\ x_3 \\ x_4 \\ \vdots \\ x_{d-1} \\ x_{d} \end{pmatrix}\otimes\begin{pmatrix}\cos m\theta_1 \\ \cos m\theta_1 \\ \cos m\theta_2 \\ \cos m\theta_2 \\ \vdots \\ \cos m\theta_{d/2} \\ \cos m\theta_{d/2} \end{pmatrix} + \begin{pmatrix}-x_2 \\ x_1 \\ -x_4 \\ x_3 \\ \vdots \\ -x_{d} \\ x_{d-1} \end{pmatrix}\otimes\begin{pmatrix}\sin m\theta_1 \\ \sin m\theta_1 \\ \sin m\theta_2 \\ \sin m\theta_2 \\ \vdots \\ \sin m\theta_{d/2} \\ \sin m\theta_{d/2} \end{pmatrix}\end{equation} \tag{34} RΘ,mdx= x1x2x3x4xd1xd cosmθ1cosmθ1cosmθ2cosmθ2cosmθd/2cosmθd/2 + x2x1x4x3xdxd1 sinmθ1sinmθ1sinmθ2sinmθ2sinmθd/2sinmθd/2 (34)
其中 ⊗ \otimes ​是逐位对应相乘,从这个实现也可以看到,RoPE可以视为是乘性位置编码的变体。

为什么可以简化成这样子,把乘 x \pmb x x带入公式(15)得到:
R Θ , m d x = ( cos ⁡ m θ 1 − sin ⁡ m θ 1 0 0 ⋯ 0 0 sin ⁡ m θ 1 cos ⁡ m θ 1 0 0 ⋯ 0 0 0 0 cos ⁡ m θ 2 − sin ⁡ m θ 2 ⋯ 0 0 0 0 sin ⁡ m θ 2 cos ⁡ m θ 2 ⋯ 0 0 ⋮ ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ 0 0 0 0 ⋯ cos ⁡ m θ d / 2 − sin ⁡ m θ d / 2 0 0 0 0 ⋯ sin ⁡ m θ d / 2 cos ⁡ m θ d / 2 ) ( x 1 x 2 x 3 x 4 ⋮ x d − 1 x d ) \pmb R_{\Theta,m}^d \pmb x= \begin{pmatrix}\begin{array}{cc:cc:cc:cc} \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 \\ \hdashline 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 \\ \hdashline \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ \hdashline 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{pmatrix} \begin{pmatrix}x_1 \\ x_2 \\ \hdashline x_3 \\ x_4 \\ \hdashline\vdots \\ \hdashline x_{d-1} \\ x_{d}\end{pmatrix} RΘ,mdx= cosmθ1sinmθ10000sinmθ1cosmθ1000000cosmθ2sinmθ20000sinmθ2cosmθ2000000cosmθd/2sinmθd/20000sinmθd/2cosmθd/2 x1x2x3x4xd1xd
根据分块矩阵的乘法,我们仅考虑左右两边矩阵的第一块,其得到(34)中向量的第1和第2个元素:
( cos ⁡ m θ 1 − sin ⁡ m θ 1 sin ⁡ m θ 1 cos ⁡ m θ 1 ) ( x 1 x 2 ) = ( x 1 cos ⁡ m θ 1 − x 2 sin ⁡ m θ 1 x 1 sin ⁡ m θ 1 + x 2 cos ⁡ m θ 1 ) \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1\\ \sin m\theta_1 & \cos m\theta_1 \end{pmatrix} \begin{pmatrix} x_1\\ x_2 \end{pmatrix} = \begin{pmatrix}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{pmatrix} (cosmθ1sinmθ1sinmθ1cosmθ1)(x1x2)=(x1cosmθ1x2sinmθ1x1sinmθ1+x2cosmθ1)
因此这是成立的。

RoPE的远程衰减

可以看到,RoPE 形式上和 Sinusoidal 位置编码有点相似,只不过 Sinusoidal 位置编码是加性的,而 RoPE 可以视为乘性的。在 θ i \theta_i θi的选择上,同样沿用了 Sinusoidal 位置编码的方案,即 θ i = 1000 0 − 2 i / d \theta_i = 10000^{-2i/d} θi=100002i/d,它可以带来一定的远程衰减性。

具体证明如下: 将 q = W q x m \pmb q= \pmb W_q \pmb x_m q=Wqxm k = W k x n \pmb k = \pmb W_k\pmb x_n k=Wkxn两两分组,它们在公式(16)加上RoPE后的内积可以用复数乘法表示为:
( R Θ , m d W q x m ) ⊤ ( R Θ , n d W k x n ) = Re [ ∑ i = 0 d / 2 − 1 q [ 2 i : 2 i + 1 ] k [ 2 i : 2 i + 1 ] ∗ e i ( m − n ) θ i ] \begin{equation} (\boldsymbol{R}_{\Theta,m}^d \pmb W_q\pmb x_m)^{\top}(\boldsymbol{R}_{\Theta,n}^d \pmb W_k \pmb x_n) = \text{Re}\left[\sum_{i=0}^{d/2-1}\boldsymbol{q}_{[2i:2i+1]}\boldsymbol{k}_{[2i:2i+1]}^* e^{\text{i}(m-n)\theta_i}\right]\tag{35}\end{equation} (RΘ,mdWqxm)(RΘ,ndWkxn)=Re i=0d/21q[2i:2i+1]k[2i:2i+1]ei(mn)θi (35)
其中 q [ 2 i : 2 i + 1 ] \boldsymbol{q}_{[2i:2i+1]} q[2i:2i+1]表示 q \pmb q q中第 2 i 2i 2i到第 2 i + 1 2i+1 2i+1个元素组成一个二维向量。其实就是(12)中的第三个等式,不过分成了 d / 2 d/2 d/2个块。

h i = q [ 2 i : 2 i + 1 ] k [ 2 i : 2 i + 1 ] ∗ , S j = ∑ i = 0 j − 1 e i ( m − n ) θ i h_i = \boldsymbol{q}_{[2i:2i+1]}\boldsymbol{k}_{[2i:2i+1]}^*, S_j = \sum\limits_{i=0}^{j-1} e^{\text{i}(m-n)\theta_i} hi=q[2i:2i+1]k[2i:2i+1],Sj=i=0j1ei(mn)θi,并令 h d / 2 = 0 , S 0 = 0 h_{d/2}=0,S_0=0 hd/2=0,S0=0,由分部求和法
∑ i = 0 d / 2 − 1 q [ 2 i : 2 i + 1 ] k [ 2 i : 2 i + 1 ] ∗ e i ( m − n ) θ i = ∑ i = 0 d / 2 − 1 h i ( S i + 1 − S i ) = [ h d / 2 S d / 2 − h 0 S 0 ] − ∑ i = 0 d / 2 − 1 S i + 1 ( h i + 1 − h i ) = − ∑ i = 0 d / 2 − 1 S i + 1 ( h i + 1 − h i ) (36) \begin{aligned}\sum_{i=0}^{d/2-1}\boldsymbol{q}_{[2i:2i+1]}\boldsymbol{k}_{[2i:2i+1]}^* e^{\text{i}(m-n)\theta_i} &= \sum_{i=0}^{d/2-1} h_i (S_{i +1} - S_i)\\ &= [h_{d/2}S_{d/2} - h_0S_0]-\sum_{i=0}^{d/2-1} S_{i+1}(h_{i+1} - h_i)\\ &= -\sum_{i=0}^{d/2-1} S_{i+1}(h_{i+1} - h_i) \end{aligned} \tag{36} i=0d/21q[2i:2i+1]k[2i:2i+1]ei(mn)θi=i=0d/21hi(Si+1Si)=[hd/2Sd/2h0S0]i=0d/21Si+1(hi+1hi)=i=0d/21Si+1(hi+1hi)(36)
我们先来看(36)的第一个等式:
S i + 1 − S i = ∑ i = 0 i e i ( m − n ) θ i − ∑ i = 0 i − 1 e i ( m − n ) θ i = e i ( m − n ) θ i S_{i+1} -S_i = \sum_{i=0}^{i} e^{\text{i}(m-n)\theta_i} -\sum\limits_{i=0}^{i-1} e^{\text{i}(m-n)\theta_i}= e^{\text{i}(m-n)\theta_i} Si+1Si=i=0iei(mn)θii=0i1ei(mn)θi=ei(mn)θi

其实不用分部求和法也能证明剩下的等式:
∑ i = 0 d / 2 − 1 h i ( S i + 1 − S i ) = h 0 ( S 1 − S 0 ) + h 1 ( S 2 − S 1 ) + ⋯ + h d / 2 − 1 ( S d / 2 − S d / 2 − 1 ) = h 0 S 1 − h 0 S 0 + h 1 S 2 − h 1 S 1 + h 2 S 3 − h 2 S 2 + ⋯ + h d / 2 − 1 S d / 2 − h d / 2 − 1 S d / 2 − 1 = S 1 ( h 0 − h 1 ) + S 2 ( h 1 − h 2 ) + S d / 2 ( h d / 2 − 1 − h d / 2 ) = ∑ i = 0 d / 2 − 1 S i + 1 ( h i − h i + 1 ) = − ∑ i = 0 d / 2 − 1 S i + 1 ( h i + 1 − h i ) \begin{aligned} \sum_{i=0}^{d/2-1} h_i (S_{i +1} - S_i)&= h_0(S_{1} -S_0) +h_1(S_2-S_1) + \cdots + h_{d/2-1}(S_{d/2} - S_{d/2 -1})\\ &=h_0S_1 - h_0S_0 + h_1S_2 - h_1S_1+ h_2S_3 - h_2S_2+\cdots + h_{d/2-1} S_{d/2} - h_{d/2-1} S_{d/2 -1} \\ &= S_1(h_0 -h_1) + S_2(h_1 -h_2) + S_{d/2}(h_{d/2-1} -h_{d/2} ) \\ &= \sum_{i=0}^{d/2-1}S_{i+1}(h_i - h_{i+1}) \\ &=-\sum_{i=0}^{d/2-1}S_{i+1}(h_{i+1} -h_i) \end{aligned} i=0d/21hi(Si+1Si)=h0(S1S0)+h1(S2S1)++hd/21(Sd/2Sd/21)=h0S1h0S0+h1S2h1S1+h2S3h2S2++hd/21Sd/2hd/21Sd/21=S1(h0h1)+S2(h1h2)+Sd/2(hd/21hd/2)=i=0d/21Si+1(hihi+1)=i=0d/21Si+1(hi+1hi)

因此
∣ ∑ i = 0 d / 2 − 1 q [ 2 i : 2 i + 1 ] k [ 2 i : 2 i + 1 ] ∗ e i ( m − n ) θ i ∣ =   ∣ ∑ i = 0 d / 2 − 1 S i + 1 ( h i + 1 − h i ) ∣ ≤   ∑ i = 0 d / 2 − 1 ∣ S i + 1 ∣ ∣ h i + 1 − h i ∣ ≤   ( max ⁡ i ∣ h i + 1 − h i ∣ ) ∑ i = 0 d / 2 − 1 ∣ S i + 1 ∣ (37) \begin{equation}\begin{aligned} \left|\sum_{i=0}^{d/2-1}\boldsymbol{q}_{[2i:2i+1]}\boldsymbol{k}_{[2i:2i+1]}^* e^{\text{i}(m-n)\theta_i}\right| =&\, \left|\sum_{i=0}^{d/2-1} S_{i+1}(h_{i+1} - h_i)\right| \\ \leq&\, \sum_{i=0}^{d/2-1} |S_{i+1}| |h_{i+1} - h_i| \\ \leq&\, \left(\max_i |h_{i+1} - h_i|\right)\sum_{i=0}^{d/2-1} |S_{i+1}| \end{aligned}\end{equation} \tag{37} i=0d/21q[2i:2i+1]k[2i:2i+1]ei(mn)θi = i=0d/21Si+1(hi+1hi) i=0d/21Si+1∣∣hi+1hi(imaxhi+1hi)i=0d/21Si+1(37)
通过设定 θ i = 1000 0 − 2 i / d \theta_i = 10000^{-2i/d} θi=100002i/d,我们可以考察 1 d / 2 ∑ i = 1 d / 2 ∣ S i ∣ \frac{1}{d/2}\sum\limits_{i=1}^{d/2} |S_i| d/21i=1d/2Si随着相对距离 m − n m-n mn的增加而衰减,如图2所示。

image-20240413150840849

从图中我们可以可以看到随着相对距离的变大,内积结果有衰减趋势的出现。因此,选择 θ i = 1000 0 − 2 i / d \theta_i = 10000^{-2i/d} θi=100002i/d,确实能带来一定的远程衰减性。当然,能带来远程衰减性的不止这个选择,几乎任意的光滑单调函数都可以,这里只是沿用了已有的选择。

实验和评估

结论

在这项工作中,作者提出了一种新的位置嵌入方法,将绝对相对位置依赖纳入自注意力机制中,以增强Transformer架构的性能。

参考


  1. 苏剑林. (Mar. 23, 2021). 《Transformer升级之路:2、博采众长的旋转式位置编码 》 ↩︎

  2. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention ↩︎ ↩︎

  3. Efficient Attention: Attention with Linear Complexities ↩︎


网站公告

今日签到

点亮在社区的每一天
去签到