Rotary position embedding

The paper was published in 2021 as a preliminary report, and in late 2023 it has been updated to v5. The core idea is the same. The only thing that’s changed is basically evaluating on English benchmarks instead of Chinese data only.

This is probably one of the papers with the worst English writing skills I’ve read. It’s kinda iconic since this is an LLM paper, but LLM is not being used to audit the writing of the paper.

Why relative

This is my understanding, not what’s in the paper. We don’t need the absolute position information that’s done in Absolute position embedding. What we care about is the relative distance when doing attention. That means the y can generalize better on longer sequence (position 5 to 8 is the same as 500 to 503).

Formulation

Say we generate the positional embedding via function $f$,

$$\begin{array}{rl}{q}_m& =f_q({\mathbf{x}}_m,m)\\ {\mathbf{k}}_n& =f_k({\mathbf{x}}_n,n)\\ {\mathbf{v}}_n& =f_v({\mathbf{x}}_n,n),\end{array}$$

$k$ and $v$ are both on $n$ position because that’s a way of understanding the weighted sum. And then they go through softmax like this

$$a_{m,n}=\frac{\exp (\frac{{\mathbf{q}}_m^{⊺}{\mathbf{k}}_n}{\sqrt{d}})}{{\sum }_{j=1}^N\exp (\frac{{\mathbf{q}}_m^{⊺}{\mathbf{k}}_j}{\sqrt{d}})}$$

For the traditional absolute position embedding, that’s

$$f_{t:t\in \{q,k,v\}}({\mathbf{x}}_i,i):={\mathbf{W}}_{t:t\in \{q,k,v\}}({\mathbf{x}}_i+{\mathbf{p}}_i),$$

Since we want to capture the relative information and attention, it would be good if that $q_m^T{k}_n$ depend only on the relative position between $m$ and $n$. So the problem becomes: can we find such an $f$, such that

$$⟨f_q({\mathbf{x}}_m,m),f_k({\mathbf{x}}_n,n)⟩=g({\mathbf{x}}_m,{\mathbf{x}}_n,m-n).$$

One can find a solution to our formulation, when $d=2$ is:

$$\begin{array}{rl}{f}_q({\mathbf{x}}_m,m)& =({\mathbf{W}}_q{\mathbf{x}}_m)e^{im\theta }\\ f_k({\mathbf{x}}_n,n)& =({\mathbf{W}}_k{\mathbf{x}}_n)e^{in\theta }\\ g({\mathbf{x}}_m,{\mathbf{x}}_n,m-n)& =\mathrm{Re}[({\mathbf{W}}_q{\mathbf{x}}_m)({\mathbf{W}}_k{\mathbf{x}}_n{)}^{\ast }{e}^{i(m-n)\theta }]\end{array}$$

where $\mathrm{Re}[\cdot ]$ is the real part of a complex number and $({\mathbf{W}}_k{\mathbf{x}}_n{)}^{\ast }$ represents the conjugate complex number of $({\mathbf{W}}_k{\mathbf{x}}_n)$. $\theta \in \mathbb{R}$ is a preset non-zero constant. We can further write $f_{\{}q,k\}$ in a multiplication matrix:

$$f_{\{q,k\}}({\mathbf{x}}_m,m)=\left(\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)\left(\begin{array}{cc}{W}_{\{q,k\}}^{(11)}& W_{\{q,k\}}^{(12)}\\ W_{\{q,k\}}^{(21)}& W_{\{q,k\}}^{(22)}\end{array}\right)\left(\begin{array}{c}{x}_m^{(1)}\\ x_m^{(2)}\end{array}\right)$$

We can see intuitively that this work since this is rotating the embedding, or, assigning them angles in 2D plane. A vector at $\theta$ and a vector at $\gamma$, when doing dot product, gives us $\cos (\theta -\gamma )$.

In order to generalize our results in 2D to any $x_i\in {\mathbb{R}}^d$ where $d$ is even, we divide the d-dimension space into $d/2$ sub-spaces and combine them in the merit of the linearity of the inner product, turning $f_{\{q,k\}}$ into:

$$f_{\{q,k\}}(x_m,m)=R_{\Theta ,m}^d{W}_{\{q,k\}}{x}_m$$

where

$$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)$$

is the rotary matrix with pre-defined parameters $\Theta =\{{\theta }_i=10000^{-2(i-1)/d},i\in [1,2,...,d/2]\}$. One can think about it as “encoding relative information for each pair of size 2 in the embedding”.

Properties of RoPE

• Similar as the OG position embedding, it has a long term decay property (that’s because only the “going down” part of the cosine is the dominant part)
• It can be used easily with linear attention

Appendix A: linearity of the inner product

Let’s consider two vectors, $\mathbf{x}$ and $\mathbf{y}$, in a $d$-dimensional space, where $d$ is an even number.

$$\mathbf{x}=(x_1,x_2,x_3,x_4,\dots ,x_{d-1},x_d)$$ $$\mathbf{y}=(y_1,y_2,y_3,y_4,\dots ,y_{d-1},y_d)$$

The standard inner product (dot product), denoted by $⟨\mathbf{x},\mathbf{y}⟩$, is defined as the sum of the element-wise products of their components:

$$⟨\mathbf{x},\mathbf{y}⟩=\sum _{i=1}^d{x}_i{y}_i=x_1{y}_1+x_2{y}_2+x_3{y}_3+\cdots +x_d{y}_d$$

Now, RoPE treats the $d$-dimensional vector as a concatenation of $d/2$ smaller, 2-dimensional vectors. Let’s denote these sub-vectors with a prime symbol (′):

• ${{\mathbf{x}}^{\prime }}_1=(x_1,x_2)$, ${{\mathbf{x}}^{\prime }}_2=(x_3,x_4)$, …, ${{\mathbf{x}}^{\prime }}_{d/2}=(x_{d-1},x_d)$
• ${{\mathbf{y}}^{\prime }}_1=(y_1,y_2)$, ${{\mathbf{y}}^{\prime }}_2=(y_3,y_4)$, …, ${{\mathbf{y}}^{\prime }}_{d/2}=(y_{d-1},y_d)$

Because of the basic rules of addition, we can simply regroup the terms in the original inner product sum:

$$⟨\mathbf{x},\mathbf{y}⟩=(x_1{y}_1+x_2{y}_2)+(x_3{y}_3+x_4{y}_4)+\cdots +(x_{d-1}{y}_{d-1}+x_d{y}_d)$$

Notice that each term in parentheses is just the inner product of the corresponding 2D sub-vectors:

• $(x_1{y}_1+x_2{y}_2)=⟨{{\mathbf{x}}^{\prime }}_1,{{\mathbf{y}}^{\prime }}_1⟩$
• $(x_3{y}_3+x_4{y}_4)=⟨{{\mathbf{x}}^{\prime }}_2,{{\mathbf{y}}^{\prime }}_2⟩$

This leads us to the core identity that RoPE exploits. The inner product in $d$-dimensions is precisely the sum of the inner products in the $d/2$ constituent 2D subspaces:

$$⟨\mathbf{x},\mathbf{y}⟩=\sum _{i=1}^{d/2}⟨{{\mathbf{x}}^{\prime }}_i,{{\mathbf{y}}^{\prime }}_i⟩$$