Transformer Upgrade Road: 12. ReRoPE with Infinite Extrapolation?

Review

We know that RoPE is formally an absolute position encoding, but in practice, it brings relative position information to the Attention mechanism, specifically the following Toeplitz matrix: \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 \\ \text{\small $L - 2$} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\ \text{\small $L - 1$} & \text{\small $L - 2$} & \ddots & \ddots & \ddots & 3 & 2 & 1 & 0 & \\ \end{pmatrix} \label{eq:rope} \end{equation} Here L is the current sample length. When L significantly exceeds the training length, the extra positions have not been trained, so their performance cannot be guaranteed. This is why direct length extrapolation usually performs poorly.

Later, researchers proposed Position Interpolation (PI), which is equivalent 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 \\ \text{\small $\frac{L-2}{k}$} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\ \text{\small $\frac{L-1}{k}$} & \text{\small $\frac{L-1}{k}$} & \ddots & \ddots & \ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\ \end{pmatrix} \end{equation} In this way, by adjusting k, we can ensure that the maximum relative position does not exceed the training length, thus avoiding extrapolation. However, it makes the position information more "crowded," so a certain number of fine-tuning steps are required for the model to work properly again. Because it avoids extrapolation, the number of fine-tuning steps required is much smaller than that for direct extrapolation (neural networks are often better at interpolation than extrapolation).

As for the subsequently proposed NTK-aware Scaled RoPE, it takes a "clever" approach by distributing the extrapolation pressure across each dimension. Thus, it performs well even without fine-tuning, but it ultimately still relies on extrapolation—something neural networks are not good at—so its performance has an upper limit. In my experiments, its Long Context performance could not closely match the training effect.

Fusion

We can also examine these methods from the perspective of the locality of language models. Locality refers to the fact that when a language model predicts the next token, it clearly depends more on neighboring tokens. Direct extrapolation maintains locality (position encodings near 0 remain unchanged), but performs poorly because it introduces position encodings beyond the training length. Position interpolation avoids extrapolating position encodings but disrupts locality (position encodings near 0 are compressed to 1/k), so it does not perform well without fine-tuning. NTK-aware Scaled RoPE implicitly combines the advantages of both through "high-frequency extrapolation, low-frequency interpolation," ensuring locality without significant extrapolation of position encodings, thus yielding good results without fine-tuning.

Is there a way to combine extrapolation and interpolation more directly? Yes. We can set a window size w. Within the window, we use a position interval of size 1; outside the window, we use a position interval of size 1/k. The entire relative position matrix is as follows: \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}{\text{\small $w - 1$}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{w} & \color{red}{\text{\small $w - 1$}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\text{\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}{\text{\small $w + \frac{2}{k}$}} & \color{green}{\text{\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}{\text{\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}{\text{\small $w + \frac{2}{k}$}} & \color{green}{\text{\small $w + \frac{1}{k}$}} & \color{green}{w} & \color{red}{\text{\small $w - 1$}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{\text{\small $w + \frac{L-1-w}{k}$}} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\text{\small $w + \frac{2}{k}$}} & \color{green}{\text{\small $w + \frac{1}{k}$}} & \color{green}{w} & \color{red}{\text{\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 all position encodings do not exceed the training length while precisely maintaining locality, simply and directly combining direct extrapolation and position interpolation.

Specifically, matrix [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}{\text{\small $w - 1$}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\ \color{green}{w} & \color{red}{\text{\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}{\text{\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}{\text{\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 the input length, the range of its position encodings does not exceed w. Therefore, this is a scheme that potentially supports Context of arbitrary length!

Formally, the relationship between matrices [eq:rerope], [eq:leaky-rerope] and the standard RoPE matrix [eq:rope] is analogous to the relationship between ReLU, Leaky ReLU, and Linear functions. Thus, I call [eq:rerope] "ReRoPE (Rectified RoPE)" and [eq:leaky-rerope] "Leaky ReRoPE."

Calculation

In fact, similar ideas are not difficult to conceive. Relative position encodings based on Attention Bias (such as classic relative position encoding or T5 position encoding) often feature such block-wise operations. However, unlike those, implementing such block-wise operations in RoPE significantly increases the computational load, which is the main reason I had previously set this idea aside.

How should we understand the increased computational load? We know that RoPE "achieves relative position through absolute position," which only yields linear relative positions. Matrices [eq:leaky-rerope] and [eq:rerope] are non-linear (or piecewise linear). To implement them, one must calculate the Attention matrix twice and then combine them. Specifically, first calculate the Attention matrix using standard RoPE (before Softmax): \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 method, and the second is the equivalent result, where \boldsymbol{\mathcal{R}} is the rotation matrix of RoPE. For simplicity, we omit the Attention scale factor. Next, we need to calculate the Attention matrix for RoPE with an interval of 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} If it is ReRoPE, it is 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, merge them based on the condition i - j < w: \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 is ReRoPE or Leaky ReRoPE, calculating the Attention matrix twice is unavoidable (if there is a more efficient implementation, please let me know). This is one of the sources of increased computation. Furthermore, the need for custom Attention matrix calculation means that existing flash attention implementations cannot be directly applied, further increasing the computational cost.

On the other hand, also due to the non-linear relative positions, the Key sequence cache during auto-regressive decoding can only store values before RoPE. Then, at each decoding step, the corresponding RoPE must be applied to the entire Key sequence. This change also increases inference computation. The only good news is that during token-by-token decoding, starting from the second step, the Query sequence length is 1. At this point, we only need to customize the RoPE for the Key sequence, so the Attention matrix only needs to be calculated 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 settings from "Transformer Upgrade Road: 11. Taking the \beta-radix Position to the End", we conducted experiments on ReRoPE. The results are shown in the table below:

As mentioned at the beginning, the extrapolation effect of ReRoPE without fine-tuning is surprisingly good. It not only significantly surpasses the previously optimal NTK-RoPE-mixed but also clearly exceeds HFWA, which was pre-trained from scratch! Here, \text{w256} refers to w=256. \log n^{\textcolor{red}{\dagger}} refers to cases where \log n scaling was not included during pre-training (like Llama), but during testing, each \boldsymbol{q}_n is multiplied by \max(1, \log_{\text{maxlen}} n). \log n refers to cases where the \log n scaling factor was included during pre-training.

Below are some ablation experiments showing that ReRoPE is quite robust regarding w. The optimal value is roughly between 1/4 and 1/2 of the training length:

The following table compares ReRoPE and Leaky ReRoPE:

As a generalization of ReRoPE, a fine-tuned Leaky ReRoPE has the potential to outperform ReRoPE, but the improvement is very slight. Furthermore, when k takes a finite value, the maximum length it can handle is also finite. Since we cannot know the total length to be generated in advance, we can only preset a sufficiently large k. However, once set to a finite value, when the input is long enough, the performance will drop significantly because the position encoding exceeds the training length. In contrast, ReRoPE does not have this risk. Overall, the value of fine-tuning Leaky ReRoPE compared to ReRoPE seems small.

The above experimental results were all tested on a GAU model with 100 million parameters. Below are the test results based on Llama2-13b (the metric is loss, lower is better), representing performance on a real LLM:

As can be seen, ReRoPE truly achieves almost no loss in training performance (RoPE-4096 represents the training effect) and satisfies the ideal characteristic of "longer context, lower loss" (more context should be more helpful for prediction). Additionally, I tested the chat performance on the Llama2-13b fine-tuned model open-sourced by OpenBuddy, and it felt quite good (I tested up to 20k tokens of Context).

Finally, I am sharing the code for implementing ReRoPE and Leaky ReRoPE based on the Llama model in the transformers library. Readers can load Llama series models for testing themselves:

Summary

In this article, I proposed ReRoPE (Rectified RoPE), which is also a post-processing scheme for RoPE. Experimental results show that its length extrapolation capability without fine-tuning not only significantly exceeds the previous NTK-aware Scaled RoPE but even surpasses the specially designed HFWA that requires training from scratch. Furthermore, unlike NTK-aware Scaled RoPE, whose capability drops sharply after exceeding a certain length, ReRoPE seems to perform well at any length. In addition to comparative experiments, the article provides a reference implementation based on transformers-llama for interested readers to test.

Original address: https://kexue.fm/archives/9708