12. Infinitely Extrapolatable ReRoPE?

Since introducing the idea of mixed radix in “The Road to Transformer Upgrade: 11. Taking $\beta$-Radix Position to the End” to further generalize NTK-aware Scaled RoPE, the author feels that the effect of similar ideas has reached its limit, and a more significant improvement must be sought through alternative means. At this point, the author recalled an idea conceived earlier, which was set aside due to its high complexity. Now that a bottleneck has been encountered, “the only way is the best way,” so it was picked up again.

Unexpectedly, although this method increases some inference complexity, its experimental results are surprisingly good—even hinting at infinite length extrapolation capability! Therefore, the author is eager to write this article to share this method. Due to its formal similarity with the ReLU activation function, the author names this method “ReRoPE (Rectified Rotary Position Embeddings)”.

Revisit

#

We know that RoPE is formally an absolute position encoding, but in practice, it brings relative position information to Attention, specifically the following Toeplitz matrix:

Here $L$ is the current sample length. When $L$ significantly exceeds the training length, the additional positions have not been trained, so the performance cannot be guaranteed. This is why direct extrapolation (Length Extrapolation) usually performs poorly.

Later, researchers proposed Position Interpolation, which is equivalent to changing the relative position matrix to:

In this way, by adjusting $k$, the maximum relative position can be guaranteed not to 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 again. Precisely because it avoids extrapolation, it requires far fewer fine-tuning steps than direct extrapolation (neural networks are often better at interpolation than extrapolation).

As for the later proposed NTK-aware Scaled RoPE, it takes a “sideways approach,” cleverly distributing the extrapolation pressure across each dimension, so it can achieve good results without fine-tuning. However, it still relies on extrapolation, which is something neural networks are not good at, so there is an upper limit to its effectiveness. In the author’s experiments, its Long Context performance still cannot get very close to the training performance.

Fusion

#

We can also examine these methods from the perspective of language model locality. Locality means that when a language model predicts the next token, it clearly relies more on neighboring tokens. Direct extrapolation maintains locality (position encoding near 0 remains unchanged), but performs poorly because it introduces position encodings beyond the training length; Position Interpolation, while not extrapolating position encodings, disrupts locality (position encodings near 0 are compressed by $1/k$), so it does not perform well without fine-tuning; NTK-aware Scaled RoPE, through “high-frequency extrapolation, low-frequency interpolation,” implicitly combines the advantages of both, ensuring locality without obvious position encoding extrapolation, thus achieving good results even without fine-tuning.

Is there a method that can 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$, and outside the window, we use a position interval of size $1/k$. The entire relative position matrix is as follows:

As long as $w$ is less than the training length, by controlling $k$, we can ensure that all position encodings do not exceed the training length, while precisely maintaining locality. This simply and directly combines direct extrapolation and position interpolation.

In particular, the matrices above have a special case: when $k\to\infty$, it simplifies to

In this case, regardless of the input length, the range of its position encoding does not exceed $w$, so this is a scheme that potentially supports arbitrary length Context!

Formally, the relationship between the matrices above and the standard RoPE matrix is equivalent to the relationship between ReLU, Leaky ReLU, and Linear. Therefore, the author calls the former “ReRoPE (Rectified RoPE)” and the latter “Leaky ReRoPE”.

Computation

#

In fact, similar ideas are not difficult to come up with. Previous relative position encodings based on Attention Bias (such as Classic Relative Position Encoding, T5 Position Encoding) often use such block operations. However, unlike these relative position encodings, implementing such block operations in RoPE will significantly increase the computation, which is the main reason why this idea was set aside by the author.

How to understand the increase in computation? We know that RoPE achieves “relative position through absolute position,” which can only result in linear relative positions. The matrices above are nonlinear (or piecewise linear). To implement it, one can only calculate the Attention matrix twice and then combine them. Specifically, first calculate the Attention matrix (before Softmax) once using standard RoPE:

Here the first equality is the implementation method, and the second equality is the equivalent result. $\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):

If it is ReRoPE, it is simpler:

Finally, based on the condition $i - j < w$, combine them:

Whether it’s ReRoPE or Leaky ReRoPE, it inevitably requires calculating the Attention matrix twice (if there is a more efficient implementation, please enlighten me), which is one of the increased computational costs. Furthermore, the need to customize the Attention matrix calculation means that existing flash attention implementations cannot be directly used, thus relatively increasing computational cost.

On the other hand, also due to the nonlinear relative position, during autoregressive decoding, the Key sequence cache can only store the sequence before RoPE. Then, at each decoding step, the corresponding RoPE needs to be added to the entire Key sequence. This modification will also increase inference computational cost. The only good news is that during token-by-token decoding, starting from the second step, the length of the Query sequence is 1. At this time, only the RoPE for the Key sequence needs to be customized, so the Attention matrix can be calculated only once:

Experiments

#

Continuing with the settings from “The Road to Transformer Upgrade: 11. Taking $\beta$-Radix Position to the End”, we conducted experiments on ReRoPE, and the results are shown in the table below:

As mentioned at the beginning of the article, the extrapolation effect of ReRoPE without fine-tuning is surprisingly good. It not only significantly surpasses the previous best NTK-RoPE-mixed but also clearly exceeds HFWA, which was pre-trained from scratch! Here, $\text{w256}$ refers to $w=256$, $\log n^{\dagger}$ means that $\log n$ scaling was not included during pre-training (like LLAMA), and during the test phase, each $\boldsymbol{q}_n$ is multiplied by $\max(1, \log_{\text{maxlen}} n)$, while $\log n$ means that the $\log n$ scaling factor was included during pre-training.

Here are some ablation experiments, showing that ReRoPE is quite robust to $w$, with the optimal value being roughly $1/4 \sim 1/2$ of the training length:

The table below compares ReRoPE and Leaky ReRoPE:

As a generalization of ReRoPE, fine-tuned Leaky ReRoPE has the potential to surpass ReRoPE, but the improvement is very slight. In addition, when $k$ is a finite value, the maximum length that can be handled is also limited, because we cannot know the total length to be generated in advance, so we can only preset a sufficiently large $k$. However, after setting it to a finite value, when the input is long enough, the performance will significantly drop 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 not significant.

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

As can be seen, ReRoPE truly achieved almost no loss in training performance (RoPE-4096 represents training performance) and satisfies the ideal characteristic of “longer context, lower loss” (more context should be more helpful for prediction). In addition, the author also tested the chat effect on the OpenBuddy open-source LLAMA2-13b fine-tuned model, and it felt quite good (tested with Context up to 20k tokens).

Finally, sharing the author’s code implementation of ReRoPE and Leaky ReRoPE based on the transformers LLAMA model. Readers can also load the LLAMA series models for testing:

Github: https://github.com/bojone/rerope

Summary (formatted)

#

In this article, the author proposes ReRoPE (Rectified RoPE), which is also a post-processing scheme for RoPE. Experimental results show that its zero-shot length extrapolation capability not only significantly surpasses the previous NTK-aware Scaled RoPE but also exceeds HFWA, which was specially designed and required training from scratch. Furthermore, unlike NTK-aware Scaled RoPE, whose performance drops sharply after a certain length, ReRoPE seems to perform well at any length. In addition to comparative experiments, the article also provides a reference implementation based on transformers-llama, which interested readers can test themselves.

@online{kexuefm-9708,
        title={The Road to Transformer Upgrade: 12. Infinitely Extrapolatable ReRoPE?},
        author={苏剑林},
        year={2023},
        month={08},
        url={\url{https://kexue.fm/archives/9708}},
}