【论文精读】RoFormer

摘要

       本文提出旋转位置嵌入(RoPE),将位置嵌入信息利用到plm的学习过程中。RoPE用旋转矩阵实现绝对位置嵌入,同时在自注意力公式中融入了显式的相对位置嵌入。RoPE具有序列长度的灵活性、随着相对距离的增加而衰减的token间依赖性,以及为线性自注意力配备相对位置编码的能力。具体改进有:

  • RoPE将位置信息利用到plm的学习过程中。通过将上下文表示与具有明确理论解释的旋转矩阵相乘来编码相对位置
  • RoPE随着词嵌入相对距离的增加而产生相关性衰减
  • RoPE兼容线性自注意力

框架

自注意力计算过程

       给定一个长度为 N N N的输入序列:
S N = { w i } i = 1 N S_N=\{w_i \}^N_{i=1} SN={wi}i=1N

        w i w_i wi表示输入序列第 i i i个token, S N S_N SN对应的词嵌入向量表示为:
E N = { x i } i = 1 N E_N=\{x_i \}^N_{i=1} EN={xi}i=1N

        x i ∈ R d x_i\in \Reals^d xiRd表示第 i i i个token对应的不含位置嵌入的词嵌入向量。自注意力计算之前首先将位置嵌入信息加入词嵌入中并计算query、key、value向量,如下:
q m = f q ( x m , m ) k n = f k ( x n , n ) v n = f v ( x n , n ) q_m=f_q(x_m,m)\\k_n=f_k(x_n,n)\\v_n=f_v(x_n,n) qm=fq(xm,m)kn=fk(xn,n)vn=fv(xn,n)

       其中 q m q_m qm表示第 m m m个token对应的词嵌入向量 x m x_m xm集成位置嵌入 m m m后的query向量, k n k_n kn v n v _n vn为类似的key、value向量。自注意力为使用query和key计算注意力权重,同时计算输出为value表示的加权和,即为:
a m , n = e x p ( q m T K n d ) ∑ j = 1 N e x p ( q m T K n d ) o m = ∑ n = 1 N a m , n v n a_{m,n}=\frac {exp(\frac {q^T_mK_n} {\sqrt d})} {\textstyle\sum^N_{j=1}exp(\frac {q^T_mK_n} {\sqrt d})}\\o_m=\sum^N_{n=1}a_{m,n}v_n am,n=j=1Nexp(d qmTKn)exp(d qmTKn)om=n=1Nam,nvn

       本文目标为基于transformer的位置嵌入编码方法着重于构造一个合适的 f { q , k , v } f\{ q,k,v\} f{q,k,v}函数。

绝对位置嵌入

       常规的位置编码是在计算query,key和value向量之前,计算一个位置编码向量 p i ∈ R d p_i\in \Reals^d piRd加到词嵌入 x i x_i xi上,再乘以对应的变换矩阵 W { q , k , v } W\{q,k,v\} W{q,k,v},即:
f t : t ∈ { q , k , v } ( x i , i ) : = W t : t ∈ { q , k , v } ( x i + p i ) f_{t:t\in\{q,k,v\}}(x_i,i):=W_{t:t\in\{q,k,v\}}(x_i+p_i) ft:t{q,k,v}(xi,i):=Wt:t{q,k,v}(xi+pi)

        p i p_i pi取决于token x i x_i xi的位置,典型的计算方法为:
p i , 2 t = s i n ( i 1000 0 2 t d ) p i , 2 t + 1 = c o s ( i 1000 0 2 t d ) p_{i,2t}=sin(\frac i {10000^{\frac {2t} d}})\\p_{i,2t+1}=cos(\frac i {10000^{\frac {2t} d}}) pi,2t=sin(10000d2ti)pi,2t+1=cos(10000d2ti)

       其中 p i , 2 t p_{i,2t} pi,2t是位置编码向量 p i p_i pi的偶数维度的元素, p i , 2 t + 1 p_{i,2t+1} pi,2t+1 p i p_i pi的奇数维度的元素。

相对位置嵌入

       典型的相对位置编码为:
f q ( x m ) : = W q x m f k ( x n , n ) : = W k ( x n + p ˜ r k ) f v ( x n , n ) : = W v ( x n + p ˜ r v ) f_q(x_m):=W_qx_m\\f_k(x_n,n):=W_k(x_n+\~p^k_r)\\f_v(x_n,n):=W_v(x_n+\~p^v_r) fq(xm):=Wqxmfk(xn,n):=Wk(xn+p˜rk)fv(xn,n):=Wv(xn+p˜rv)

       其中 p ˜ r k , p ˜ r v ∈ R d \~p^k_r,\~p^v_r\in\Reals^d p˜rkp˜rvRd为可学习的相对位置嵌入参数, r = c l i p ( m − n , r m i n , r m a x ) r=clip(m-n,r_{min},r_{max}) r=clip(mn,rmin,rmax)表示位置 m m m n n n的相对距离。该方法假设,相对位置在超过一定距离后会失效,因此限制了相对距离的取值范围。类似方法诸多,略。

RoPE目标

       本文所提出的方法旨在通过绝对位置编码的方式实现相对位置编码,其在实践上的实用之处,比如可以拓展到线性Attention等。为达到这个目的,假设通过下述运算来给 q , k q,k q,k添加绝对位置信息:
q ˜ m = f ( q , m ) , k ˜ n = f ( k , n ) \~q_m=f(q,m),\~k_n=f(k,n) q˜m=f(q,m)k˜n=f(k,n)

       即分别为 q , k q,k q,k设计操作 f ( ⋅ , m ) , f ( ⋅ , n ) f(\cdot,m),f(\cdot,n) f(,m),f(,n),使得经过该操作后, q ˜ m , k ˜ n \~q_m,\~k_n q˜m,k˜n就带有了位置 m , n m,n m,n的绝对位置信息。而 q m T k n q^T_m k_n qmTkn可以在不同位置的token之间进行知识传递,为了纳入相对位置信息,假设 q m q_m qm k n k_n kn的内积由函数 g g g制定,该函数将单词嵌入 x m x_m xm, x n x_n xn和它们的相对位置 m − n m − n mn作为输入变量,其内积可编码相对位置信息。即:
< f q ( x m , m ) , f k ( x n , n ) > = g ( x m , x n , m − n ) <f_q(x_m,m),f_k(x_n,n)>=g(x_m,x_n,m-n) <fq(xm,m),fk(xn,n)>=g(xm,xn,mn)

       故目标为找到一种等价的编码机制 f q ( x m , m ) f_q(x_m, m) fq(xm,m) f k ( x n , n ) f_k(x_n, n) fk(xn,n),以符合上述关系。

RoPE构建

       考虑 d = 2 d=2 d=2的二维情况,给定query和key两个词嵌入向量 x q , x k x_q,x_k xq,xk及其对应位置 m , n m,n m,n。对应的位置编码为:
q m = f q ( x q , m ) k n = f k ( x k , n ) q_m=f_q(x_q,m)\\k_n=f_k(x_k,n) qm=fq(xq,m)kn=fk(xk,n)

       假设存在函数 g g g,定义了 f { q , k } f\{q,k\} f{q,k}之间的内积,则有:
q m T k n = < f q ( x q , m ) , f k ( x k , n ) > = g ( x q , x k , m − n ) q^T_mk_n=<f_q(x_q,m),f_k(x_k,n)>=g(x_q,x_k,m-n) qmTkn=<fq(xq,m),fk(xk,n)>=g(xq,xk,mn)

       有复数公式 < q , k > = R e [ q k ∗ ] <q,k>=Re[qk^*] <q,k>=Re[qk],则有 g ( x q , x k , m − n ) = R e [ f ( x q , m ) f ∗ ( x k , n ) ] g(x_q,x_k,m-n)=Re[f(x_q,m)f^*(x_k,n)] g(xq,xk,mn)=Re[f(xq,m)f(xk,n)],简单起见,令 g ( x q , x k , m − n ) = f ( x q , m ) f ∗ ( x k , n ) g(x_q,x_k,m-n)=f(x_q,m)f^*(x_k,n) g(xq,xk,mn)=f(xq,m)f(xk,n),则 q m , k n , q m T k n q_m,k_n,q^T_mk_n qm,kn,qmTkn的指数形式可分解为:
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 , m − n ) = R g ( x q , x k , m − n ) e i Θ g ( x q , x k , m − n ) f_q(x_q,m)=R_q(x_q,m)e^{i\Theta_q(x_q,m)}\\f_k(x_k,n)=R_k(x_k,n)e^{i\Theta_k(x_k,n)}\\g(x_q,x_k,m-n)=R_g(x_q,x_k,m-n)e^{i\Theta_g(x_q,x_k,m-n)} fq(xq,m)=Rq(xq,m)eiΘq(xq,m)fk(xk,n)=Rk(xk,n)eiΘk(xk,n)g(xq,xk,mn)=Rg(xq,xk,mn)eiΘg(xq,xk,mn)

       则有:
R q ( x q , m ) R k ( x k , n ) = R g ( x q , x k , m − n ) Θ q ( x q , m ) − Θ k ( x k , n ) = Θ g ( x q , x k , m − n ) R_q(x_q,m)R_k(x_k,n)=R_g(x_q,x_k,m-n)\\ \Theta_q(x_q,m)-\Theta_k(x_k,n)=\Theta_g(x_q,x_k,m-n) Rq(xq,m)Rk(xk,n)=Rg(xq,xk,mn)Θq(xq,m)Θk(xk,n)=Θg(xq,xk,mn)

       令空位置嵌入编码的向量为:
q = f q ( x q , 0 ) , k = f k ( x k , 0 ) q=f_q(x_q,0),k=f_k(x_k,0) q=fq(xq,0),k=fk(xk,0)

       其不依赖于 m m m n n n,故认为有:
R q ( x q , m ) = ∥ q ∥ , R k ( x k , n ) = ∥ k ∥ Θ q ( x q , 0 ) = θ q , Θ k ( x k , 0 ) = θ k R_q(x_q,m)=\lVert q \rVert,R_k(x_k,n)=\lVert k \rVert\\\Theta_q(x_q,0)=\theta_q,\Theta_k(x_ k,0)=\theta_k Rq(xq,m)=qRk(xk,n)=kΘq(xq,0)=θqΘk(xk,0)=θk

       则对应空位置嵌入有:
q = R q ( x q , 0 ) e i Θ q ( x q , 0 ) = ∥ q ∥ e i θ q k = R k ( x k , 0 ) e i Θ k ( x k , 0 ) = ∥ k ∥ e i θ k q=R_q(x_q,0)e^{i\Theta_q(x_q,0)}=\lVert q \rVert e^{i\theta_q}\\k=R_k(x_k,0)e^{i\Theta_k(x_k,0)}=\lVert k \rVert e^{i\theta_k} q=Rq(xq,0)eiΘq(xq,0)=qeiθqk=Rk(xk,0)eiΘk(xk,0)=keiθk

       若 n = m n=m n=m,则有:
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 ∥ Θ q ( x q , m ) − Θ k ( x k , m ) = Θ g ( x q , x k , 0 ) = ∥ Θ q ( x q , 0 ) − Θ k ( x k , 0 ) ∥ = ∥ θ q − θ k ∥ 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)=\lVert q \rVert \lVert k \rVert\\\Theta_q(x_q,m)-\Theta_k(x_k,m)=\Theta_g(x_q,x_k,0)=\lVert \Theta_q(x_q,0)-\Theta_k(x_k,0) \rVert=\lVert \theta_q-\theta_k \rVert Rq(xq,m)Rk(xk,m)=Rg(xq,xk,0)=Rq(xq,0)Rk(xk,0)=qkΘq(xq,m)Θk(xk,m)=Θg(xq,xk,0)=Θq(xq,0)Θk(xk,0)∥=θqθk

       进一步推导有:
Θ q ( x q , m ) − θ q = Θ k ( x k , m ) − θ k \Theta_q(x_q,m)-\theta_q=\Theta_k(x_k,m)-\theta_k Θq(xq,m)θq=Θk(xk,m)θk

       得知 Θ q ( x q , m ) − θ q \Theta_q(x_q,m)-\theta_q Θq(xq,m)θq Θ k ( x k , m ) − θ k \Theta_k(x_k,m)-\theta_k Θk(xk,m)θk只与 m m m相关,于 q q q k k k无关,则定义:
ϕ ( m ) = Θ f ( x { q , k } , m ) − θ { q , k } \phi(m)=\Theta_f(x_{\{q,k\}},m)-\theta_{\{q,k\}} ϕ(m)=Θf(x{q,k},m)θ{q,k}

       若 n = m + 1 n=m+1 n=m+1,则有:
ϕ ( m + 1 ) − ϕ ( m ) = ( Θ q ( x q , m + 1 ) − θ q ) − ( Θ k ( x k , m ) − θ k ) = Θ g ( x q , x k , 1 ) + θ k − θ q \phi(m+1)-\phi(m) \\ =(\Theta_q(x_q,m+1)-\theta_q)-(\Theta_k(x_k,m)-\theta_k)\\ =\Theta_g(x_q,x_k,1)+\theta_k-\theta_q ϕ(m+1)ϕ(m)=(Θq(xq,m+1)θq)(Θk(xk,m)θk)=Θg(xq,xk,1)+θkθq

       得知 ϕ ( m + 1 ) − ϕ ( m ) \phi(m+1)-\phi(m) ϕ(m+1)ϕ(m)是一个与 m m m无关的常数,故 ϕ ( m ) \phi(m) ϕ(m)使用连续整数输入可成为一个等差数列,则得到
ϕ ( m ) = m ( Θ g ( x q , x k , 1 ) + θ k − θ q ) + γ = m θ + γ \phi(m)=m(\Theta_g(x_q,x_k,1)+\theta_k-\theta_q)+\gamma\\ =m\theta+\gamma ϕ(m)=m(Θg(xq,xk,1)+θkθq)+γ=mθ+γ

       其中 θ , γ ∈ R \theta,\gamma\in\Reals θ,γR为常量, θ \theta θ非零。整合上述推导,可得到解决方案:
q m = f q ( x q , m ) = R q ( x q , m ) e i Θ q ( x q , m ) = ∥ q ∥ e i ( θ q + ϕ ( m ) ) = q e i ϕ ( m ) = q e i ( m θ + γ ) k n = f k ( x k , n ) = R k ( x k , n ) e i Θ k ( x k , n ) = ∥ k ∥ e i ( θ k + ϕ ( n ) ) = k e i ϕ ( n ) = k e i ( n θ + γ ) q_m =f_q(x_q,m) =R_q(x_q,m)e^{i\Theta_q(x_q,m)} =\lVert q \rVert e^{i(\theta_q+\phi(m))} =qe^{i\phi(m)} =qe^{i(m\theta+\gamma)} \\ k_n=f_k(x_k,n)=R_k(x_k,n)e^{i\Theta_k(x_k,n)}=\lVert k \rVert e^{i(\theta_k+\phi(n))}=ke^{i\phi(n)}=ke^{i(n\theta+\gamma)} qm=fq(xq,m)=Rq(xq,m)eiΘq(xq,m)=qei(θq+ϕ(m))=qeiϕ(m)=qei(mθ+γ)kn=fk(xk,n)=Rk(xk,n)eiΘk(xk,n)=kei(θk+ϕ(n))=keiϕ(n)=kei(nθ+γ)

       整合经典绝对位置嵌入函数,令:
q = W q x m , k = W k x n q=W_qx_m,k=W_kx_n q=Wqxmk=Wkxn

       令 γ = 0 \gamma=0 γ=0,则最终算法为:
q m = ( W q x m ) e i m θ k n = ( W k x n ) e i n θ q_m=(W_qx_m)e^{im\theta}\\k_n=(W_kx_n)e^{in\theta} qm=(Wqxm)eimθkn=(Wkxn)einθ

       因为:
e i ( m ) θ = c o s ( ( m ) θ ) + i s i n ( ( m ) θ ) e^{i(m)\theta}=cos((m)\theta)+isin((m)\theta) ei(m)θ=cos((m)θ)+isin((m)θ)

       则算法转换为矩阵形式有:
f { q , k } ( x m , m ) = ( c o s   m θ − s i n   m θ s i n   m θ c o s   m θ ) ( W { q , k } ( 11 ) W { q , k } ( 12 ) W { q , k } ( 21 ) W { q , k } ( 22 ) ) ( x m ( 1 ) x m ( 2 ) ) f_{\{q,k\}}(x_m,m)=\begin{pmatrix} cos \ m\theta & -sin \ m\theta \\ sin \ m\theta & cos \ m\theta \end{pmatrix} \begin{pmatrix} W^{(11)}_{\{q,k\}} & W^{(12)}_{\{q,k\}} \\ W^{(21)}_{\{q,k\}} & W^{(22)}_{\{q,k\}} \end{pmatrix} \begin{pmatrix} x^{(1)}_{m} \\ x^{(2)}_{m} \end{pmatrix} f{q,k}(xm,m)=(cos mθsin mθsin mθcos mθ)(W{q,k}(11)W{q,k}(21)W{q,k}(12)W{q,k}(22))(xm(1)xm(2))

       即算法的本质为将仿射变换的词嵌入向量旋转其位 置索引的角度倍数,从而导致词嵌入向量间的相对位置,故称之为旋转式位置编码。如下图所示:

image

证明

       根据上述推导,算法需满足:
< f q ( x q , m ) , f k ( x k , n ) > = g ( x q , x k , m − n ) <f_q(x_q,m),f_k(x_k,n)>=g(x_q,x_k,m-n) <fq(xq,m),fk(xk,n)>=g(xq,xk,mn)

       根据上述推导存在:
g ( x q , x k , m − n ) = R e [ ( W q x m ) ( W k x n ) ∗ e i ( m − n ) θ ] g(x_q,x_k,m-n)=Re[(W_qx_m)(W_kx_n)^*e^{i(m-n)\theta}] g(xq,xk,mn)=Re[(Wqxm)(Wkxn)ei(mn)θ]

       其中 R e [ ⋅ ] Re[\cdot] Re[]表示复数实部, ( W k x n ) ∗ (W_kx_n)^* (Wkxn) ( W k x n ) (W_kx_n) (Wkxn)的共轭复数。又因为有:
z = a + i b , z ∗ = a − i b z=a+ib,z^*=a-ib z=a+ibz=aib

       故有:
W q x m = q = q ( 1 ) + i q ( 2 ) ( W k x n ) ∗ = k ∗ = k ( 1 ) − i k ( 2 ) e i ( m − n ) θ = c o s ( ( m − n ) θ ) + i s i n ( ( m − n ) θ ) W_qx_m=q=q^{(1)}+iq^{(2)}\\(W_kx_n)^*=k^*=k^{(1)}-ik^{(2)}\\e^{i(m-n)\theta}=cos((m-n)\theta)+isin((m-n)\theta) Wqxm=q=q(1)+iq(2)(Wkxn)=k=k(1)ik(2)ei(mn)θ=cos((mn)θ)+isin((mn)θ)

       则:
g ( x q , x k , m − n ) = R e [ ( q ( 1 ) + i q ( 2 ) ) ( k ( 1 ) − i k ( 2 ) ) ( c o s ( ( m − n ) θ ) + i s i n ( ( m − n ) θ ) ) ] = R e [ ( ( q ( 1 ) k ( 1 ) + q ( 2 ) k ( 2 ) ) + i ( q ( 2 ) k ( 1 ) − q ( 1 ) k ( 2 ) ) ) ( c o s ( ( m − n ) θ ) + i s i n ( ( m − n ) θ ) ) ] = ( q ( 1 ) k ( 1 ) + q ( 2 ) k ( 2 ) ) c o s ( ( m − n ) θ ) − ( q ( 2 ) k ( 1 ) − q ( 1 ) k ( 2 ) ) s i n ( ( m − n ) θ ) ) g(x_q,x_k,m-n)=Re[(q^{(1)}+iq^{(2)})(k^{(1)}-ik^{(2)})(cos((m-n)\theta)+isin((m-n)\theta))] \\ =Re[((q^{(1)}k^{(1)}+q^{(2)}k^{(2)})+ i(q^{(2)}k^{(1)}-q^{(1)}k^{(2)}))(cos((m-n)\theta)+isin((m-n)\theta))] \\ =(q^{(1)}k^{(1)}+q^{(2)}k^{(2)})cos((m-n)\theta)-(q^{(2)}k^{(1)}-q^{(1)}k^{(2)})sin((m-n)\theta)) g(xq,xk,mn)=Re[(q(1)+iq(2))(k(1)ik(2))(cos((mn)θ)+isin((mn)θ))]=Re[((q(1)k(1)+q(2)k(2))+i(q(2)k(1)q(1)k(2)))(cos((mn)θ)+isin((mn)θ))]=(q(1)k(1)+q(2)k(2))cos((mn)θ)(q(2)k(1)q(1)k(2))sin((mn)θ))

       对于 < f q ( x q , m ) , f k ( x k , n ) > <f_q(x_q,m),f_k(x_k,n)> <fq(xq,m),fk(xk,n)>,则有:
f q ( x q , m ) = [ q ( 1 ) c o s ( m θ ) − q ( 2 ) s i n ( m θ ) , q ( 2 ) c o s ( m θ ) + q ( 1 ) s i n ( m θ ) ] f k ( x k , n ) = [ k ( 1 ) c o s ( n θ ) − k ( 2 ) s i n ( n θ ) , k ( 2 ) c o s ( n θ ) + k ( 1 ) s i n ( n θ ) ] < f q ( x q , m ) , f k ( x k , n ) > = ( q ( 1 ) c o s ( m θ ) − q ( 2 ) s i n ( m θ ) ) ( k ( 1 ) c o s ( n θ ) − k ( 2 ) s i n ( n θ ) ) + ( q ( 2 ) c o s ( m θ ) + q ( 1 ) s i n ( m θ ) ) ( k ( 2 ) c o s ( n θ ) + k ( 1 ) s i n ( n θ ) ) = q ( 1 ) c o s ( m θ ) k ( 1 ) c o s ( n θ ) − q ( 1 ) c o s ( m θ ) k ( 2 ) s i n ( n θ ) − q ( 2 ) s i n ( m θ ) k ( 1 ) c o s ( n θ ) + q ( 2 ) s i n ( m θ ) k ( 2 ) s i n ( n θ ) q ( 2 ) c o s ( m θ ) k ( 2 ) c o s ( n θ ) + q ( 2 ) c o s ( m θ ) k ( 1 ) s i n ( n θ ) q ( 1 ) s i n ( m θ ) k ( 2 ) c o s ( n θ ) + q ( 1 ) s i n ( m θ ) k ( 1 ) s i n ( n θ ) = q ( 1 ) k ( 1 ) ( c o s ( m θ ) c o s ( n θ ) + s i n ( m θ ) s i n ( n θ ) ) + q ( 1 ) k ( 2 ) ( − c o s ( m θ ) s i n ( n θ ) + s i n ( m θ ) c o s ( n θ ) ) + q ( 2 ) k ( 1 ) ( − s i n ( m θ ) c o s ( n θ ) + c o s ( m θ ) s i n ( n θ ) ) + q ( 2 ) k ( 2 ) ( s i n ( m θ ) s i n ( n θ ) + c o s ( m θ ) c o s ( n θ ) ) f_q(x_q,m)=[q^{(1)}cos(m\theta)-q^{(2)}sin(m\theta),q^{(2)}cos(m\theta)+q^{(1)}sin(m\theta)]\\f_k(x_k,n)=[k^{(1)}cos(n\theta)-k^{(2)}sin(n\theta),k^{(2)}cos(n\theta)+k^{(1)}sin(n\theta)]\\<f_q(x_q,m),f_k(x_k,n)>=(q^{(1)}cos(m\theta)-q^{(2)}sin(m\theta))(k^{(1)}cos(n\theta)-k^{(2)}sin(n\theta))\\ +(q^{(2)}cos(m\theta)+q^{(1)}sin(m\theta))(k^{(2)}cos(n\theta)+k^{(1)}sin(n\theta))\\ = q^{(1)}cos(m\theta)k^{(1)}cos(n\theta) - q^{(1)}cos(m\theta)k^{(2)}sin(n\theta)\\ - q^{(2)}sin(m\theta)k^{(1)}cos(n\theta) + q^{(2)}sin(m\theta)k^{(2)}sin(n\theta)\\ q^{(2)}cos(m\theta)k^{(2)}cos(n\theta) + q^{(2)}cos(m\theta)k^{(1)}sin(n\theta)\\ q^{(1)}sin(m\theta)k^{(2)}cos(n\theta)+q^{(1)}sin(m\theta)k^{(1)}sin(n\theta) \\ =q^{(1)}k^{(1)} (cos(m\theta)cos(n\theta) + sin(m\theta)sin(n\theta))\\ + q^{(1)}k^{(2)}(-cos(m\theta)sin(n\theta) + sin(m\theta)cos(n\theta))\\ +q^{(2)}k^{(1)}(-sin(m\theta)cos(n\theta) + cos(m\theta)sin(n\theta))\\ +q^{(2)}k^{(2)}(sin(m\theta)sin(n\theta)+cos(m\theta)cos(n\theta)) fq(xq,m)=[q(1)cos(mθ)q(2)sin(mθ),q(2)cos(mθ)+q(1)sin(mθ)]fk(xk,n)=[k(1)cos(nθ)k(2)sin(nθ),k(2)cos(nθ)+k(1)sin(nθ)]<fq(xq,m),fk(xk,n)>=(q(1)cos(mθ)q(2)sin(mθ))(k(1)cos(nθ)k(2)sin(nθ))+(q(2)cos(mθ)+q(1)sin(mθ))(k(2)cos(nθ)+k(1)sin(nθ))=q(1)cos(mθ)k(1)cos(nθ)q(1)cos(mθ)k(2)sin(nθ)q(2)sin(mθ)k(1)cos(nθ)+q(2)sin(mθ)k(2)sin(nθ)q(2)cos(mθ)k(2)cos(nθ)+q(2)cos(mθ)k(1)sin(nθ)q(1)sin(mθ)k(2)cos(nθ)+q(1)sin(mθ)k(1)sin(nθ)=q(1)k(1)(cos(mθ)cos(nθ)+sin(mθ)sin(nθ))+q(1)k(2)(cos(mθ)sin(nθ)+sin(mθ)cos(nθ))+q(2)k(1)(sin(mθ)cos(nθ)+cos(mθ)sin(nθ))+q(2)k(2)(sin(mθ)sin(nθ)+cos(mθ)cos(nθ))

       根据三角函数:
s i n ( a − b ) = s i n   a   c o s   b − c o s   a   s i n   b c o s ( a − b ) = c o s   a   c o s   b + s i n   a   s i n   b sin(a-b)=sin \ a \ cos \ b - cos \ a \ sin \ b \\cos(a-b)=cos \ a \ cos \ b + sin \ a \ sin \ b sin(ab)=sin a cos bcos a sin bcos(ab)=cos a cos b+sin a sin b

       则有:
< f q ( x q , m ) , f k ( x k , n ) > = q ( 1 ) k ( 1 ) c o s ( ( m − n ) θ ) + q ( 1 ) k ( 2 ) s i n ( ( m − n ) θ ) − q ( 2 ) k ( 1 ) s i n ( ( m − n ) θ ) + q ( 2 ) k ( 2 ) c o s ( ( m − n ) θ ) = ( q ( 1 ) k ( 1 ) + q ( 2 ) k ( 2 ) ) c o s ( ( m − n ) θ ) − ( q ( 2 ) k ( 1 ) − q ( 1 ) k ( 2 ) ) s i n ( ( m − n ) θ ) = g ( x q , x k , m − n ) <f_q(x_q,m),f_k(x_k,n)>\\=q^{(1)}k^{(1)}cos((m-n)\theta) +q^{(1)}k^{(2)}sin((m-n)\theta) \\ -q^{(2)}k^{(1)}sin((m-n)\theta) +q^{(2)}k^{(2)}cos((m-n)\theta)\\ =(q^{(1)}k^{(1)}+q^{(2)}k^{(2)})cos((m-n)\theta)-(q^{(2)}k^{(1)}-q^{(1)}k^{(2)})sin((m-n)\theta)\\ =g(x_q,x_k,m-n) <fq(xq,m),fk(xk,n)>=q(1)k(1)cos((mn)θ)+q(1)k(2)sin((mn)θ)q(2)k(1)sin((mn)θ)+q(2)k(2)cos((mn)θ)=(q(1)k(1)+q(2)k(2))cos((mn)θ)(q(2)k(1)q(1)k(2))sin((mn)θ)=g(xq,xk,mn)

       故算法 q m = ( W q x m ) e i m θ q_m=(W_qx_m)e^{im\theta} qm=(Wqxm)eimθ k n = ( W k x n ) e i n θ k_n=(W_kx_n)e^{in\theta} kn=(Wkxn)einθ满足 < f q ( x q , m ) , f k ( x k , n ) > = g ( x q , x k , m − n ) <f_q(x_q,m),f_k(x_k,n)>=g(x_q,x_k,m-n) <fq(xq,m),fk(xk,n)>=g(xq,xk,mn)限制的条件。

一般形式

       将算法推广到任意 x i ∈ R d x_i\in\Reals^d xiRd d d d为偶数,将 d d d维空间划分维 d / 2 d/2 d/2子空间,由于内积满足线性叠加性,因此任意偶数维的RoPE,我们都可以表示为二维情形的拼接,即可转化为:
f { q , k } ( x m , m ) = R Θ , m d W { q , k } x m f_{\{q,k\}}(x_m,m)=R^d_{\Theta,m}W_{\{q,k\}}x_m f{q,k}(xm,m)=RΘ,mdW{q,k}xm
R Θ , m d = ( c o s   m θ 1 − s i n   m θ 1 0 0 … 0 0 s i n   m θ 1 c o s   m θ 1 0 0 … 0 0 0 0 c o s   m θ 2 − s i n   m θ 2 … 0 0 0 0 s i n   m θ 2 c o s   m θ 2 … 0 0 ⋮ ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ 0 0 0 0 … c o s   m θ d / 2 − s i n   m θ d / 2 0 0 0 0 … s i n   m θ d / 2 c o s   m θ d / 2 ) R^d_{\Theta,m}=\begin{pmatrix} cos \ m\theta_1 & -sin \ m\theta_1 & 0 & 0 & \dots &0 &0\\ sin \ m\theta_1 & cos \ m\theta_1 & 0 & 0 & \dots &0 &0 \\ 0 & 0 & cos \ m\theta_2 & -sin \ m\theta_2 & \dots &0 &0 \\ 0 & 0 &sin \ m\theta_2 & cos \ m\theta_2 & \dots &0 &0 \\ \vdots & \vdots& \vdots & \vdots & \ddots & \vdots & \vdots\\ 0 & 0 &0 &0 & \dots & cos \ m\theta_{d/2} & -sin \ m\theta_{d/2} \\ 0 & 0 &0 &0 & \dots &sin \ m\theta_{d/2} & cos \ m\theta_{d/2} \end{pmatrix} RΘ,md= cos mθ1sin mθ10000sin mθ1cos mθ1000000cos mθ2sin mθ20000sin mθ2cos mθ2000000cos mθd/2sin mθd/20000sin mθd/2cos mθd/2

       则有:
q m T k n = ( R Θ , m d W q x m ) T ( R Θ , n d W k x n ) = x T W q R Θ , ( m − n ) d W k x n q^T_mk_n=(R^d_{\Theta,m}W_{q}x_m)^T(R^d_{\Theta,n}W_{k}x_n)=x^TW_qR^d_{\Theta,(m-n)}W_{k}x_n qmTkn=(RΘ,mdWqxm)T(RΘ,ndWkxn)=xTWqRΘ,(mn)dWkxn

       即为给位置 m m m n n n的词嵌入向量 x m x_m xm x n x_n xn乘以旋转矩阵 R Θ , m d R^d_{\Theta,m} RΘ,md R Θ , n d R^d_{\Theta,n} RΘ,nd得到的自注意力自动包含位置信息。 R Θ d R^d_{\Theta} RΘd是一个正交矩阵,不会改变向量的模长,因此它不会改变原模型的稳定性。

       由于 R Θ d R^d_{\Theta} RΘd是个包含大量 0 0 0元素的稀疏矩阵,直接用矩阵乘法来实现会很浪费算力,故可采用如下计算方法加速计算:
R Θ , m d x = ( x 1 x 2 x 3 x 4 ⋮ x d − 1 x d ) ⊗ ( c o s   m θ 1 c o s   m θ 1 c o s   m θ 2 c o s   m θ 2 ⋮ c o s   m θ d / 2 c o s   m θ d / 2 ) + ( − x 2 x 1 − x 4 x 3 ⋮ − x d x d − 1 ) ⊗ ( s i n   m θ 1 s i n   m θ 1 s i n   m θ 2 s i n   m θ 2 ⋮ s i n   m θ d / 2 s i n   m θ d / 2 ) R^d_{\Theta,m}x=\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} RΘ,mdx= x1x2x3x4xd1xd cos mθ1cos mθ1cos mθ2cos mθ2cos mθd/2cos mθd/2 + x2x1x4x3xdxd1 sin mθ1sin mθ1sin mθ2sin mθ2sin mθd/2sin mθd/2

远程衰减

       对向量 q q q k k k按维度 d d d成对分组,则RoPE的内积可转换为:
( R Θ , m d W q x m ) T ( R Θ , n d W k x n ) = R e [ ∑ i = 0 d / 2 − 1 q [ 2 i : 2 i + 1 ] k [ 2 i : 2 i + 1 ] ∗ e i ( m − n ) θ i ] (R^d_{\Theta,m}W_{q}x_m)^T(R^d_{\Theta,n}W_{k}x_n)=Re[\sum^{d/2-1}_{i=0}q_{[2i:2i+1]}k^*_{[2i:2i+1]}e^{i(m-n)\theta_i}] (RΘ,mdWqxm)T(RΘ,ndWkxn)=Re[i=0d/21q[2i:2i+1]k[2i:2i+1]ei(mn)θi]

       其中 q [ 2 i : 2 i + 1 ] q_{[2i:2i+1]} q[2i:2i+1]表示 q q q 2 i ∽ 2 i + 1 2i \backsim 2i+1 2i2i+1维度的值。令 h i = q [ 2 i : 2 i + 1 ] k [ 2 i : 2 i + 1 ] ∗ h_i=q_{[2i:2i+1]}k^*_{[2i:2i+1]} hi=q[2i:2i+1]k[2i:2i+1] S j = ∑ i = 0 j − 1 e i ( m − n ) θ i S_j=\textstyle\sum^{j-1}_{i=0}e^{i(m-n)\theta_i} Sj=i=0j1ei(mn)θi,则有:
∑ 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 ) \sum^{d/2-1}_{i=0}q_{[2i:2i+1]}k^*_{[2i:2i+1]}e^{i(m-n)\theta_i}= \sum^{d/2-1}_{i=0}h_i(S_{i+1}-S_i) i=0d/21q[2i:2i+1]k[2i:2i+1]ei(mn)θi=i=0d/21hi(Si+1Si)

       令 h d / 2 = 0 h_{d/2}=0 hd/2=0 S 0 = 0 S_0=0 S0=0,利用阿贝尔变换得:
∑ 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 ) \sum^{d/2-1}_{i=0}h_i(S_{i+1}-S_i)=[h_{d/2}S_{d/2}-h_0S_0]-\sum^{d/2-1}_{i=0}S_{i+1}(h_{i+1}-h_i) \\ =-\sum^{d/2-1}_{i=0}S_{i+1}(h_{i+1}-h_i) i=0d/21hi(Si+1Si)=[hd/2Sd/2h0S0]i=0d/21Si+1(hi+1hi)=i=0d/21Si+1(hi+1hi)

       根据不等式 ∣ a + b ∣ ≤ ∣ a ∣ + ∣ b ∣ |a+b|≤|a|+|b| a+ba+b得:
∣ ∑ 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 ) ∣ = ∑ 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 ∣ \begin{vmatrix} \sum^{d/2-1}_{i=0}q_{[2i:2i+1]}k^*_{[2i:2i+1]}e^{i(m-n)\theta_i} \end{vmatrix}= \begin{vmatrix} \sum^{d/2-1}_{i=0}S_{i+1}(h_{i+1}-h_i) \end{vmatrix} \\ \le \sum^{d/2-1}_{i=0} \begin{vmatrix} S_{i+1}(h_{i+1}-h_i)\end{vmatrix} =\sum^{d/2-1}_{i=0} \begin{vmatrix} S_{i+1}\end{vmatrix} \begin{vmatrix}(h_{i+1}-h_i)\end{vmatrix} \\ \le ( \max_i\begin{vmatrix} h_{i+1}-h_i\end{vmatrix})\sum^{d/2-1}_{i=0} \begin{vmatrix} S_{i+1}\end{vmatrix} i=0d/21q[2i:2i+1]k[2i:2i+1]ei(mn)θi = i=0d/21Si+1(hi+1hi) i=0d/21 Si+1(hi+1hi) =i=0d/21 Si+1 (hi+1hi) (imax hi+1hi )i=0d/21 Si+1

       故设 θ i = 1000 0 − 2 i / d \theta_i=10000^{-2i/d} θi=100002i/d,如下图可视化 1 d / 2 ∑ i = 0 d / 2 − 1 ∣ S i ∣ \frac 1 {d/2}\sum^{d/2-1}_{i=0} \begin{vmatrix} S_{i}\end{vmatrix} d/21i=0d/21 Si 的值,可知其随相对距离 m − n m-n mn的增加而衰减。故可知随着相对距离的变大,RoPE的内积结果有衰减趋势的出现,这满足长距token间相关性减弱的假设。

image

整合线性自注意力

       线性Attention的常见形式为:
A t t e n t i o n ( Q , K , V ) i = ∑ j = 1 n s i m ( q i , k j ) v j ∑ j = 1 n s i m ( q i , k j ) = ∑ j = 1 n ϕ ( q i ) T φ ( k j ) v j ∑ j = 1 n ϕ ( q i ) T φ ( k j ) Attention(Q,K,V)_i=\frac {\sum^n_{j=1}sim(q_i,k_j)v_j} {\sum^n_{j=1}sim(q_i,k_j)}=\frac {\sum^n_{j=1}\phi(q_i)^T\varphi(k_j)v_j} {\sum^n_{j=1}\phi(q_i)^T\varphi(k_j)} Attention(Q,K,V)i=j=1nsim(qi,kj)j=1nsim(qi,kj)vj=j=1nϕ(qi)Tφ(kj)j=1nϕ(qi)Tφ(kj)vj

        ϕ , φ \phi,\varphi ϕ,φ为非负激活函数,由于其计算方式也为内积,故可直接将RoPE整合进去:
A t t e n t i o n ( Q , K , V ) i = ∑ j = 1 n [ W i ϕ ( q i ) ] T [ W j φ ( k j ) ] v j ∑ j = 1 n [ W i ϕ ( q i ) ] T [ W j φ ( k j ) ] Attention(Q,K,V)_i=\frac {\sum^n_{j=1}[W_i\phi(q_i)]^T[W_j\varphi(k_j)]v_j} {\sum^n_{j=1}[W_i\phi(q_i)]^T[W_j\varphi(k_j)]} Attention(Q,K,V)i=j=1n[Wiϕ(qi)]T[Wjφ(kj)]j=1n[Wiϕ(qi)]T[Wjφ(kj)]vj

       但 [ W i ϕ ( q i ) ] T [ W j φ ( k j ) ] [W_i\phi(q_i)]^T[W_j\varphi(k_j)] [Wiϕ(qi)]T[Wjφ(kj)]可能为负数,分母有为0的风险,考虑 W i , W j W_i,W_j Wi,Wj是正交矩阵,不改变向量的模长,因此可以抛弃常规的注意力归一化要求,使用如下形式的线性Attention:
A t t e n t i o n ( Q , K , V ) i = ∑ j = 1 n [ W i ϕ ( q i ) ] T [ W j φ ( k j ) ] v j ∑ j = 1 n ϕ ( q i ) T φ ( k j ) Attention(Q,K,V)_i=\frac {\sum^n_{j=1}[W_i\phi(q_i)]^T[W_j\varphi(k_j)]v_j} {\sum^n_{j=1}\phi(q_i)^T\varphi(k_j)} Attention(Q,K,V)i=j=1nϕ(qi)Tφ(kj)j=1n[Wiϕ(qi)]T[Wjφ(kj)]vj

实验

机器翻译

       实验对self-attention层进行了一些修改,以使RoPE能够适应其学习过程。翻译任务基于联合源和目标字节对编码(BPE) ,用37k的词汇表复制了英德翻译任务的设置。在评估过程中,通过平均最后5个checkpoint得到模型。结果使用大小为4,长度惩罚为0.6的beam search。用Adam优化器 β 1 = 0.9 β_1 = 0.9 β1=0.9, β 2 = 0.98 β_2 = 0.98 β2=0.98进行优化,学习率从 1 e − 7 1e^ {− 7} 1e7 5 e − 4 5e^{ − 4} 5e4线性增加,然后与步骤数的反平方根成比例衰减,采用了0.1的label-smoothing。用测试集上的BLEU分数作为最终指标。

image

       上图显示在WMT 2014英语到德语翻译任务中,与Transformer相比,RoPE给出了更好的BLEU分数。

预训练语言模型

       使用Huggingface数据集的BookCorpus和Wikipedia语料库Foundation 进行预训练,语料以8:2的比例划分为训练集和验证集,使用训练过程的掩码语言建模(MLM)损失值作为评估指标。采用BERT作为基线模型。对于RoFormer,将BETR的自注意力块中的正弦位置编码替换为RoPE。对BERT和RoFormer进行了训练,batch size大小为64,最大序列长度为512,步骤为100k。使用学习率为1e-5的AdamW优化器。

image

       上图显示预训练期间的MLM损失。与普通的BERT相比,RoFormer的收敛速度更快。

对GLUE任务进行微调

       实验对MRPC,SST-2,QNLI,STS-B,QQP和MNLI预训练后微调查看GLUE。在MRPC和QQP数据集上使用F1-score,在STS-B数据集上使用spearman相关性,在剩余数据集上使用准确率作为评价指标。使用Huggingface Transformers库对上述每个下游任务进行3个epoch的微调,最大序列长度为512,批量大小为32,学习率为2,3,4,5e-5。

image

       上图显示RoFormer在6个数据集中的3个数据集上的性能明显优于BERT。

线性注意力整合RoPE

       实验在Enwik8数据集Mahoney上进行测试,该数据集来自英文维基百科,除英文文本外,还包括标记、特殊字符和其他语言的文本。将RoPE合并到12层12个头768维的PerFormer。报告了在相同设置下,有和没有RoPE时预训练过程的损失曲线。学习率 1 e − 4 1e-4 1e4、 batch size128和固定的最大序列长度1024。

       结果如figure3右图,观察到在相同的训练步数下,将RoPE代入Performer可以快速收敛和更低的损失。

中文数据测试

       在长度超过512个字符的中文长文本上进行了实验,对WoBERT进行了一些修改,用RoPE替换绝对位置嵌入。作为与其他基于transformer的中文预训练模型BERT、WoBERT和NEZHA的交叉比较。实验从中文维基百科、新闻和论坛收集的大约34GB数据上对RoFormer进行了预训练。预训练分多个阶段进行,改变batch size和最大输入序列长度,以使模型适应各种场景。

image

       上图为实验配置,abs和rel表示绝对位置嵌入和相对位置嵌入。

image

       上图显示RoFormer的准确率随着序列长度上界的增加而提高,这验证了RoFormer在处理长文本方面的能力。

下游任务

       实验选择Chinese AI and Law 2019 Similar Case Matching数据集在语义文本匹配方面的能力。CAIL2019-SCM收录最高人民法院公布的案件三联稿8964篇。输入三元组表示为(A, B, C),是对三种情况的事实描述。任务是在预定义的相似性度量下预测(A, B)是否比(A, C)更接近。按照6:2:2的比例划分训练集、验证集和测试集。

image

       上图显示在相同的预训练数据上,将该模型与预训练的BERT和WoBERT模型进行对比。对于短文本截断,即512个,RoFormer的结果与WoBERT相当,略优于BERT的实现。当最大输入文本长度增加到1024时,RoFormer的性能比WoBERT提高了1.5%。

reference

Su, J. , Lu, Y. , Pan, S. , Wen, B. , & Liu, Y. . (2021). Roformer: enhanced transformer with rotary position embedding.

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值