In a previous article, "The Road to Transformer Upgrade: 2. Rotary Position Embedding (RoPE) Integrating Various Strengths", we proposed the Rotary Position Embedding (RoPE) and the corresponding Transformer model, RoFormer. Since my primary research field is NLP, I initially thought this matter was settled. However, recently, Transformer models have become very popular in the field of Computer Vision (CV), with various Vision Transformers (ViTs) emerging one after another. This raises a question: what should RoPE look like in a two-dimensional (2D) scenario?
At first glance, this seems to be a simple generalization of the 1D case, but the derivations and understanding involved are far more complex than one might imagine. This article provides an analysis of this topic to deepen our understanding of RoPE.
2D RoPE
What is a 2D position? What does the corresponding 2D RoPE look like? Where does the difficulty lie? In this section, we first briefly introduce 2D positions, then directly provide the results and derivation logic for 2D RoPE. In the subsequent sections, we will provide the detailed derivation process.
2D Positions
In NLP, linguistic position information is one-dimensional; in other words, we need to tell the model which word in a sentence a specific token is. However, in CV, image position information is two-dimensional; we need to tell the model which row and column a feature is located in. Here, "two-dimensional" refers to the fact that two numbers are required to fully describe the position information, not the dimensionality of the position vector itself.
Some readers might think: can’t we just flatten the image and treat it as 1D? That doesn’t work well. For example, in an h \times h feature map, the position (x, y) becomes xh + y after flattening. The positions (x+1, y) and (x, y+1) become xh + y + h and xh + y + 1, respectively. The differences between these and xh + y are h and 1. However, according to our intuitive understanding, the distances from (x+1, y) and (x, y+1) to (x, y) should be the same. Flattening them results in inconsistent differences of h and 1, which is inherently unreasonable.
Therefore, we need to design position embeddings specifically for the 2D case, rather than simply flattening them into 1D.
The Standard Answer
After the derivations presented later, one solution for 2D RoPE is: \boldsymbol{\mathcal{R}}_{x,y}=\left( \begin{array}{cc:cc} \cos x\theta & -\sin x\theta & 0 & 0 \\ \sin x\theta & \cos x\theta & 0 & 0 \\ \hdashline 0 & 0 & \cos y\theta & -\sin y\theta \\ 0 & 0 & \sin y\theta & \cos y\theta \\ \end{array}\right) \label{eq:rope-2d} This solution is easy to understand: it is a block matrix composed of two 1D RoPEs. In implementation, it involves splitting the input vector into two halves, applying 1D RoPE for x to one half and 1D RoPE for y to the other. From this form, it is not difficult to generalize to 3D, 4D, or higher-dimensional positions.
The matrix in Eq. [eq:rope-2d] is an orthogonal matrix that satisfies two key properties:
1. Relativity: That is, \boldsymbol{\mathcal{R}}_{x_1,y_1}^{\top}\boldsymbol{\mathcal{R}}_{x_2,y_2}=\boldsymbol{\mathcal{R}}_{x_2-x_1,y_2-y_1}. It is precisely because of this property that RoPE possesses the ability to realize relative positions through absolute positions.
2. Invertibility: Given \boldsymbol{\mathcal{R}}_{x,y}, one can solve for x and y. This means the encoding of position information is lossless.
In a sense, Eq. [eq:rope-2d] is the simplest solution satisfying these two properties. While other slightly different solutions exist, they are more complex in form and implementation.
Derivation Logic
In hindsight, RoPE essentially found a matrix \boldsymbol{\mathcal{R}}_n=\begin{pmatrix}\cos n\theta & -\sin n\theta\\ \sin n\theta & \cos n\theta\end{pmatrix} that satisfies the "relativity" condition: \boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n=\boldsymbol{\mathcal{R}}_{n-m} \label{eq:re} Thus, it is natural to expect that the basic requirement for 2D RoPE is also to satisfy relativity, i.e., to find a matrix \boldsymbol{\mathcal{R}}_{x,y} such that it satisfies the 2D relativity condition \boldsymbol{\mathcal{R}}_{x_1,y_1}^{\top}\boldsymbol{\mathcal{R}}_{x_2,y_2}=\boldsymbol{\mathcal{R}}_{x_2-x_1,y_2-y_1}. However, if this were the only requirement, there would be many feasible solutions, such as simply letting: \boldsymbol{\mathcal{R}}_{x,y} = \begin{pmatrix}\cos (x+y)\theta & -\sin (x+y)\theta\\ \sin (x+y)\theta & \cos (x+y)\theta\end{pmatrix} The problem with this solution is that we cannot uniquely determine (x, y) from x+y, meaning this choice is lossy for position information. Therefore, we need the additional "invertibility" to ensure that the original position signals can be losslessly reconstructed from the position matrix.
To this end, we have two natural paths: 1. Quaternions; 2. Matrix Exponentials. We will introduce them one by one.
Quaternions
In the derivation of 1D RoPE, we primarily used complex numbers as a tool. Quaternions are a generalization of complex numbers and retain many of their properties, so using them to derive 2D RoPE seems like a natural idea. Unfortunately, this is a dead end, but I am still including the thought process here for reference.
Complex Numbers and Matrices
In high school, we learned that a complex number a+b\boldsymbol{i} corresponds one-to-one with a 2D vector (a,b) (to align with quaternions later, the imaginary unit \boldsymbol{i} is bolded). However, this correspondence only holds for addition and subtraction (since vectors lack a general multiplication operation). A more sophisticated correspondence is between complex numbers and matrices: a+b\boldsymbol{i} \quad \leftrightarrow \quad \begin{pmatrix} a & -b \\ b & a \end{pmatrix} Under this mapping, the addition, subtraction, multiplication, and division of complex numbers correspond one-to-one with those of matrices. For example: \begin{array}{ccc} (a+b\boldsymbol{i})(c+d\boldsymbol{i}) &=& (ac - bd) + (ad + bc)\boldsymbol{i} \\[5pt] \begin{pmatrix} a & -b \\ b & a \end{pmatrix}\begin{pmatrix} c & -d \\ d & c \end{pmatrix} &=& \begin{pmatrix} ac - bd & - ad - bc \\ ad + bc & ac - bd \end{pmatrix} \end{array} Thus, the matrix mapping is a complete isomorphism of the field of complex numbers, while the vector mapping is merely an intuitive geometric interpretation.
The matrix mapping of complex numbers is also an important foundation for RoPE. In the previous post, we derived that the complex representation of RoPE is \boldsymbol{q}e^{n\boldsymbol{i}\theta}=(\cos n\theta + \boldsymbol{i}\sin n\theta)\boldsymbol{q}. According to the matrix mapping, \cos n\theta + \boldsymbol{i}\sin n\theta corresponds to the matrix: \boldsymbol{\mathcal{R}}_n=\begin{pmatrix}\cos n\theta & -\sin n\theta\\ \sin n\theta & \cos n\theta\end{pmatrix} which gives the matrix form of 1D RoPE.
Introduction to Quaternions
As mentioned, quaternions are a generalization of complex numbers. In fact, they are the "ancestors" of matrices; historically, quaternions existed before general matrix operations and inspired many of them.
While a complex number is a 2D vector, a quaternion is a 4D vector, represented as a+b\boldsymbol{i}+c\boldsymbol{j}+d\boldsymbol{k}, where \boldsymbol{i}^2=\boldsymbol{j}^2=\boldsymbol{k}^2=-1, but they are distinct from each other. The operation rules between the bases are:
| \times | 1 | \boldsymbol{i} | \boldsymbol{j} | \boldsymbol{k} |
|---|---|---|---|---|
| 1 | 1 | \boldsymbol{i} | \boldsymbol{j} | \boldsymbol{k} |
| \boldsymbol{i} | \boldsymbol{i} | -1 | \boldsymbol{k} | -\boldsymbol{j} |
| \boldsymbol{j} | \boldsymbol{j} | -\boldsymbol{k} | -1 | \boldsymbol{i} |
| \boldsymbol{k} | \boldsymbol{k} | \boldsymbol{j} | -\boldsymbol{i} | -1 |
At the time, the biggest shock they gave people was non-commutativity, e.g., \boldsymbol{i}\boldsymbol{j}=-\boldsymbol{j}\boldsymbol{i}\neq \boldsymbol{j}\boldsymbol{i}. Otherwise, they are highly similar to complex number operations.
For instance, there is an Euler’s formula similar to that of complex numbers: e^{a+b\boldsymbol{i}+c\boldsymbol{j}+d\boldsymbol{k}} = e^a\left(\cos r + \frac{b\boldsymbol{i}+c\boldsymbol{j}+d\boldsymbol{k}}{r}\sin r\right) \label{eq:euler} where r = \Vert b\boldsymbol{i}+c\boldsymbol{j}+d\boldsymbol{k}\Vert = \sqrt{b^2+c^2+d^2}. Furthermore, there is a similar matrix mapping: a+b\boldsymbol{i}+c\boldsymbol{j}+d\boldsymbol{k} \quad \leftrightarrow \quad \begin{pmatrix} a & -b & -c & -d \\ b & a & -d & c \\ c & d & a & -b \\ d & -c & b & a \end{pmatrix} \label{eq:mapping}
Violation of Relativity
The origins of these formulas are a long story. With Euler’s formula and the exponential mapping, one might think: since 1D RoPE is just the matrix mapping of e^{n\boldsymbol{i}\theta}, wouldn’t 2D RoPE just be the matrix mapping of e^{x\boldsymbol{i}\theta + y\boldsymbol{j}\theta}?
I thought so at first, but unfortunately, it is wrong. Why? In 1D RoPE, we utilized the complex representation of the inner product: \langle\boldsymbol{q},\boldsymbol{k}\rangle=\text{Re}[\boldsymbol{q}\boldsymbol{k}^*] This identity also holds for quaternions. Then we used complex exponentials: \langle\boldsymbol{q}e^{m\boldsymbol{i}\theta},\boldsymbol{k}e^{n\boldsymbol{i}\theta}\rangle=\text{Re}\left[\left(\boldsymbol{q}e^{m\boldsymbol{i}\theta}\right)\left(\boldsymbol{k}e^{n\boldsymbol{i}\theta}\right)^*\right]=\text{Re}\left[\boldsymbol{q}e^{m\boldsymbol{i}\theta}e^{-n\boldsymbol{i}\theta}\boldsymbol{k}^*\right]=\text{Re}\left[\boldsymbol{q}e^{(m-n)\boldsymbol{i}\theta}\boldsymbol{k}^*\right] The first two equalities can be carried over to quaternions, but the key is that the third equality does not hold generally in quaternions! Generally, for two quaternions \boldsymbol{p} and \boldsymbol{q}, the equation e^{\boldsymbol{p}+\boldsymbol{q}}=e^{\boldsymbol{p}}e^{\boldsymbol{q}} is not true. More broadly, for any two objects whose multiplication is non-commutative, e^{\boldsymbol{p}+\boldsymbol{q}}\neq e^{\boldsymbol{p}}e^{\boldsymbol{q}} usually holds.
Therefore, because exponential multiplication cannot be converted into addition, relativity cannot be guaranteed. Thus, the path of derivation via quaternions was aborted.
Matrix Exponentials
The matrix mapping of quaternions shows that quaternions represent a specific cluster of 4\times 4 matrices. If the quaternion derivation fails, perhaps general matrix analysis will work. This is indeed the case. In this section, we will use matrix exponentials to provide a derivation.
Matrix Exponentials
The matrix exponential here is not an element-wise operation using the exponential function as an activation function in neural networks, but rather an operation defined by a power series: \exp \boldsymbol{B} = \sum_{k=0}^{\infty}\frac{\boldsymbol{B}^k}{k!} where \boldsymbol{B}^k refers to the matrix multiplication of k copies of \boldsymbol{B}.
The matrix exponential is a very important matrix operation. It can directly express the solution to a system of linear differential equations with constant coefficients \frac{d}{dt}\boldsymbol{x}_t=\boldsymbol{A}\boldsymbol{x}_t: \boldsymbol{x}_t = \big(\exp t\boldsymbol{A}\big)\boldsymbol{x}_0 For the derivation of RoPE, we primarily use the following property of the matrix exponential: \boldsymbol{A}\boldsymbol{B} = \boldsymbol{B}\boldsymbol{A} \quad\Rightarrow\quad \big(\exp \boldsymbol{A}\big)\big(\exp \boldsymbol{B}\big) = \exp \big(\boldsymbol{A} + \boldsymbol{B}\big) \label{eq:expm-ex} That is, if the multiplication of \boldsymbol{A} and \boldsymbol{B} is commutative, the matrix exponential can convert multiplication into addition, just like the exponential of numbers. Note that this is a sufficient but not necessary condition.
1D General Solution
Why can we link RoPE with matrix exponentials? Because 1D RoPE has a simple exponential expression: \boldsymbol{\mathcal{R}}_n=\begin{pmatrix}\cos n\theta & -\sin n\theta\\ \sin n\theta & \cos n\theta\end{pmatrix}=\exp\left\{n\theta\begin{pmatrix}0 & -1\\ 1 & 0\end{pmatrix}\right\} \label{eq:rope-exp} So I began to consider matrices of the following form as solutions for RoPE: \boldsymbol{\mathcal{R}}_n=\exp n\boldsymbol{B} where \boldsymbol{B} is a matrix independent of n. A necessary condition for RoPE is satisfying the "relativity" condition [eq:re], so we analyze: \big(\exp m\boldsymbol{B}\big)^{\top}\big(\exp n\boldsymbol{B}\big) = \big(\exp m\boldsymbol{B}^{\top}\big)\big(\exp n\boldsymbol{B}\big) Assuming \boldsymbol{B}^{\top} and \boldsymbol{B} are commutative, then according to Eq. [eq:expm-ex]: \big(\exp m\boldsymbol{B}^{\top}\big)\big(\exp n\boldsymbol{B}\big) = \exp \big(m\boldsymbol{B}^{\top} + n\boldsymbol{B}\big) To make m\boldsymbol{B}^{\top} + n\boldsymbol{B}=(n-m)\boldsymbol{B}, we only need: \boldsymbol{B}^{\top} = - \boldsymbol{B} This is the constraint given by "relativity." We previously assumed \boldsymbol{B}^{\top} and \boldsymbol{B} are commutative; now we can verify that any \boldsymbol{B} satisfying \boldsymbol{B}^{\top} = -\boldsymbol{B} (skew-symmetric) must commute with its transpose, so the result is self-consistent.
This means that for any matrix \boldsymbol{B} satisfying \boldsymbol{B}^{\top} + \boldsymbol{B} = 0, \exp n\boldsymbol{B} is a solution to Eq. [eq:re], and it can be proven to be an orthogonal matrix. More directly, since \exp n\boldsymbol{B}=\left(\exp \boldsymbol{B}\right)^n, we find that for any orthogonal matrix \boldsymbol{O}, \boldsymbol{\mathcal{R}}_n=\boldsymbol{O}^n is a solution to Eq. [eq:re].
For 2\times 2 matrices, the general solution to \boldsymbol{B}^{\top} + \boldsymbol{B} = 0 is \boldsymbol{B}=\begin{pmatrix}0 & -\theta\\ \theta & 0\end{pmatrix}, which leads to the solution in Eq. [eq:rope-exp].
2D Constraints
Similarly, for 2D RoPE, we consider: \boldsymbol{\mathcal{R}}_{x,y}=\exp \big(x\boldsymbol{B}_1 + y\boldsymbol{B}_2\big) as a candidate solution. Repeating the derivation for the "relativity" condition: assuming x_1\boldsymbol{B}_1^{\top} + y_1\boldsymbol{B}_2^{\top} and x_2\boldsymbol{B}_1 + y_2\boldsymbol{B}_2 are commutative, we obtain the following constraints: \boldsymbol{B}_1^{\top} + \boldsymbol{B}_1 = 0,\quad \boldsymbol{B}_2^{\top} + \boldsymbol{B}_2 = 0 However, the commutativity of x_1\boldsymbol{B}_1^{\top} + y_1\boldsymbol{B}_2^{\top} and x_2\boldsymbol{B}_1 + y_2\boldsymbol{B}_2 implies that (\boldsymbol{B}_1,\boldsymbol{B}_1^{\top}), (\boldsymbol{B}_2,\boldsymbol{B}_2^{\top}), (\boldsymbol{B}_1,\boldsymbol{B}_2^{\top}), and (\boldsymbol{B}_2,\boldsymbol{B}_1^{\top}) must all be commutative. The first two constraints only guarantee the commutativity of the first two pairs. We must add the remaining commutativity as a constraint: \left\{\begin{aligned} &\boldsymbol{B}_1^{\top} + \boldsymbol{B}_1 = 0\\ &\boldsymbol{B}_2^{\top} + \boldsymbol{B}_2 = 0\\ &\boldsymbol{B}_1 \boldsymbol{B}_2^{\top} = \boldsymbol{B}_2^{\top} \boldsymbol{B}_1 \end{aligned}\right. \label{eq:2d-conds} It is easy to prove that under the first two conditions, the new constraint is equivalent to \boldsymbol{B}_1 \boldsymbol{B}_2 = \boldsymbol{B}_2 \boldsymbol{B}_1.
RoPE Emerges
Since a 2\times 2 matrix satisfying the first two conditions has only one independent parameter, it cannot satisfy "invertibility" for two variables (x, y). Therefore, we must consider at least 3\times 3 matrices, which have 3 independent parameters: \begin{pmatrix}0 & -a & -b \\ a & 0 & -c \\ b & c & 0\end{pmatrix} To ensure invertibility, we might assume \boldsymbol{B}_1 and \boldsymbol{B}_2 are "orthogonal," for example: \boldsymbol{B}_1=\begin{pmatrix}0 & -a & 0 \\ a & 0 & 0 \\ 0 & 0 & 0\end{pmatrix},\quad\boldsymbol{B}_2=\begin{pmatrix}0 & 0 & -b \\ 0 & 0 & -c \\ b & c & 0\end{pmatrix} Without loss of generality, let a=1. Solving Eq. [eq:2d-conds] yields b=0, c=0, meaning \boldsymbol{B}_2 must be a zero matrix, which does not meet our requirements. The Mathematica code for this is:
B[a_, b_, c_] = {{0, -a, -b}, {a, 0, -c}, {b, c, 0}};
B1 = B[1, 0, 0];
B2 = B[0, b, c];
Solve[{Dot[B1, B2] == Dot[B2, B1]}, {b, c}]Therefore, we must consider at least 4\times 4 matrices, which have 6 independent parameters. Without loss of generality, consider an orthogonal decomposition: \boldsymbol{B}_1=\begin{pmatrix}0 & -a & -b & 0 \\ a & 0 & -c & 0 \\ b & c & 0 & 0 \\ 0 & 0 & 0 & 0\end{pmatrix},\quad\boldsymbol{B}_2=\begin{pmatrix}0 & 0 & 0 & -d \\ 0 & 0 & 0 & -e \\ 0 & 0 & 0 & -f \\ d & e & f & 0\end{pmatrix} Solving this yields: d=cf,\quad e=-bf Solving code:
B[a_, b_, c_, d_, e_, f_] = {{0, -a, -b, -d}, {a, 0, -c, -e}, {b, c, 0, -f}, {d, e, f, 0}};
B1 = B[1, b, c, 0, 0, 0];
B2 = B[0, 0, 0, d, e, f];
Solve[{Dot[B1, B2] == Dot[B2, B1]}, {b, c, d, e, f}]We find that the result places no constraints on f. For simplicity, we can let f=1 and set the remaining b, c, d, e to 0. At this point: \boldsymbol{\mathcal{R}}_{x,y}=\exp \,\begin{pmatrix}0 & -x & 0 & 0 \\ x & 0 & 0 & 0 \\ 0 & 0 & 0 & -y \\ 0 & 0 & y & 0\end{pmatrix} Adding a parameter \theta and expanding gives: \boldsymbol{\mathcal{R}}_{x,y}=\exp \,\left\{\begin{pmatrix}0 & -x & 0 & 0 \\ x & 0 & 0 & 0 \\ 0 & 0 & 0 & -y \\ 0 & 0 & y & 0\end{pmatrix}\theta\right\}=\left( \begin{array}{cc:cc} \cos x\theta & -\sin x\theta & 0 & 0 \\ \sin x\theta & \cos x\theta & 0 & 0 \\ \hdashline 0 & 0 & \cos y\theta & -\sin y\theta \\ 0 & 0 & \sin y\theta & \cos y\theta \\ \end{array}\right)
Extended Story
This concludes the introduction to the derivation of 2D RoPE. Now, readers might ask: how effective is it? Unfortunately, there are no complete experimental results yet. After all, I haven’t worked on ViT before, and the derivation of 2D RoPE was only recently completed, so progress is slow. I can only say that preliminary results show it is quite effective. Members of the EleutherAI team have also experimented with this scheme, and the results were better than other existing position encodings.
Speaking of the EleutherAI team, I’ll add a few more words. EleutherAI is the team that gained fame for attempting to "reproduce GPT-3." After we proposed RoPE and RoFormer in our previous article, we were fortunate to catch their attention. They conducted many supplementary experiments, confirming that RoPE is more effective than many other position encodings (refer to their blog post "Rotary Embeddings: A Relative Revolution"). This prompted us to complete the English paper "RoFormer: Enhanced Transformer with Rotary Position Embedding" and submit it to Arxiv. The questions regarding 2D RoPE also initially originated from the EleutherAI team.
Summary
This article introduced our 2D generalization of RoPE, primarily using "relativity" and "invertibility" as starting points to determine the final form of 2D RoPE. We attempted two derivation processes—quaternions and matrix exponentials—and ultimately provided the solution through matrix exponentials. Through this derivation process, we can further deepen our understanding of RoPE.