Transformer Upgrade Path: 13. Inverting Leaky ReRoPE

Transformer升级之路：13、逆用Leaky ReRoPE 苏剑林 · 2023-08-14 · Machine Learning attention positional encoding generalization extrapolation RoPE

Last week, in Transformer Upgrade Path: 12. Infinitely Extrapolatable ReRoPE?, I proposed ReRoPE and Leaky ReRoPE. Numerous experimental results showed that they can extend an LLM's context length without fine-tuning, at almost no cost to the trained performance, and achieve the desirable property of "longer context, lower loss." Furthermore, unlike NTK-aware Scaled RoPE, ReRoPE in particular seemed to exhibit an essentially unbounded context-processing capability.

All in all, ReRoPE looks quite satisfying, but the fly in the ointment is that it increases inference cost—specifically, the first inference step needs to compute attention twice, and every subsequent step requires recomputing the positional encoding. This post attempts to solve this problem by inverting the use of Leaky ReRoPE during training.

Recap

Let's not tire of revisiting this: RoPE is formally an absolute positional encoding, but the effect it actually achieves is that of a relative positional encoding, with the corresponding relative position matrix given by:

\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}more

To avoid the problem of position indices going out of range in long contexts while retaining locality, Leaky ReRoPE changes the relative position matrix used at inference time 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}{\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}

where $w$ is the window width, typically taken to be $\frac{1}{4}$ to $\frac{1}{2}$ times the training length, and $k$ is used to control the maximum length that can be handled—generally it's best to keep $w + \frac{L-1-w}{k}$ from exceeding half the training length. As for ReRoPE, it is simply obtained by taking the limit $k\to\infty$:

\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}

Reversal

Judging from the evaluation results in the previous post, as a training-free extrapolation scheme, both ReRoPE and Leaky ReRoPE performed quite satisfyingly—neither losing performance within the training length nor failing to achieve "longer context, lower loss." The one drawback is that their inference speed is slower than that of ordinary attention, and they are not currently compatible with acceleration techniques such as Flash Attention.

So, can we flip this around? ReRoPE/Leaky ReRoPE run at normal RoPE speed during training but slow down at inference; flipping this means: can we make training slower so that inference reverts to conventional RoPE? Some readers might wonder: why would we want to make training slower—doesn't that increase training cost? The reason is that ReRoPE/Leaky ReRoPE are length-extrapolation methods, used in the "train short, test long" setting. The slowdown in training is short-term and controllable, whereas a slowdown in inference is long-term and hard to bear. So, given an equal amount of slowdown, we'd much rather place it in the training phase.

Let's look at Leaky ReRoPE again: during training it uses the relative position matrix of Eq. $\eqref{eq:rope}$ with a step size of 1, while during inference it uses a step size of $1$ within the window $w$ and a step size of $\frac{1}{k} < 1$ outside the window—that is, Eq. $\eqref{eq:leaky-rerope}$. In other words, the difference is that inference uses a smaller step size outside the window. If we flip this around and use Leaky ReRoPE during training with a step size outside the window that's larger than $1$, then following the principle "inference uses a smaller step size outside the window," could inference just use a step size equal to $1$ outside the window, thereby degenerating into ordinary RoPE?

I call this idea "InvLeaky ReRoPE" (Inverse Leaky ReRoPE). Without further ado, let's run some experiments to test it.

Experiments

Continuing with the previous experimental combination of "GAU + Deep Norm + Tiger + language model," we use Leaky ReRoPE with $k=1/16, w=128$ during training and ordinary RoPE during inference. The test results are as follows:

$$\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{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{InvLeaky ReRoPE-w128-}\log n & 49.38\% & 82.25\% & 48.32\% \\ \text{InvLeaky ReRoPE-w128-b8-}\log n & 49.62\% & 81.15\% & 48.85\% \\ \hline \text{HFWA} & 48.70\% & 80.84\% & 48.15\% \\ \hline \end{array}$$

Here $\text{b8}$ refers to changing RoPE's frequency base from 10000 to 80000. As we can see, while "Leaky ReRoPE → RoPE" InvLeaky ReRoPE doesn't perform as well as "RoPE → ReRoPE/Leaky ReRoPE," it still beats HFWA, and since inference uses conventional RoPE, it can take advantage of existing acceleration techniques—making it still quite competitive. In addition, I did some simple tuning of parameters such as $k,w,b$, and found that the optimal solution is basically the two combinations above—that is, setting $k$ to "the reciprocal of twice the extension factor," setting $w$ to $\frac{1}{4}$ times the training length, with $b$ optionally multiplied by the extension factor.

So how much does InvLeaky ReRoPE affect training speed? In the experiments above, the model has 100 million parameters and a training length of 512; the training time per 1000 steps increased from 330 seconds to 350 seconds—less than a 10% increase. Of course, part of this is due to GAU, since GAU uses single-head attention, which is inherently faster than multi-head attention. With multi-head attention or a longer training length, the increase would likely be larger, but based on rough estimates, anything under 50% should be acceptable.

Summary

This post proposed an "inverted" use of Leaky ReRoPE: by using a larger step size for Leaky ReRoPE during training, inference can revert to conventional RoPE, thereby keeping inference speed unchanged. Experimental results show that this approach still has a fair degree of competitiveness.