The Road to Upgrading Transformer: 11. Carrying β-ary Positions to the End

Transformer升级之路：11、将β进制位置进行到底 苏剑林 · 2023-07-31 · Machine Learning attention positional encoding generalization extrapolation RoPE

In the article Transformer Upgrade Path: 10, RoPE is a β-ary Encoding, we gave a base-$\beta$ interpretation of RoPE and, based on the idea of base conversion, derived NTK-aware Scaled RoPE, which can extend context length without fine-tuning. It has to be said that understanding positional encoding via the analogy with base-$\beta$ representation is a truly beautiful and inspiring perspective—so much so that every time I think it over again, I seem to gain some new insight from it.

In this article, we will revisit the base-$\beta$ interpretation of RoPE and try to generalize the existing NTK-aware Scaled RoPE, in the hope of finding a better strategy for extending the LLM context length without fine-tuning.

Base Analogy

We know that RoPE's parametrization follows the form of the Sinusoidal positional encoding. Whether by coincidence or by design, the Sinusoidal positional encoding for the integer $n$ has a lot in common with its base-$\beta$ representation.

Specifically, the $m$-th digit (counting from the right) of the base-$\beta$ representation of the integer $n$ is:

\begin{equation}\left\lfloor\frac{n}{\beta^{m-1}}\right\rfloor\bmod\beta\label{eq:mod}\end{equation}

while its Sinusoidal positional encoding is

\begin{equation}\boldsymbol{p}_n=\big[\cos\theta_1,\sin\theta_1,\cos\theta_2,\sin\theta_2,\cdots,\cos\theta_{d/2},\sin\theta_{d/2}\big]\\[5pt] \theta_m = \frac{n}{\beta^{m-1}},\quad \beta=10000^{2/d} \label{eq:sinu}\end{equation}

We can see that the two share the same $\frac{n}{\beta^{m-1}}$, and both $\bmod$ and $\cos,\sin$ are periodic functions, so the only difference between them is the inconsequential rounding $\lfloor\cdot\rfloor$. This is why it is such an intuitive and reasonable result to interpret RoPE/Sinusoidal positional encoding as a base-$\beta$ representation.

Fixing NTK

Following the reasoning of Transformer Upgrade Path: 10, RoPE is a β-ary Encoding, direct extrapolation concentrates the extrapolation burden on the "high-order digits" (where $m$ is large), whereas positional interpolation makes the representation of the "low-order digits" (where $m$ is small) denser, which is unfavorable for distinguishing relative distances. NTK-aware Scaled RoPE is essentially a base conversion: it spreads the extrapolation burden evenly across every digit while keeping the adjacent spacing unchanged. These properties are very friendly and important for LLMs, which clearly tend to rely on relative position, which is why it can achieve some effect even without fine-tuning.

Looking closely at equation $\eqref{eq:sinu}$, $\cos,\sin$ is in fact a single unit, so it actually only has $d/2$ digits—that is, it is equivalent to a $d/2$-digit base-$\beta$ encoding of $n$. If we want to extend to $k$ times the context length, converting from base $\beta$ to base $\beta\lambda$, then at the very least we should have

\begin{equation}\lambda^{d/2}=k\quad\Rightarrow\quad\lambda = k^{2/d}\end{equation}

so the new RoPE becomes

\begin{equation}\boldsymbol{p}_n=\big[\cos\theta_1,\sin\theta_1,\cos\theta_2,\sin\theta_2,\cdots,\cos\theta_{d/2},\sin\theta_{d/2}\big]\\[5pt] \theta_m = \frac{n}{(\beta\lambda)^{m-1}},\quad \beta=10000^{2/d},\quad \lambda = k^{2/d}\label{eq:ntk-old}\end{equation}

This is the NTK-RoPE that we proposed in the previous article.

However, after further thought, I realized that this is actually not quite reasonable. Going back to equation $\eqref{eq:mod}$, if we want to compute the $m$-th digit in base $\beta\lambda$, it should be

\begin{equation}\left\lfloor\frac{n}{(\beta\lambda)^{m-1}}\right\rfloor\bmod(\beta\lambda)\end{equation}

That is to say, besides changing $\frac{n}{\beta^{m-1}}$ to $\frac{n}{(\beta\lambda)^{m-1}}$, the period used to compute $\bmod$ should also be enlarged by a factor of $\lambda$, which is equivalent to dividing by an extra $\lambda$ before computing $\cos,\sin$:

\begin{equation}\boldsymbol{p}_n=\big[\cos\theta_1,\sin\theta_1,\cos\theta_2,\sin\theta_2,\cdots,\cos\theta_{d/2},\sin\theta_{d/2}\big]\\[5pt] \theta_m = \frac{n}{\lambda(\beta\lambda)^{m-1}},\quad \beta=10000^{2/d},\quad \lambda = k^{2/d}\label{eq:ntk-fixed}\end{equation}

In the experiments below, we call the version proposed in the previous article (equation $\eqref{eq:ntk-old}$) "NTK-RoPE-old", and the version in equation $\eqref{eq:ntk-fixed}$ "NTK-RoPE-fixed".

Mixed Base

Now, let's be a bit more adventurous—since we can use base $\beta$ to represent position, why not go ahead and use a more general "mixed base" scheme? By mixed base here, I mean that the carrying base used for each digit is not necessarily the same. This is actually not unfamiliar to us: for instance, 60 seconds make a minute, 60 minutes make an hour, but 24 hours make a day and 7 days make a week—here 60, 60, 24, 7 are different base values, in other words seconds, minutes, hours, days, weeks form an example that uses mixed bases.

Suppose, counting from the right, the 1st digit uses base $\beta_1$, the 2nd digit uses base $\beta_2$, the 3rd digit uses base $\beta_3$, and so on; then the $m$-th digit of $n$ is given by

\begin{equation}\left\lfloor\frac{n}{\beta_1\beta_2\cdots\beta_{m-1}}\right\rfloor\bmod\beta_m\label{eq:mod2}\end{equation}

Why consider a mixed base at all? Because one day I noticed an interesting fact: RoPE is essentially a form of relative positional encoding, and relative position is a special case of a Toeplitz matrix, which looks like this (since this article is mainly concerned with language models, the upper-right portion is omitted):

\begin{equation}\begin{pmatrix}0 & \\ 1 & 0 & \\ 2 & 1 & 0 &\\ 3 & 2 & 1 & 0 & \\ 4 & 3 & 2 & 1 & 0 & \\ 5 & 4 & 3 & 2 & 1 & 0 & \\ 6 & 5 & 4 & 3 & 2 & 1 & 0 & \\ \end{pmatrix}\end{equation}

From the above we can see that the distribution of relative positions is not uniform! 0 occurs the most often, 1 the next most, 2 less still, and so on—that is, the larger $n$ is, the fewer times it occurs. This means that, as a base-$\beta$ encoding, the "high-order digits" of RoPE are very likely undertrained—in other words, the generalization ability of the high-order digits is likely worse than that of the low-order digits. As mentioned earlier, NTK-RoPE spreads the extrapolation burden evenly across every digit; if this guess is correct, then "spreading it evenly" is not optimal—the low-order digits should bear a larger share and the high-order digits a smaller share, which naturally leads to a mixed base.

Optimizing the Allocation

Specifically, we extend the context length by a factor of $k$ by converting from base $\beta$ to a mixed base $\beta_1,\beta_2,\cdots,\beta_{d/2}$, where $\beta_m = \beta \lambda_m$. In this case, equation $\eqref{eq:mod2}$ becomes

\begin{equation}\left\lfloor\frac{n}{\beta^{m-1}(\lambda_1\lambda_2\cdots\lambda_{m-1})}\right\rfloor\bmod(\beta\lambda_m)\end{equation}

and equation $\eqref{eq:ntk-fixed}$ correspondingly becomes

\begin{equation}\boldsymbol{p}_n=\big[\cos\theta_1,\sin\theta_1,\cos\theta_2,\sin\theta_2,\cdots,\cos\theta_{d/2},\sin\theta_{d/2}\big]\\[5pt] \theta_m = \frac{n}{\beta^{m-1}(\lambda_1\lambda_2\cdots\lambda_m)},\quad \beta=10000^{2/d}\end{equation}

According to the principles of "extending by a factor of $k$" and "the low-order digits should bear a larger share", the constraint is

\begin{equation}\lambda_1\lambda_2\cdots\lambda_{d/2} = k,\quad \lambda_1 \geq \lambda_2 \geq \cdots \geq \lambda_{d/2} \geq 1\end{equation}

We consider a solution of the following form (interested readers are welcome to try other forms—there is quite a lot of freedom here):

\begin{equation}\lambda_1\lambda_2\cdots\lambda_m = \exp(am^b)\end{equation}

When $a > 0, b\leq 1$, it satisfies the condition $\lambda_1 \geq \lambda_2 \geq \cdots \geq \lambda_{d/2} \geq 1$; when $b=1$, this is actually just the earlier "NTK-RoPE-fixed"; and when $b=0$, it becomes "Positional Interpolation (PI)". $\lambda_1\lambda_2\cdots\lambda_{d/2} = k$ gives the constraint

\begin{equation}a\left(\frac{d}{2}\right)^b = \log k\end{equation}

so there is only one degree of freedom left to tune. After a simple binary search, I found that in my own experiments, $b=0.625$ tends to give, on average, a fairly good extension effect (different models may have different optima, so please tune it yourself). This version is called "NTK-RoPE-mixed".

Experimental Results

Building on the experiments in Transformer Upgrade Path: 10, RoPE is a β-ary Encoding, I ran additional experiments for "NTK-RoPE-fixed" and "NTK-RoPE-mixed", with the comparison shown 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\text{-fixed} & 49.40\% & 62.85\% & 44.14\% \\ \text{NTK-RoPE-mixed} & 49.41\% & 53.09\% & 40.12\% \\ \text{NTK-RoPE-}\log n\text{-mixed} & 49.40\% & \boldsymbol{68.91\%} & \boldsymbol{45.41\%} \\ \hline \end{array}$$

We can see that, compared with the equal-base "NTK-RoPE-old" and "NTK-RoPE-fixed", the improvement brought by the mixed-base "NTK-RoPE-mixed" is quite substantial—and since it requires no fine-tuning at all, it is essentially a "free lunch". In addition, we can see that the $\log n$ version indeed performs better at extension, but the $\log n$ trick needs to be incorporated at the pretraining stage. Some readers have previously asked whether models like LLAMA, which did not include the $\log n$ trick during pretraining, can still enjoy the "benefits" of $\log n$. After testing, I found that it can, by adding the following scale factor:

\begin{equation}\max(1, \log_{\text{maxlen}} n)\label{eq:plogn}\end{equation}

Here $\text{maxlen}$ is the maximum pretraining length—512 in the experiments of this article, 2048 for LLAMA, and 4096 for LLAMA2. In implementation, one can simply multiply each $\boldsymbol{q}_n$ by the corresponding factor. In this way, the part within $\text{maxlen}$ is unaffected, while the part beyond it is scaled by $\log n$, which serves as a simple transition. The results are as follows (with a $\color{red}{\dagger}$ added to distinguish it from the original $\log n$):

$$\begin{array}{c|cc} \hline \text{test length} & 512(\text{training}) & 4096(\text{repeat}) & 4096(\text{no repeat})\\ \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-mixed} & 49.41\% & 53.09\% & 40.12\% \\ \text{NTK-RoPE-}\log n^{\color{red}{\dagger}}\text{-mixed} & 49.41\% & 59.11\% & 42.38\% \\ \hline \end{array}$$

As we can see, this $\log n^{\color{red}{\dagger}}$ can also be considered a free lunch. In short, if you plan to pretrain from scratch, you might as well include the $\log n$ trick from the start; and if training is already complete, you can use equation $\eqref{eq:plogn}$ as a substitute, then further apply NTK-RoPE-mixed on top, to achieve a better context-extension effect.

Summary

In this article, we revisited the base-$\beta$ perspective on RoPE and attempted to generalize NTK-aware Scaled RoPE. Inspired by the idea of a mixed base, we obtained a better strategy for extending context length without fine-tuning, and finally demonstrated its effectiveness through experiments.