Transformer Upgrade Path: 10. RoPE as a β-ary Encoding

For readers who care about extending the Context length of LLMs, last week was undoubtedly an exciting one, with the open-source community producing one exhilarating result after another. First, the netizen @kaiokendev experimented with a "positional linear interpolation" scheme in his project SuperHOT, showing that with very little fine-tuning on long texts, an existing LLM could be made to handle Long Context. Almost simultaneously, Meta proposed the same idea, backed by extensive experimental results, in the paper Extending Context Window of Large Language Models via Positional Interpolation. But the surprises didn't stop there: shortly after, the netizen @bloc97 proposed NTK-aware Scaled RoPE, which achieved the effect of extending Context length without fine-tuning at all!

All these developments, and NTK-aware Scaled RoPE in particular, pushed me to rethink the meaning of RoPE. After some analysis, I found that RoPE's construction can be seen as a kind of $\beta$-ary encoding, and from this perspective, all these advances from the open-source community can be understood as different ways of extending this base-encoding scheme. More below.

Base Representations

Suppose we have an integer $n$ less than 1000 (not including 1000) that we want to feed into a model as a conditioning input. What would be a good way to do this?

The most naive idea is to feed it directly as a one-dimensional floating-point vector. However, the range 0–999 spans nearly a thousand units, which is not easy for a gradient-based optimizer to handle well. What if we scale it into the range 0–1 instead? That's not great either, because then the gap between adjacent integers shrinks from 1 to 0.001, making it hard for both the model and the optimizer to distinguish neighboring numbers. In general, gradient-based optimizers are a bit "finicky" — they only handle inputs of moderate scale well, and both too large and too small tend to cause problems.

So, to avoid this issue, we need to come up with a new way of encoding the input. Before figuring out how to make a machine handle this, let's think about how humans handle it. For an integer like 759, this is a three-digit number in base 10, with each digit ranging from 0 to 9. Since we ourselves represent numbers in base 10, why not just feed this base-10 representation directly into the model? That is, we represent the integer $n$ as a three-dimensional vector $[a,b,c]$, where $a,b,c$ are respectively the hundreds, tens, and units digits of $n$. This way, we shrink the numerical range without shrinking the gap between adjacent numbers, at the cost of increasing the input dimensionality — which, conveniently, is exactly what neural networks are good at handling.

If we want to shrink the numerical range even further, we can also reduce the base further, e.g. using base 8, base 6, or even base 2, at the cost of further increasing the input dimensionality.

Direct Extrapolation

Suppose we trained a model using this three-dimensional base-10 representation, and it works reasonably well. Then suddenly a new requirement comes in: increase the upper bound of $n$ to under 2000. How should we handle this?

If we still use a base-10 representation as input to the model, the input now becomes a four-dimensional vector. But the original model was designed and trained for three-dimensional vectors, so adding a new dimension leaves the model unable to cope. Some readers might ask: why not just reserve enough extra dimensions in advance? Indeed, we could reserve a few extra dimensions, set them to 0 during training, and then change them to other numbers at inference time — this is extrapolation.

Direct ExtrapolationDirect Extrapolation

However, if the reserved dimensions are always 0 during training, changing them to other numbers at inference time is unlikely to work well, since the model has no guarantee of being able to adapt to situations it has never been trained on. In other words, because the training data for certain dimensions is insufficient, direct extrapolation typically causes a serious drop in model performance.

Linear Interpolation

This led people to think of replacing extrapolation with interpolation: simply put, compress the range under 2000 down to under 1000 — for example, by dividing by 2, so 1749 becomes 874.5, which is then converted into the three-dimensional vector [8, 7, 4.5] and fed into the original model. In absolute terms, the new $[7,4,9]$ actually corresponds to 1498, twice the original mapping — inconsistent with the original mapping. In relative terms, the gap between adjacent numbers used to be 1, but is now 0.5, making the last dimension more "crowded." So, after making this interpolation modification, some fine-tuning is usually needed so the model can re-adapt to this crowded mapping.

Linear InterpolationLinear Interpolation

Of course, some readers will say the extrapolation scheme can also be fine-tuned. True, but the interpolation scheme requires far fewer fine-tuning steps, because in many scenarios (such as positional encoding), relative magnitude (or perhaps we should say ordering information) matters more — in other words, the model only needs to know that 874.5 is greater than 874, not what actual quantity it represents. Since the model has already learned that 875 is greater than 874, and models have some degree of generalization ability, learning the additional fact that 874.5 is greater than 874 isn't too hard.

That said, the interpolation scheme isn't perfect either. As the range being handled grows further, the adjacent gaps become even smaller, and this shrinking gap is concentrated in the units digit, while the hundreds and tens digits still retain an adjacent gap of 1. In other words, interpolation causes the distribution across different dimensions to become uneven — each dimension is no longer treated equally, which makes further learning harder for the model.

Base Conversion

Is there a scheme that doesn't add new dimensions yet still preserves the adjacent gap? Yes — and it's probably something we're already familiar with: base conversion! Three digits in base 10 can represent 0–999. What about base 16? At most it can represent $16^3 - 1 = 4095 > 1999$. So, we only need to convert to base 16 — e.g. 1749 becomes $[6,13,5]$ — and a three-dimensional vector can cover the target range, at the cost that each dimension's digits now range from 0–9 to 0–15.

Base ConversionBase Conversion

If you think about it carefully, this turns out to be a rather brilliant idea. As mentioned above, the scenarios we care about mainly rely on ordering information, and the previously trained model has already learned $875 > 874$. Under base 16, we likewise have $875 > 874$, with exactly the same comparison rule (the model has no idea what base you're using in the first place). The only concern is whether the model can still compare correctly once each digit exceeds 9 (i.e. 10–15), but in practice models generally have some generalization ability, so a slight extrapolation in each dimension isn't a problem. So this idea of converting the base might even work without fine-tuning the original model at all! Additionally, to further narrow the extrapolation range, we could switch to a smaller base $\left\lceil\sqrt[3]{2000}\right\rceil =13$ instead of 16.

As we'll see next, this idea of base conversion actually corresponds exactly to the NTK-aware Scaled RoPE mentioned at the start of the article!

Positional Encoding

To establish the connection, we first need the following result:

The rotary position embedding (RoPE) for position $n$ is, in essence, exactly the $\beta$-ary encoding of the number $n$!

This might seem surprising at first, since the two appear completely different on the surface. But in fact, the two operations share the same key property. To understand this, let's first recall that for a base-10 number $n$, if we want to find the $m$-th digit (counting from the right) of its $\beta$-ary representation, the method is

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

that is, first divide by the $\beta^{k-1}$-th power, then take the modulus (remainder). Now let's recall RoPE, whose construction is based on Sinusoidal positional encoding, which can be rewritten as

\begin{equation}\left[\cos\left(\frac{n}{\beta^0}\right),\sin\left(\frac{n}{\beta^0}\right),\cos\left(\frac{n}{\beta^1}\right),\sin\left(\frac{n}{\beta^1}\right),\cdots,\cos\left(\frac{n}{\beta^{d/2-1}}\right),\sin\left(\frac{n}{\beta^{d/2-1}}\right)\right]\label{eq:sinu}\end{equation}

where $\beta=10000^{2/d}$. Now, comparing with equation $\eqref{eq:mod}$ — doesn't equation $\eqref{eq:sinu}$ have exactly the same $\frac{n}{\beta^{m-1}}$? As for the modulus operation, its most important property is periodicity — and isn't the $\cos,\sin$ in equation $\eqref{eq:sinu}$ also precisely a periodic function? So, apart from the inconsequential difference of the floor function, RoPE (or Sinusoidal positional encoding) is really just the $\beta$-ary encoding of the number $n$!

Having established this connection, the schemes for extending the integer $n$ discussed in the previous sections can now be mapped onto the various developments mentioned at the beginning of the article. Among these, the direct extrapolation scheme is simply leaving things unchanged, while the interpolation scheme replaces $n$ with $n/k$, where $k$ is the scaling factor by which we want to extend the range — this is exactly the Positional Interpolation experimented with in Meta's paper, and the experimental results there indeed confirm that extrapolation requires more fine-tuning steps than interpolation.

As for base conversion, this amounts to extending the representable range by a factor of $k$, which requires the original base $\beta$ to be extended to at least base $\beta (k^{2/d})$ (note that although equation $\eqref{eq:sinu}$ is a $d$-dimensional vector, the $\cos,\sin$ appear in pairs, so it's effectively a $\beta$-digit representation in base $d/2$, hence we need to take the $d/2$-th root rather than the $d$-th root), or equivalently, replace the original base $10000$ with $10000k$ — this is essentially NTK-aware Scaled RoPE. As discussed earlier, since positional encoding depends more on ordering information, and base conversion basically doesn't change the ordering comparison rule, NTK-aware Scaled RoPE achieves decent results on longer contexts even without fine-tuning.

Tracing the Origins

Some readers might be curious what this has to do with NTK. NTK stands for "Neural Tangent Kernel," which we touched on briefly in Viewing Optimization Algorithms Through Dynamics (7): SGD ≈ SVM?. The connection between the above result and NTK is more a matter of the proposer's academic background — the proposer was quite familiar with results such as Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains, which uses NTK-related results to prove that neural networks cannot directly learn high-frequency signals, and that the solution is to convert them into Fourier features — a form quite similar to the Sinusoidal positional encoding in equation $\eqref{eq:mod}$.

So, based on intuition from NTK-related results, the proposer derived NTK-aware Scaled RoPE. I asked the proposer about the derivation, and it turns out to be quite simple: it combines extrapolation and interpolation — extrapolating high frequencies and interpolating low frequencies. Specifically, the lowest frequency term in equation $\eqref{eq:sinu}$ is $\frac{n}{\beta^{d/2-1}}$; introducing a parameter $\lambda$ turns it into $\frac{n}{(\beta\lambda)^{d/2-1}}$, and setting it equal to the interpolation case gives

\begin{equation}\frac{n}{(\beta\lambda)^{d/2-1}} = \frac{n/k}{\beta^{d/2-1}}\end{equation}

Solving this yields $\lambda=k^{2/(d-2)}$. As for the highest frequency term, $\frac{n}{\beta}$, after introducing $\lambda$ it becomes $\frac{n}{\beta\lambda}$; since $d$ is usually large, $\lambda$ is very close to 1, so this term remains close to $\frac{n}{\beta}$, i.e. it's equivalent to extrapolation.

So this scheme cleverly and simply combines extrapolation and interpolation. Moreover, since $d$ is typically fairly large (64 for BERT, 128 for LLAMA), $k^{2/(d-2)}$ is not very different from $k^{2/d}$, so it is basically consistent with the $k^{2/d}$ solution I proposed based on the base-conversion idea. Also, from the proposer's underlying idea, any scheme that achieves "high-frequency extrapolation, low-frequency interpolation" would work — it need not be limited to the specific scheme introducing $\lambda$ described above. Readers are welcome to try this themselves.

My Own Tests

As a scheme purported to extend an LLM's Context length without any fine-tuning, I was quite shocked the first time I saw NTK-aware Scaled RoPE, and couldn't wait to test it. After all, based on my experience in Transformer Upgrade Path: 9. A New Approach to Global Length Extrapolation, on my preferred "GAU + Post-Norm" combination, many mainstream schemes had failed — so how would this one fare?

With $k$ set to 8, the comparison results are as follows (regarding the distinction between "repeated" and "non-repeated," see here):

$$\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} & 49.41\% & 51.28\% & 39.27\% \\ \text{NTK-RoPE-}\log n & 49.40\% & 61.71\% & 43.75\% \\ \hline \end{array}$$

All the results reported above are without any fine-tuning on long texts, where Baseline refers to extrapolation, PI (Positional Interpolation) refers to switching Baseline to interpolation, and NTK-RoPE refers to switching Baseline to NTK-aware Scaled RoPE. The option marked with $\log n$ refers to adding the scale from Viewing Attention's Scale Operation Through Entropy Invariance during pretraining; this variant is considered because I felt that while NTK-RoPE solves RoPE's length generalization problem, it doesn't solve the issue of attention dispersion.

The experimental results in the table fully match expectations:

1. Direct extrapolation doesn't work well;
2. Interpolation without fine-tuning also performs poorly;
3. NTK-RoPE achieves non-trivial (though somewhat degraded) extrapolation results without fine-tuning;
4. Adding $\log n$ to concentrate attention does indeed help.

So, NTK-RoPE has successfully become the second scheme I've personally found effective for extending an LLM's Context length without fine-tuning (the first, of course, being NBCE) — hats off once again to the proposer's remarkable insight! Even more encouraging, NTK-RoPE performs noticeably better on "repeated" extrapolation than "non-repeated" extrapolation, showing that this modification preserves global dependencies rather than simply localizing attention.

Closing Remarks

This article approached RoPE from the perspective of $\beta$-ary encoding, and used this lens to introduce some recent developments in the open-source community around Long Context, including a modification scheme that can extend Context length without any fine-tuning.

In just one week, the open-source community's progress on Long Context has been almost overwhelming to keep up with — and thoroughly gratifying. So much so that the netizen @ironborn123 commented:

Last week looked like revenge of the interpolators :) OpenClosedAI had better watch out.
English translation of a post from 科学空间 | Scientific Spaces by 苏剑林. Original: https://kexue.fm/archives/9675
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.