Enhanced Transformer with Rotary Position Embedding (RoPE)
When we process sequences, the order of tokens matters (e.g words in a sentence).
Pre-transformer approaches, such as RNNs, encode order implicitly by sequentially updating their hidden state.
Transformers, on the other hand, are equivariant to token order permutation when this position information is not explicitly injected.
That is, for a transformer model \(\operatorname{encoder}\), input matrix \(\mathbf{X} \in \mathbb{R}^{T \times d}\) and permutation \(\mathbf{P}\), it holds that \[\operatorname{encoder}(\mathbf{P} \cdot \mathbf{X}) = \mathbf{P} \cdot \operatorname{encoder}(\mathbf{X}),\] meaning that changing token order does not impact the computed representations.
Therefore, we must find a way to inject positional information.
Recall the attention mechanism \[\mathbf{O} = \operatorname{softmax}\left(\frac{\mathbf{QK}^T}{\sqrt{d}}\right)\mathbf{V},\] a useful way to think about this “position problem” is to say that individual query and key vectors are generated by functions \[ \mathbf{q}_i = \mathbf{f}_q(\mathbf{x}_i, i) \qquad \mathbf{k}_j = \mathbf{f}_k(\mathbf{x}_j, j), \] where each \(\mathbf{f}_{\{q, k\}}\) takes as input the “raw” vector \(\mathbf{x}_i\) and position \(i\), and outputs a new vector with injected positional information.
In the original “Attention Is All You Need” paper1 the authors use absolute positional encoding \[ \mathbf{f}_q(\mathbf{x}_i, i) = \mathbf{W}^{(q)}(\mathbf{x}_i + \mathbf{p}_i) \qquad \mathbf{f}_k(\mathbf{x}_j, j) = \mathbf{W}^{(k)}(\mathbf{x}_j + \mathbf{p}_j) \] where \(\mathbf{p}_i\) is a vector generated deterministically (but it can also be learned).
The RoPE approach differs in two ways
- we use relative position (between \(i\) and \(j\)) instead of absolute
- we multiply (rotate) the input vector rather than adding
Let’s first look at what happens in our product \(\mathbf{QK}^T\), for a single pair of vectors \((\mathbf{q}_i,\mathbf{k}_j)\) \[\mathbf{q}_i^T\mathbf{k}_j = \mathbf{f}_q(\mathbf{x}_i, i)^T \mathbf{f}_k(\mathbf{x}_j, j).\]
Now, what is the constraint that the equation above should satisfy if we want to encode relative position information ?
\(\to\) The product of each pair of query and key vectors should be the result of a function \(g\) that only takes the relative position \(j - i\) as input \[\mathbf{f}_q(\mathbf{x}_i, i)^T \mathbf{f}_k(\mathbf{x}_j, j) = \mathbf{g}(\mathbf{x}_i, \mathbf{x}_j, j - i).\]
The authors further require \[\mathbf{f}_q(\mathbf{x}, 0) = \mathbf{W}^{(q)}\mathbf{x} \qquad \mathbf{f}_k(\mathbf{x}, 0) = \mathbf{W}^{(k)}\mathbf{x}.\]
Then, they show that a solution for the 2D case \(\mathbf{x} = [x_1, x_2]\) is \[\mathbf{f}_{\{q, k\}}(\mathbf{x}_j, j) = \begin{bmatrix} \cos(j\theta) & -\sin(j\theta)\\ \sin(j\theta) & \cos(j\theta) \end{bmatrix} \cdot \mathbf{W}^{\{q, k\}} \cdot \mathbf{x}_j, \] and prove that this formulation satisfies the above constraints.
To understand why, we denote the introduced rotation matrix as \(\mathbf{R}(i, \theta)\) and again take a look at \[\mathbf{f}_q(\mathbf{x}_i, i)^T \mathbf{f}_k(\mathbf{x}_j, j) = \left(\mathbf{R}(i, \theta) \cdot \mathbf{W}^{(q)} \cdot \mathbf{x}_i\right)^T\left(\mathbf{R}(j, \theta) \cdot \mathbf{W}^{(k)} \cdot \mathbf{x}_j\right),\] with middle product \[ \mathbf{R}(i, \theta)^T\, \mathbf{R}(j, \theta) = \begin{bmatrix} \cos(i\theta) & \sin(i\theta)\\ -\sin(i\theta) & \cos(i\theta) \end{bmatrix} \cdot \begin{bmatrix} \cos(j\theta) & -\sin(j\theta)\\ \sin(j\theta) & \cos(j\theta) \end{bmatrix} = \begin{bmatrix} \cos((i-j)\theta) & \sin((i-j)\theta)\\ -\sin((i-j)\theta) & \cos((i-j)\theta) \end{bmatrix} = \mathbf{R}(j-i, \theta). \]
This is pretty neat: using individual rotations that are function of absolute position, the dot-product ends up encoding relative position only \[ \mathbf{q}_i^T\mathbf{k}_j = \mathbf{x}_i^T \cdot \mathbf{W}^{(q)T} \cdot \mathbf{R}(j-i, \theta) \cdot \mathbf{W}^{(k)} \cdot \mathbf{x}_j. \]
The authors then proceed to generalize for the \(d = 2k\) case with a sparse block-diagonal matrix, and also give a form that is more efficient to compute.
They also prove a long-term decay property, matching the intuition that tokens that are further apart should have less connection.
Core properties of RoPE
- acts on query and key vectors, after their projections, this is why it is applied at every layer and not just once at the beginning
- does not inject position information to value vectors
- it is multiplicative (rotates vectors) rather than additive
- injects relative positional information, attention scores depend only on relative offset
Footnotes
Where the authors already highlight some desirable properties of positional encodings, such as the model being able to attend to relative positions and extrapolating to longer sequences than seen during training.↩︎