Transformer Upgrade Path: 12. ReRoPE for Infinite Length Extrapolation?
Ever since Transformer Upgrade Path: 11. Pushing β-ary Positional Encodings to the Limit introduced the mixed-radix idea to further generalize NTK-aware Scaled RoPE, I've felt that the effectiveness of such approaches has hit its ceiling — to make substantial further progress, we need a genuinely different route. This reminded me of an idea I'd conceived some time ago but shelved because of its high implementation complexity. Now that we've hit a bottleneck, though, "the only way forward is the best way forward," so I decided to pick it back up.
Much to my surprise, although this method does add some inference overhead, its experimental performance is astonishingly good — it even seems to hint at unlimited length extrapolation capability! So I couldn't wait to write this post to share it. Given its formal resemblance to the ReLU activation function, I've named this method "ReRoPE (Rectified Rotary Position Embeddings)."
A Quick Recap
We know that RoPE is formally an absolute positional encoding, but in practice it endows attention with relative position information, in the form of the following Toeplitz matrix: more
\begin{equation}\begin{pmatrix}0 & \\ 1 & 0 & \\ 2 & 1 & 0 &\\ 3 & 2 & 1 & 0 & \\ \ddots & 3 & 2 & 1 & 0 & \\ \ddots & \ddots & 3 & 2 & 1 & 0 & \\ \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\ \small{L - 2} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\ \small{L - 1} & \small{L - 2} & \ddots & \ddots & \ddots & 3 & 2 & 1 & 0 & \\ \end{pmatrix}\label{eq:rope}\end{equation}
Here $L$ is the length of the current sample. When $L$ noticeably exceeds the training length, the extra positions that appear cannot be guaranteed to work well since they were never seen during training — this is why direct extrapolation (Length Extrapolation) usually performs poorly.
Later, researchers proposed Position Interpolation, which amounts to changing the relative position matrix to:
\begin{equation}\begin{pmatrix}0 & \\ \frac{1}{k} & 0 & \\ \frac{2}{k} & \frac{1}{k} & 0 &\\ \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\ \ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\ \ddots & \ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\ \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\ \small{\frac{L-2}{k}} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\ \small{\frac{L-1}{k}} & \small{\frac{L-1}{k}} & \ddots & \ddots & \ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\ \end{pmatrix}\end{equation}
This way, by simply adjusting $k$, one can ensure that the largest relative position never exceeds the training length, thereby avoiding extrapolation. However, it makes the positional information more "crowded," so some amount of fine-tuning is still needed to get the model working properly again. Precisely because extrapolation is avoided, though, the number of fine-tuning steps required is much smaller than with direct extrapolation (neural networks tend to be much better at interpolation than extrapolation).
As for the subsequently proposed NTK-aware Scaled RoPE, it takes a more unconventional route — cleverly spreading the extrapolation burden across every dimension, so that it works reasonably well even without fine-tuning. But it still fundamentally relies on extrapolation, which is something neural networks aren't good at, so its performance has a ceiling. In my experiments, its long-context performance still falls noticeably short of what's achieved with actual training.
Combining the Two
We can also examine these methods through the lens of the locality of language models. By locality, I mean that when a language model predicts the next token, it relies much more heavily on nearby tokens. Direct extrapolation preserves locality (the positional encoding near 0 is unchanged), and its poor performance stems from introducing positional encodings beyond the training length. Position interpolation, while not introducing any out-of-range positional encodings, disrupts locality (the positional encoding near 0 gets compressed to $1/k$), so it also performs poorly without fine-tuning. NTK-aware Scaled RoPE, through "high-frequency extrapolation, low-frequency interpolation," implicitly combines the advantages of both — it preserves locality while avoiding conspicuous extrapolated positional encodings, so it performs reasonably well even without fine-tuning.
Is there a more direct way to combine extrapolation and interpolation? Yes — we can set a window size $w$: within the window we use a position spacing of $1$, and outside the window we use a position spacing of $1/k$. The full relative position matrix then looks like this:
\begin{equation}\begin{pmatrix} \color{red}{0} & \\ \color{red}{1} & \color{red}{0} & \\ \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \\ \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\small{w + \frac{L-1-w}{k}}} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \end{pmatrix}\label{eq:leaky-rerope}\end{equation}
As long as $w$ is smaller than the training length, by controlling $k$ we can ensure that no positional encoding exceeds the training length, all while exactly preserving locality — a simple and direct combination of direct extrapolation and position interpolation.
In particular, matrix $\eqref{eq:leaky-rerope}$ has a special case: when $k\to\infty$, it simplifies to
\begin{equation}\begin{pmatrix} \color{red}{0} & \\ \color{red}{1} & \color{red}{0} & \\ \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{w} & \color{green}{w} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{w} & \color{green}{w} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\ddots} & \color{green}{w} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \\ \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{w} & \color{green}{w} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{w} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{w} & \color{green}{w} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \end{pmatrix}\label{eq:rerope}\end{equation}
In this case, no matter how long the input is, the range of its positional encodings never exceeds $w$, so this is potentially a scheme that supports context of arbitrary length!
Formally, the relationship between matrices $\eqref{eq:rerope}$, $\eqref{eq:leaky-rerope}$ and the standard RoPE matrix $\eqref{eq:rope}$ is analogous to the relationship between ReLU, Leaky ReLU, and a linear function — so I call $\eqref{eq:rerope}$ "ReRoPE (Rectified RoPE)," and $\eqref{eq:leaky-rerope}$ "Leaky ReRoPE."
Computation
Actually, ideas along these lines aren't hard to come up with — such piecewise computations often appear in earlier relative positional encodings based on attention bias (e.g., classic relative positional encoding, T5 positional encoding). However, unlike those relative positional encodings, implementing this kind of piecewise computation in RoPE noticeably increases the computational cost, which is the main reason I had shelved this idea.
Why does it increase the computational cost? Recall that RoPE "achieves relative position through absolute position," which can only yield a linear relative position — whereas matrices $\eqref{eq:leaky-rerope}$ and $\eqref{eq:rerope}$ are nonlinear (or rather, piecewise linear). To implement them we have no choice but to compute the attention matrix twice and then combine the results. Specifically, we first compute an attention matrix (before softmax) using standard RoPE:
\begin{equation}a_{i,j}^{(1)} = \left(\boldsymbol{\mathcal{R}}^i\boldsymbol{q}_i\right)^{\top}\left(\boldsymbol{\mathcal{R}}^j\boldsymbol{k}_j\right) = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^{j-i}\boldsymbol{k}_j\end{equation}
Here the first equality is the implementation, and the second equality is the equivalent result, where $\boldsymbol{\mathcal{R}}$ is the RoPE rotation matrix; for simplicity we omit the attention scale factor. Next, we need to compute the attention matrix for RoPE with spacing $1/k$ (Leaky ReRoPE):
\begin{equation}a_{i,j}^{(2)} = \left(\boldsymbol{\mathcal{R}}^{(i-w)/k+w}\boldsymbol{q}_i\right)^{\top}\left(\boldsymbol{\mathcal{R}}^{j/k}\boldsymbol{k}_j\right) = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^{(j-i+w)/k-w}\boldsymbol{k}_j\end{equation}
For ReRoPE, this is somewhat simpler:
\begin{equation}a_{i,j}^{(2)} = \left(\boldsymbol{\mathcal{R}}^w\boldsymbol{q}_i\right)^{\top}\boldsymbol{k}_j = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^w\boldsymbol{k}_j\end{equation}
Finally, based on the condition $i - j < w$, we merge the two:
\begin{equation}a_{i,j} = \left\{\begin{aligned} &a_{i,j}^{(1)},\quad (i - j < w) \\[8pt] &a_{i,j}^{(2)}, \quad (i - j \geq w) \end{aligned}\right.\end{equation}
Whether it's ReRoPE or Leaky ReRoPE, computing the attention matrix twice is unavoidable (if anyone knows of a more efficient implementation, I'd love to hear about it), and this is one source of the added computational cost. Moreover, the need for a custom attention-matrix computation means we can't directly plug in an off-the-shelf flash attention implementation, which adds further overhead.
On top of this, again because of the nonlinear relative position, during autoregressive decoding the key-sequence cache can only store the keys before RoPE is applied, and then at each decoding step, RoPE must be reapplied to the whole key sequence — this change also increases inference cost. The only good news is that during token-by-token decoding, from the second step onward the query sequence has length 1, so we only need to apply the customized RoPE to the key sequence, meaning the attention matrix only needs to be computed once:
\begin{equation}a_{i,j} = \left\{\begin{aligned} &\boldsymbol{q}_i^{\top}\left(\boldsymbol{\mathcal{R}}^{\max(j-i,-w)}\boldsymbol{k}_j\right), \quad(\text{ReRoPE})\\[8pt] &\boldsymbol{q}_i^{\top}\left(\boldsymbol{\mathcal{R}}^{\max(j-i,(j-i+w)/k-w)}\boldsymbol{k}_j\right), \quad(\text{Leaky ReRoPE}) \end{aligned}\right.\end{equation}
Experiments
Continuing with the setup from Transformer Upgrade Path: 11. Pushing β-ary Positional Encodings to the Limit, we ran experiments on ReRoPE, with results as shown in the table below:
$$\begin{array}{c|cc} \hline \text{test length} & 512(\text{training}) & 4096(\text{repeat}) & 4096(\text{no repeat})\\ \hline \text{Baseline} & 49.41\% & 24.17\% & 23.16\% \\ \text{Baseline-}\log n & 49.40\% & 24.60\% & 24.02\% \\ \hline \text{PI-RoPE} & 49.41\% & 15.04\% & 13.54\% \\ \text{PI-RoPE-}\log n & 49.40\% & 14.99\% & 16.51\% \\ \hline \text{NTK-RoPE-old} & 49.41\% & 51.28\% & 39.27\% \\ \text{NTK-RoPE-}\log n\text{-old} & 49.40\% & 61.71\% & 43.75\% \\ \hline \text{NTK-RoPE-fixed} & 49.41\% & 51.86\% & 39.61\% \\ \text{NTK-RoPE-}\log n^{\color{red}{\dagger}}\text{-fixed} & 49.41\% & 55.94\% & 41.11\% \\ \text{NTK-RoPE-}\log n\text{-fixed} & 49.40\% & 62.85\% & 44.14\% \\ \text{NTK-RoPE-mixed} & 49.41\% & 53.09\% & 40.12\% \\ \text{NTK-RoPE-}\log n^{\color{red}{\dagger}}\text{-mixed} & 49.41\% & 59.11\% & 42.38\% \\ \text{NTK-RoPE-}\log n\text{-mixed} & 49.40\% & 68.91\% & 45.41\% \\ \hline \text{ReRoPE-w256} & 49.41\% & 77.90\% & 48.48\% \\ \text{ReRoPE-w256-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.40\% & 48.85\% \\ \text{ReRoPE-w256-}\log n & 49.40\% & \boldsymbol{85.12\%} & \boldsymbol{49.07\%} \\ \hline \text{HFWA} & 48.70\% & 80.84\% & 48.15\% \\ \hline \end{array}$$
As mentioned at the start of this post, ReRoPE's performance under length extrapolation without any fine-tuning is remarkably good — it not only clearly surpasses the previously best NTK-RoPE-mixed, but also clearly surpasses HFWA, which was trained from scratch! Here $\text{w256}$ refers to $w=256$, $\log n^{\color{red}{\dagger}}$ means that no $\log n$ scaling was applied during pretraining (as in LLAMA), and at test time every $\boldsymbol{q}_n$ is multiplied by $\max(1, \log_{\text{maxlen}} n)$, while $\log n$ means that the $\log n$ scaling factor was already incorporated during pretraining.
Below are some ablation experiments showing that ReRoPE is fairly robust with respect to $w$, with the optimal value being roughly $1/4\sim 1/2$ times the training length:
$$\begin{array}{c|cc} \hline \text{test length} & 512(\text{training}) & 4096(\text{repeat}) & 4096(\text{no repeat})\\ \hline \text{ReRoPE-w64} & 49.41\% & 69.39\% & 45.19\% \\ \text{ReRoPE-w64-}\log n^{\color{red}{\dagger}} & 49.41\% & 78.58\% & 47.42\% \\ \text{ReRoPE-w64-}\log n & 49.40\% & 84.38\% & 48.14\% \\ \hline \text{ReRoPE-w128} & 49.41\% & 76.11\% & 47.82\% \\ \text{ReRoPE-w128-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.28\% & 48.78\% \\ \text{ReRoPE-w128-}\log n & 49.40\% & \boldsymbol{85.47\%} & 48.87\% \\ \hline \text{ReRoPE-w256} & 49.41\% & 77.90\% & 48.48\% \\ \text{ReRoPE-w256-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.40\% & 48.85\% \\ \text{ReRoPE-w256-}\log n & 49.40\% & 85.12\% & \boldsymbol{49.07\%} \\ \hline \text{ReRoPE-w384} & 49.41\% & 70.72\% & 48.15\% \\ \text{ReRoPE-w384-}\log n^{\color{red}{\dagger}} & 49.41\% & 76.42\% & 48.31\% \\ \text{ReRoPE-w384-}\log n & 49.40\% & 83.24\% & 48.62\% \\ \hline \text{ReRoPE-w512} & 49.41\% & 7.09\% & 8.25\% \\ \text{ReRoPE-w512-}\log n^{\color{red}{\dagger}} & 49.41\% & 7.08\% & 8.25\% \\ \text{ReRoPE-w512-}\log n & 49.40\% & 15.84\% & 10.83\% \\ \hline \end{array}$$
The table below compares ReRoPE and Leaky ReRoPE:
$$\begin{array}{c|cc} \hline \text{test length} & 512(\text{training}) & 4096(\text{repeat}) & 4096(\text{no repeat})\\ \hline \text{ReRoPE-w128-}\log n & 49.40\% & \boldsymbol{85.47\%} & 48.87\% \\ \text{Leaky ReRoPE-w128-k64-}\log n & 49.40\% & 85.29\% & 48.96\% \\ \text{Leaky ReRoPE-w128-k32-}\log n & 49.40\% & 85.31\% & 49.03\% \\ \text{Leaky ReRoPE-w128-k16-}\log n & 49.40\% & 85.15\% & \boldsymbol{49.10\%} \\ \text{Leaky ReRoPE-w128-k8-}\log n & 49.40\% & 80.00\% & 48.11\% \\ \hline \text{ReRoPE-w256-}\log n & 49.40\% & 85.12\% & 49.07\% \\ \text{Leaky ReRoPE-w256-k64-}\log n & 49.40\% & 84.60\% & 49.03\% \\ \text{Leaky ReRoPE-w256-k32-}\log n & 49.40\% & 84.30\% & 48.97\% \\ \text{Leaky ReRoPE-w256-k16-}\log n & 49.40\% & 83.59\% & 48.87\% \\ \text{Leaky ReRoPE-w256-k8-}\log n & 49.40\% & 69.80\% & 45.72\% \\ \hline \end{array}$$
As a generalization of ReRoPE, a carefully tuned Leaky ReRoPE has a chance of surpassing ReRoPE, but the gain is marginal. Furthermore, when $k$ takes a finite value, the maximum length that can be handled is also finite, since we can't know in advance the total length that will be generated — we can only preset a sufficiently large $k$. But once this is set to a finite value, when the input becomes long enough, performance drops sharply because the positional encoding exceeds the training length; ReRoPE, by contrast, carries no such risk. Overall, fine-tuned Leaky ReRoPE doesn't seem to offer much value over ReRoPE.
All the experimental results above were tested on a 100-million-parameter GAU model. Below are results based on llama2-13b (the metric is loss, lower is better), which reflect performance on an actual LLM:
$$\begin{array}{c|cc} \hline \text{test length} & 4096(\text{training}) & 8192 & 16384\\ \hline \text{RoPE} & 1.4967 & 8.8615 & \text{-} \\ \text{NTK-RoPE} & 1.6081 & 1.5417 & 1.5163 \\ \text{ReRoPE} & 1.4996 & 1.4267 & 1.4001 \\ \hline \end{array}$$
As we can see, ReRoPE truly manages to preserve almost all of the original training performance (RoPE-4096 represents the training performance), and it satisfies the ideal property of "longer context, lower loss" (more context should indeed help prediction). I've also tested chat performance on the OpenBuddy open-source LLAMA2-13b fine-tuned model, and it feels quite good to me (I've tested contexts of up to 20k tokens so far).
Finally, here's the code where I implemented ReRoPE and Leaky ReRoPE on top of the LLAMA model in the transformers library; readers are welcome to load LLAMA-family models themselves and test it out:
Github: https://github.com/bojone/rerope
Summary
In this post, I've proposed ReRoPE (Rectified RoPE), which is likewise a post-processing scheme for RoPE. Experimental results show that its length-extrapolation ability without any fine-tuning not only clearly exceeds the previous NTK-aware Scaled RoPE, but even surpasses HFWA, a method specifically designed and requiring training from scratch. Moreover, unlike NTK-aware Scaled RoPE, whose performance drops sharply beyond a certain length, ReRoPE seems to perform well at essentially any length. Besides the comparative experiments, this post also provides a reference implementation based on transformers-LLAMA, which interested readers are welcome to try out for themselves.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.