Transformer Upgrade Path: 19. The Second Kind of Rotary Position Embedding
Readers who have been following the "Transformer Upgrade Path" series up to this post should already be familiar with rotary position embedding (RoPE). In short, RoPE is a rotation transformation applied to the Query ($\boldsymbol{Q}$) and Key ($\boldsymbol{K}$) in attention. Formally it is an absolute position encoding, but combined with the dot-product nature of attention, it automatically achieves the effect of relative position encoding.
So, can RoPE also be applied to the Value ($\boldsymbol{V}$)? At first glance, it seems not, because rotating $\boldsymbol{V}$ would break the relative position property. However, things are not so absolute — this post discusses applying RoPE to $\boldsymbol{V}$, which we might call "the second kind of rotary position embedding."
Basics Recap
Let's decompose dot-product attention as
\begin{equation}\boldsymbol{o}_i = \sum_j a_{i,j}\boldsymbol{v}_j,\qquad a_{i,j} = \frac{e^{s_{i,j}}}{\sum\limits_j e^{s_{i,j}}},\qquad s_{i,j} = \boldsymbol{q}_i^{\top}\boldsymbol{k}_j\end{equation}more
For simplicity, we omit the scaling factor for $s_{i,j}$ here. RoPE is applied on $\boldsymbol{q}_i,\boldsymbol{k}_j$:
\begin{equation}\boldsymbol{q}_i \to \boldsymbol{\mathcal{R}}_i\boldsymbol{q}_i,\qquad \boldsymbol{k}_j \to \boldsymbol{\mathcal{R}}_j\boldsymbol{k}_j\end{equation}
This causes the attention logits, i.e., $s_{i,j}$, to become
\begin{equation}s_{i,j} = (\boldsymbol{\mathcal{R}}_i\boldsymbol{q}_i)^{\top} (\boldsymbol{\mathcal{R}}_j\boldsymbol{k}_j) = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}_i^{\top}\boldsymbol{\mathcal{R}}_j\boldsymbol{k}_j=\boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}_{j-i}\boldsymbol{k}_j\end{equation}
In other words, $s_{i,j}$ depends only on the relative position $j-i$, thereby achieving the effect of relative position encoding through an absolute-position form. This transformation exploits a property of rotation matrices, $\boldsymbol{\mathcal{R}}_i^{\top}\boldsymbol{\mathcal{R}}_j=\boldsymbol{\mathcal{R}}_{j-i}$.
Besides rotation matrices, in Transformer Upgrade Path: 4. Two-Dimensional Rotary Position Embedding we proved that the general solution is $\boldsymbol{\mathcal{R}}_i = \boldsymbol{O}^i$, where $\boldsymbol{O}$ is an arbitrary orthogonal matrix and the superscript denotes matrix exponentiation. Later, however, in Transformer Upgrade Path: 6. Completeness Analysis of Rotary Position Embedding, we also showed that the general orthogonal-matrix solution is essentially isomorphic to the rotation-matrix solution.
A New Usage
What if we apply RoPE to $\boldsymbol{v}_j$, i.e., $\boldsymbol{v}_j\to\boldsymbol{\mathcal{R}}_j\boldsymbol{v}_j$? Clearly, the attention result becomes
\begin{equation}\boldsymbol{o}_i = \sum_j a_{i,j} \boldsymbol{\mathcal{R}}_j\boldsymbol{v}_j\label{eq:v-rope-abs}\end{equation}
This causes attention to depend explicitly on the absolute position $j$. If we only want some form of position encoding, this may not be a big deal, but if we want relative position encoding specifically, this fails our purpose.
However, there is a simple trick that fixes this defect! We can apply an inverse RoPE to $\boldsymbol{o}_i$ once more:
\begin{equation}\boldsymbol{o}_i = \boldsymbol{\mathcal{R}}_i^{\top}\left(\sum_j a_{i,j} \boldsymbol{\mathcal{R}}_j\boldsymbol{v}_j\right)=\sum_j a_{i,j} \boldsymbol{\mathcal{R}}_i^{\top}\boldsymbol{\mathcal{R}}_j\boldsymbol{v}_j=\sum_j a_{i,j} \boldsymbol{\mathcal{R}}_{j-i}\boldsymbol{v}_j\label{eq:vo-rope}\end{equation}
This turns it back into a relative position encoding! Formally, it is likewise composed of two absolute position encodings, in the same spirit as the existing RoPE, so we call it "the second kind of rotary position embedding," or more intuitively, "VO-RoPE," since it applies RoPE once each to Value and Output. Correspondingly, the standard RoPE can be called "QK-RoPE."
A Simple Experiment
We ran a quick round of experiments on a roughly 1B-parameter LLAMA-like model, comparing the following settings:
1. NoPE: no position encoding at all;
2. QK-RoPE: the standard rotary position embedding;
3. VO-RoPE: the second kind of rotary position embedding proposed in this post;
4. Q/K/V/O-RoPE: RoPE applied to only one of Q, K, V, or O;
5. QKV-RoPE: RoPE applied to all of Q, K, and V;
6. QKVO-RoPE: RoPE applied to all of Q, K, V, and O.
Note that settings 4 and 5 both count as absolute position encoding. The rough conclusion is:
$$\text{QK-RoPE}\approx \text{QKVO-RoPE} > \text{K-RoPE}\approx \text{VO-RoPE} > \text{QKV-RoPE} > \text{NoPE} > \text{Q/V/O-RoPE}$$
The specific loss-function differences are:
$$\begin{array}{c|c} \hline & \text{Loss} \\ \hline \text{QK-RoPE} & 2.712 \\ \text{QKVO-RoPE} & 2.719 \\ \text{K-RoPE} & 2.769 \\ \text{VO-RoPE} & 2.770 \\ \text{QKV-RoPE} & 2.783 \\ \text{NoPE} & 2.795 \\ \text{O-RoPE} & 2.841 \\ \text{Q-RoPE} & 2.851 \\ \text{V-RoPE} & 2.856 \\ \hline \end{array}$$
Some Reflections
From the results above, we can see that VO-RoPE outperforms NoPE but falls short of QK-RoPE, and stacking VO-RoPE on top of QK-RoPE brings no additional gain. So it might seem that VO-RoPE doesn't need to be proposed at all?
In my view, completing the picture of how RoPE can be used — answering the question "can RoPE be applied to Value?" — and then experimentally establishing that "there's no real benefit" is valuable in its own right. Moreover, in the long run, it may not necessarily remain useless forever; it's just that under our current mainstream language-model setups, its effect may not show up. When I originally proposed RoPE, the motivation was simply that it was fun — I never expected it to become such a competitive position encoding (what happened afterward was just good luck).
At present, VO-RoPE does have one potential application, related to the MLA discussed in The Ultimate Tug-of-War Between Cache and Performance: From MHA, MQA, GQA to MLA. We know that during inference, MLA is roughly equivalent to an MQA with shared K and V:
\begin{equation}\boldsymbol{o}_i = \sum_{j=1}^i a_{i,j}\boldsymbol{c}_j,\qquad a_{i,j} = \frac{e^{s_{i,j}}}{\sum\limits_{j=1}^i e^{s_{i,j}}},\qquad s_{i,j} = \exp(\boldsymbol{q}_i^{\top}\boldsymbol{c}_j)\end{equation}
This property means its KV cache only needs to store a single $\boldsymbol{c}$. However, this important property is incompatible with QK-RoPE, because once we apply RoPE to $\boldsymbol{c}_j$ inside the attention matrix, there are two possible outcomes:
1. If $\boldsymbol{c}_j$ on the Value side does not get RoPE, then K and V are no longer fully shared, which means either the KV cache has to be doubled (caching both pre- and post-RoPE versions), or K has to have RoPE injected on the fly (introducing latency);
2. If $\boldsymbol{c}_j$ on the Value side does get RoPE, K and V sharing can be preserved, but then it's no longer relative position encoding.
To resolve this issue, MLA adopts a "mostly NoPE + a small portion of RoPE" concatenation scheme. But, as we now know from the second kind of rotary position embedding described in this post, all that's needed is to additionally apply an O-RoPE to the Output:
\begin{equation}\boldsymbol{o}_i = \boldsymbol{\mathcal{R}}_i^{\top}\sum_{j=1}^i a_{i,j}(\boldsymbol{\mathcal{R}}_j\boldsymbol{c}_j),\qquad a_{i,j} = \frac{e^{s_{i,j}}}{\sum\limits_{j=1}^i e^{s_{i,j}}},\qquad s_{i,j} = (\boldsymbol{\mathcal{R}}_i\boldsymbol{q}_i)^{\top} (\boldsymbol{\mathcal{R}}_j\boldsymbol{c}_j)\end{equation}
That said, this idea hasn't been fully worked out yet — it can't be directly applied to MLA's training-time formulation. I'm just writing it down here for reference.
Related Work
In fact, VO-RoPE also elegantly provides an intermediate form between attention and complex-valued linear RNNs (such as LRU and RetNet). Starting from equation $\eqref{eq:vo-rope}$, consider the causal setting, and take the special case $a_{i,j}=\gamma^{i-j}$, where $0 < \gamma < 1$. Then we get
\begin{equation}\boldsymbol{o}_i = \sum_{j=1}^i \gamma^{i-j} \boldsymbol{\mathcal{R}}_{j-i}\boldsymbol{v}_j\end{equation}
We know that the rotation matrix $\boldsymbol{\mathcal{R}}_{j-i}$, written in complex form, is simply the diagonal matrix of $e^{\mathbb{I}\theta (j - i)}$, where $\mathbb{I}$ is the imaginary unit (i.e., $\mathbb{I}^2=-1$) — to distinguish it from $i,j$'s $i$, we write it here as $\mathbb{I}$. In that case, the above equation is equivalent to
\begin{equation}\boldsymbol{o}_i = \sum_{j=1}^i \gamma^{i-j} e^{\mathbb{I}\theta (j - i)} \boldsymbol{v}_j = \sum_{j=1}^i (\gamma e^{-\mathbb{I}\theta})^{i-j} \boldsymbol{v}_j\end{equation}
which is precisely the simplest linear RNN with complex-valued decay. Based on the derivation in Google's New Work Tries to "Revive" RNNs: Can RNNs Shine Again?, such an RNN is theoretically more complete than one with purely real-valued decay.
So, adding the VO-RoPE form to RoPE amounts to a general extension from real-valued linear RNNs to complex-valued linear RNNs, theoretically making it more complete in capability — even though this greater completeness may not necessarily help on language modeling tasks, much as the complex-valued LRU shows no advantage over the purely real-valued RWKV. Still, theoretical completeness might carry hidden value for certain specific scenarios — who knows.
Side note: After sharing this post on Twitter, several readers mentioned that they had previously tried VO-RoPE, including:
1. @gharik, who said he had tried VO-RoPE before and obtained some positive results, calling it "RoPER" at the time — more details can be found here and here;
2. @vinam_arora, who noted that he tried QKVO-RoPE on a "brain decoding task" and also got positive results, in the paper A Unified, Scalable Framework for Neural Population Decoding.
Summary
This post has been built around the question "can RoPE be applied to V?", discussing this second way of using RoPE.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.