Transformer Upgrade Path: 16. "Revisiting" Length Extrapolation Techniques
Looking back, I realize that starting from Part 7, Transformer Upgrade Path: 7. Length Extrapolation and Local Attention, this "Transformer Upgrade Path" series has been "stuck" on length extrapolation, with nine consecutive articles (not counting this one) revolving around it. It's now just a little over a year since that seventh article, and during this year, research on length extrapolation in the open-source community has made significant progress. I've also gradually developed some of my own understanding — for instance, that this problem is far from as simple as it initially seemed, that many of the earlier local-attention-based approaches aren't always effective, and that this hints that a lot of the older analytical work never really got at the heart of the problem.
In this article, I want to try to "revisit" the mainstream results on length extrapolation, combining my own findings and understanding, in an attempt to uncover the key to training-free length extrapolation.
Problem Definition
As the name suggests, training-free length extrapolation means obtaining a model capable of handling and predicting long sequences without any additional training on long-sequence data — training only on short-sequence corpora — i.e., "Train Short, Test Long." So how do we judge whether a model can be used for long sequences? The most basic metric is whether the model's long-sequence loss or PPL blows up; a more practically-oriented evaluation feeds the model a sufficiently long context and has it predict an answer, which is then compared against the ground truth using metrics like BLEU or ROUGE. LongBench is one such benchmark. more
But note that length extrapolation should not come at the cost of long-range dependency — otherwise there'd be no point considering length extrapolation at all, and we might as well just truncate the text. This means that schemes which explicitly truncate long-range dependencies need to be chosen carefully — for instance ALIBI, and most of the schemes listed in Transformer Upgrade Path: 7. Length Extrapolation and Local Attention, as well as linear RNNs with explicit decay. All of these schemes behave like local attention once the sequence length is large enough, so even if they achieve length extrapolation, there's a risk of insufficient long-range dependency, and one needs to weigh this carefully depending on the use case.
How do we tell whether long-range dependency has been preserved alongside length extrapolation? A relatively rigorous approach is the evaluation scheme proposed at the end of Transformer Upgrade Path: 12. Infinitely Extrapolatable ReRoPE?: prepare a sufficiently long piece of text, but for each model, only compute the metric on the last segment of each sample, as shown below:
An evaluation method focused on long-range dependency
For example, suppose the model's training length is 4K, and we want to see how well it extrapolates to 16K. Then we prepare a 16K-token test set: the 4K model takes the last 4K tokens of each sample as input and computes the metric on them; the 8K model takes the last 8K tokens as input but computes the metric only on the last 4K tokens; the 12K model takes the last 12K tokens as input but computes the metric only on the last 4K tokens; and so on. This way, models of different lengths are all being evaluated on the metric for the same span of tokens, the only difference being the length of context fed in. If long-range dependency is effectively preserved, then the longer the context, the better the metric should get.
Rotary Position
Having finished discussing evaluation, let's return to methods. At the start of this article we mentioned "older analytical work" — one key distinction between "old" and "new" here is that most of the "old" work tried to design new architectures or positional encodings from scratch to achieve length extrapolation, whereas the "new" work of the past year has mainly studied length extrapolation for Decoder-only Transformer models equipped with rotary position embedding (RoPE).
As a side note — why have most LLMs today settled on RoPE as their positional encoding? I think there are a few main reasons:
1. RoPE has no explicit long-range decay, which is crucial for models aiming at long context;
2. RoPE is a genuine positional encoding: through trigonometric functions of different frequencies, it effectively distinguishes long-range from short-range positions, achieving an effect similar to hierarchical positional encoding — again a key ingredient for long context;
3. RoPE acts directly on Q and K without changing the form of attention itself, which meshes well with Flash Attention and makes it easier to scale up.
By contrast, things like ALIBI and KERPLE, although sometimes also called positional encodings, are in fact just a form of attention bias, carrying little positional information and not applicable to encoders. The reason they work for decoders is largely that the decoder's own lower-triangular mask already provides a fairly substantial positional bias, and the extra attention bias is just icing on the cake. Moreover, they cannot effectively distinguish long-range from short-range within a single head — they have to rely on setting different decay factors across different heads — which also means they tend to perform poorly with single-head attention (e.g., GAU).
Having gone on at such length comparing pros and cons, this might look like tooting my own horn, but that's not really the intent — it's meant as an exchange of views with readers, since some of you have raised the same questions before. As the original proposer of RoPE, my own understanding of it isn't necessarily any deeper than anyone else's — after all, the original motivation for proposing RoPE was purely for fun, and at the time I'd have been happy just to get something that worked, let alone something that could rival learnable absolute positional encodings. So, since this was all rather "unexpected" to begin with, it's only natural that "the author himself doesn't have that thorough an understanding" either.
Windowed Truncation
I seem to have drifted off topic again. To put it simply, the point of the last two sections was this: RoPE currently appears to be sufficient for long context, so studying its length extrapolation is worthwhile, and when choosing a length-extrapolation scheme, we should not sacrifice long-range dependency.
In the earliest article on this site discussing length extrapolation, Transformer Upgrade Path: 7. Length Extrapolation and Local Attention, we judged length extrapolation to be an out-of-distribution (OOD) problem at inference time. Although some of the commentary in that article looks a bit dated from today's vantage point, this basic judgment still holds up fairly well. Cast in terms of RoPE, it means that at inference time we encounter relative distances that were never seen during training. A seemingly viable fix, then, is to introduce a sliding-window attention mask, as shown on the left below:
Sliding Window Mask
Λ-shape Window Mask
Of course, since this forcibly truncates attention outside the window, this scheme doesn't satisfy the principle of "not sacrificing long-range dependency" — but let's treat it purely as a baseline for now. Unfortunately, even with such a sacrifice made, this scheme still doesn't work — it can't even keep PPL from blowing up! In-depth analysis of this phenomenon gave rise to two papers, LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models and Efficient Streaming Language Models with Attention Sinks, which arrived at almost identical conclusions. But interestingly, a few months earlier, an "outsider" had already discovered the same result and published it in a Zhihu column article, Perpetual Sampling Technical Report.
The answer might be surprising: the first few tokens matter a lot and cannot be discarded. So the usable window mask ends up looking like the right-hand panel above (the LM-Infinite paper calls this the "$\Lambda$-mask").
Why do the leading tokens occupy such an important position? There are currently two different ways of understanding this:
1. The leading tokens serve as an "anchor" for absolute position: as the name suggests, relative positional encoding can in principle only identify relative position, but some tasks may actually depend quite heavily on absolute position. By using the leading few tokens — whose absolute positions are approximately 0 — as a "reference," every token can work out its own absolute position. Remove the leading tokens, and this mechanism is lost, completely disrupting the attention pattern and causing PPL to blow up;
2. The leading tokens serve as a "recycle bin" for attention: since attention weights must sum to 1, attention has to be assigned to some tokens, but in many cases the model may find that "there's nothing here worth attending to." In that case it chooses to dump a chunk of attention on the leading, low-information tokens, effectively implementing "not attending." Remove them, and the model is forced to redistribute that attention onto other irrelevant tokens, disrupting the attention pattern.
To put it plainly: empirically, in most cases the attention share on the first few tokens is quite large, so they can't be removed — removing them throws everything into disarray. As for why it's so large, that's left to everyone's imagination.
Position Interpolation
Windowed truncation is certainly a decent baseline for length extrapolation, and the "anchor" and "recycle bin" interpretations have also deepened our understanding of how the attention mechanism works. But as noted, this comes at the cost of forcibly truncating attention outside the window and sacrificing long-range dependency, so it's not the final solution.
The OOD nature of relative position manifests directly as relative positions at inference time exceeding the range seen during training; since these were never trained on, the behavior of the "out-of-bounds" part is unpredictable. To address this, a blogger going by "kaiokendev" proposed an extremely simple fix in his blog post https://kaiokendev.github.io/til#extending-context-to-8k — "position interpolation" — scaling the positional encodings of the long text being predicted by a factor of $\frac{L_{train}}{L_{test}}$, compressing them back within the training length range, as in the formula below (all positions in the formula are relative positions). Shortly after, Meta published the same idea in the paper Extending Context Window of Large Language Models via Positional Interpolation, naming it "Positional Interpolation (PI)" and supplying more extensive experimental results.
\begin{equation}\begin{aligned}&\text{training stage}:\,(1,2,\cdots,n-1,n)\\[5pt] &\text{prediction phase}:\,(1,2,\cdots,n,\underbrace{n+1,\cdots,4n-1,4n}_{\text{out of bounds afar}})\xrightarrow{\quad\text{interpolation}\quad} \big(\underbrace{\frac{1}{4},\frac{2}{4},\frac{3}{4}}_{\text{local distortion}},\cdots,n-\frac{1}{4},n\big)\end{aligned}\end{equation}
However, position interpolation isn't really a length-extrapolation scheme, or at least not a training-free one, since PPL still blows up after applying it. The reason isn't hard to understand: although position interpolation avoids the problem of distant positions going out of range, it simultaneously compresses the distances between nearby tokens, severely disrupting the model's local resolution — and as we all know, language modeling is a task heavily reliant on local relationships, so disrupting locality naturally ruins prediction accuracy.
That said, this doesn't mean position interpolation is worthless. Readers who need length extrapolation generally fall into two camps: one group lacks the resources to do long-text fine-tuning and hopes to go straight from a short-text model to a usable long-text model — this use case demands fairly high extrapolation quality, and position interpolation isn't suited to it. The other group does have the resources for long-text fine-tuning, and studies length extrapolation purely to get a better initialization for that fine-tuning — this case can tolerate a higher initial loss from model modification, as long as fine-tuning can quickly recover it, and position interpolation fits neatly into this category. Meta's paper shows that after applying PI, only about 1000 steps of long-text training are needed to obtain an effective long-text model, which is far more training-efficient than fine-tuning directly without any modification.
Preserve the Near, Compress the Far
Direct extrapolation runs into the problem of distant positions going out of range, while position interpolation runs into the problem of local distortion. These look complementary — can we combine the strengths of both? This is exactly what Leaky ReRoPE, proposed in Transformer Upgrade Path: 12. Infinitely Extrapolatable ReRoPE?, does, along with its limiting case, ReRoPE.
Building on the analysis in the previous section, we can guess that the key to training-free length extrapolation is "preserve the near, compress the far" — that is, "keep local positions undistorted" and "compress distant positions so they don't go out of range." Leaky ReRoPE achieves this via a very direct idea: it sets a window size $w$, splits relative positions into two parts, leaves relative positions unchanged within the window to "keep local positions undistorted," and applies position interpolation outside the window to "compress distant positions so they don't go out of range," as in the formula below:
\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}\end{equation}
If we push the interpolation factor $k$ to infinity, we get the maximally simplified ReRoPE, where the positional encoding outside the window becomes constant at $w$, meaning that sequences of any length whatsoever will never go out of range — in theory, it has the potential for infinite extrapolation! And indeed, both Leaky ReRoPE and ReRoPE perform very well in practice: judging by loss, they achieve almost no loss of performance within the training length, achieve length extrapolation, and moreover, the longer the context, the lower the loss — showing that they genuinely preserve long-range dependency while extrapolating.
The main drawback of Leaky ReRoPE and ReRoPE is that their code is somewhat troublesome to implement. Unlike attention-bias-style positional encodings, RoPE can't be implemented by first constructing a relative position matrix and then computing the relative positional encoding from it (that would be far too inefficient) — it can only implement relative positional encoding via absolute positional encoding, which means it can only realize linearly growing relative positions. But the relative positions in Leaky ReRoPE and ReRoPE are piecewise linear, which means that a naive implementation would require computing the attention matrix twice (to get the two different linear segments) and then splicing them together — clearly reducing efficiency.
The good news, however, is that mainstream attention acceleration techniques such as Flash Attention already compute attention block by block — say, in blocks of 128 — so when the sequence is sufficiently long, the proportion of piecewise-linear blocks becomes very small (occurring only near the window boundary), as shown in the figure below: only the blocks with mixed red-and-green coloring need attention computed twice, while all the remaining single-colored blocks only need to be computed once. So combined with blockwise attention computation, the extra computational cost incurred by Leaky ReRoPE and ReRoPE is essentially negligible. Reader @chu-tianxiang previously shared a Triton-based implementation in the comments section, which interested readers can check out.
\begin{equation}\left(\begin{array}{cccc:cccc:cccc} \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} & \\ \hdashline \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} & \\ \hdashline \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{array}\right)\end{equation}
Coincidentally, earlier this month a paper was submitted to Arxiv, LLM Maybe LongLM: Self-Extend LLM Context Window Without Tuning, proposing a training-free length extrapolation method called "Self-Extend." It's essentially Leaky ReRoPE with a rounding operation added on top, snapping every relative position back to an integer to further alleviate the OOD issue with relative positions. The results reported in the paper are also very good, which further confirms the effectiveness of Leaky ReRoPE.
The circle-turning perspective
Although Leaky ReRoPE and ReRoPE work quite well in practice (at least in terms of loss), like position interpolation they operate directly on the position indices (Position Ids), which feels a bit like "treating the head when the head hurts, treating the foot when the foot hurts" — lacking a deeper analysis of the underlying mechanism. Because for the model, position indices themselves are not what matters; it is the position embeddings that interact directly with the model. So if we want to get "right to the root of the disease," we should try approaching from the position embeddings instead.
Some readers might ask: aren't position indices and position embeddings in one-to-one correspondence? Isn't operating on position indices equivalent to operating on position embeddings? That's true in principle, but the two behave differently in practice. For instance, position indices are unbounded, whereas position embeddings are bounded (RoPE is built from trigonometric functions, which are bounded). It's the position embeddings that interact directly with the model, so even when the position indices go out of distribution, the position embeddings need not. Analyzing from the position-embedding perspective therefore gives us a clearer picture of what "OOD due to length extrapolation" actually looks like, so that we can treat the real cause rather than the symptom.
In Transformer Upgrade Path: 2. Rotary Position Embedding That Draws on the Best of All Worlds, when we derived RoPE we first solved the two-dimensional case using complex numbers, then concatenated multiple two-dimensional solutions into a higher-dimensional solution. As a result, the $\boldsymbol{q},\boldsymbol{k}$ inner product after adding RoPE can be written in complex form as
\begin{equation} (\boldsymbol{\mathcal{R}}_m \boldsymbol{q})^{\top}(\boldsymbol{\mathcal{R}}_n \boldsymbol{k}) = \text{Re}\left[\sum_{i=0}^{d/2-1}\boldsymbol{q}_{[2i:2i+1]}\boldsymbol{k}_{[2i:2i+1]}^* e^{\text{i}(m-n)\theta_i}\right]\end{equation}
where $\theta_i$ defaults to $10000^{-2i/d}$, which is a function that gradually decreases from 1 to nearly 0. From Euler's formula $e^{\text{i}t}=\cos t + \text{i}\sin t$ we know that $e^{\text{i}(m-n)\theta_i}$ is essentially a point on the unit circle; as $m-n$ gradually grows, this point rotates around the unit circle (literally rotating), and the larger $\theta_i$ is, the faster it rotates, and vice versa.
More than one full turn
Less than one full turn
Suppose the training length is $L_{train}$, so that $m-n\in[0, L_{train}-1]$. Now let's let our imagination run a bit: a larger $\theta_i$ means a faster rotation speed and a shorter period, so as $m-n$ goes from $0$ to $L_{train}-1$, it has already turned many full circles — meaning almost every point on the circle has been "seen" during training. Hence these $\theta_i$ hardly suffer from any OOD problem. Conversely, for a smaller $\theta_i$, as $m-n$ goes from $0$ to $L_{train}-1$ it may not even complete one full turn, in which case the points that were actually trained form at most an arc on the circle. If at test time we encounter a larger $L_{test}$, it falls outside that trained arc, leading to unpredictable behavior — and this is exactly when we need interpolation to squeeze it back into the original arc. In other words, whether the position index $m-n$ itself is OOD doesn't really matter; what matters is whether the corresponding point on the unit circle has been adequately trained. If it has, we can leave it untouched (i.e., extrapolate directly); if not, we need to find a way to compress it back onto the arc that has already been well trained (i.e., position interpolation).
Concretely, for $\theta_i$ we can compute its period as $T_i=2\pi/\theta_i$, and from this compute how many "turns" it makes during training, namely $r_i=\frac{L_{train}}{T_i}=\frac{\theta_i L_{train}}{2\pi}$. We can set a threshold number of turns $\tau$: any component whose turn count exceeds $\tau$ is considered to have been adequately trained and can be left unchanged; any component whose turn count is below 1 has $\theta_i$ changed to $\frac{\theta_i L_{train}}{L_{test}}$, meaning that whatever falls outside the trained arc gets rescaled back into it; and for everything in between, we linearly interpolate. In formula form:
\begin{equation}\theta_i^{new} = \left[\gamma_i + (1 - \gamma_i)\frac{L_{train}}{L_{test}}\right]\theta_i,\quad \gamma_i = \left\{\begin{aligned}&1,&r_i > \tau \\ &0,&r_i < 1 \\ &\frac{r_i - 1}{\tau - 1},&\text{others} \end{aligned}\right.\end{equation}
This is exactly the training-free length extrapolation scheme called "YaRN," proposed in YaRN: Efficient Context Window Extension of Large Language Models. In my own tests, its extrapolation performance is excellent, only slightly behind Leaky ReRoPE and ReRoPE. But it's worth noting that YaRN only changes the value of $\theta_i$, without altering the form of Attention or RoPE itself, so it incurs no extra implementation or inference cost. Given this constraint (i.e., that it must be a drop-in replacement within existing implementations), YaRN is the best-performing length extrapolation method I have tested.
A side note
Actually, the YaRN story isn't quite over yet, but the previous section already felt long enough, so it's better to give this its own section. Besides the change to $\theta_i$, YaRN also multiplies the Attention logits by an additional scale factor:
\begin{equation}\lambda = \left(1 + 0.1 \log \frac{L_{test}}{L_{train}}\right)^2\label{eq:scale-yarn}\approx 1 + 0.2 \log \frac{L_{test}}{L_{train}}\end{equation}
The derivation of this scale factor is, frankly, a bit amusing — the answer is that there is no derivation at all. The authors say they were not able to derive it theoretically either; it was found purely by experiment that PPL was lower after adding this scale, and the formula above was fitted empirically from experiments.
Actually, this logarithmic result looks quite similar to the $\log n$ scale factor derived in Attention Scale Operations Viewed Through Entropy Invariance, except the latter depends on the specific position, while the former, once $L_{test}$ is fixed, is just a constant. Considering that when $n$ is relatively large, the $\log n$ function changes rather slowly, treating it as a constant over a certain range is not unreasonable. So it's not hard to guess that YaRN's scale factor and the entropy-invariance-derived $\log n$ scale share a common origin. I've also done a comparison: replacing the constant $\lambda$ with a factor depending on the absolute position $n$ as follows achieves a similar effect:
\begin{equation}\lambda_n = \max\left(1, \frac{\log n}{\log L_{train}}\right)\label{eq:clip-logn}\end{equation}
Note that
\begin{equation}\frac{\log L_{test} }{\log L_{train}} = 1 + \frac{1}{\log L_{train}} \log\left(\frac{L_{test}}{L_{train}}\right)\end{equation}
YaRN's experiments were based on LLAMA and LLAMA2, with training lengths of 2K and 4K respectively, giving us $\frac{1}{\log 2048}\approx 0.13$, $\frac{1}{\log 4096}\approx 0.12$ — a coefficient roughly half that of formula $\eqref{eq:scale-yarn}$. The difference isn't large, and in fact the precise value of this coefficient probably isn't too important, since I have also found datasets where formula $\eqref{eq:clip-logn}$ works better. So in this way we can consider formula $\eqref{eq:scale-yarn}$ to have been approximately derived.
Compared to YaRN itself, the story of YaRN's author Bowen Peng is perhaps even more "captivating." His earlier proposal, NTK-RoPE, was the very first training-free length extrapolation scheme for RoPE, and two posts in this series, Transformer Upgrade Path: 10. RoPE Is a Base-β Encoding and Transformer Upgrade Path: 11. Pushing Base-β Positional Encoding to the End, were both directly inspired by it. While NTK-RoPE's performance may not look particularly great by today's standards (compared to YaRN, ReRoPE, etc.), it was the first to demonstrate that training-free length extrapolation was even possible — a milestone achievement. One could even say that all subsequent research on length extrapolation has, directly or indirectly, benefited from the imaginative space that NTK-RoPE opened up.
The idea behind NTK-RoPE is very simple: just change RoPE's base. That is, what was originally $\theta_i = 10000^{-2i/d}$ becomes $\theta_i = (10000\kappa)^{-2i/d}$. How is $\kappa$ chosen? At the time, based on his own experience with results related to the Neural Tangent Kernel (NTK), Bowen Peng judged that high frequencies ($i\to 0$) are responsible for learning relative distance and thus shouldn't be changed, while low frequencies ($i\to d/2-1$) are responsible for learning absolute distance and thus need interpolation. In summary: "extrapolate at high frequencies, interpolate at low frequencies." He then set the scale at $i = d/2-1$ equal to the interpolation scale $\frac{L_{train}}{L_{test}}$, giving the equation
\begin{equation}(10000\kappa)^{-2i/d}|_{i=d/2-1} = \left.\frac{L_{train}}{L_{test}}10000^{-2i/d}\right|_{i=d/2-1}\end{equation}
which solves to
\begin{equation}\kappa = \left(\frac{L_{test}}{L_{train}}\right)^{d/(d-2)}\label{eq:kappa}\end{equation}
This simple yet clever derivation opened the "Pandora's box" of training-free length extrapolation.
From YaRN's point of view, it's not only the $i = d/2-1$ term whose $\theta_i$ fails to complete a full turn — so having NTK-RoPE apply full interpolation only to the very last $i = d/2-1$ is insufficient. And indeed, that's exactly what happens in practice: setting $\kappa$ as in formula $\eqref{eq:kappa}$ only lets the model extrapolate to around $L_{test}/2$ in length without PPL exploding; beyond that, PPL rises noticeably. It was precisely because of this problem that the author went on to propose the later, upgraded scheme, YaRN.
That said, even though NTK-RoPE underperforms YaRN, for the second group of readers mentioned earlier — those with the resources to do long-text fine-tuning — NTK-RoPE might actually be more appealing. Since they're fine-tuning anyway just to get a better initialization, they don't much care about the difference in initial performance between NTK-RoPE and YaRN; they'd rather go with the simpler-to-implement NTK-RoPE. For example, CodeLLAMA was built on top of LLAMA2 by changing the base to $10^6$ and continuing training from there. In addition, Meta, in their paper Effective Long-Context Scaling of Foundation Models, renamed NTK-RoPE as RoPE-ABF (Adjusted Base Frequency) — compared to the somewhat mysterious "NTK," the name ABF more directly conveys what it actually means.
Refusing to pay the tax
I wonder if you've noticed that none of the training-free length extrapolation methods mentioned above manage to keep the model's performance unchanged within the training length $L_{train}$. Specifically, let the original model be $f(x)$, and let the model after the extrapolation modification be $f^+(x)$. When the length of $x$ does not exceed $L_{train}$, there's no guarantee that $f(x)\equiv f^+(x)$. Since $f(x)$ was trained precisely within $L_{train}$, it's reasonable to assume that $f(x)$ is optimal for samples of length at most $L_{train}$. So $f^+(x)\neq f(x)$ means that while length extrapolation improves performance on longer samples, it degrades performance within the original range $L_{train}$. We might vividly call this loss the "extrapolation tax."
Back when NTK-RoPE was first proposed, the open-source community already recognized this "extrapolation tax" problem, and proposed a corresponding solution: dynamically adjusting the scale factor of each extrapolation method as the training length changes — this is Dynamic Scaling, first proposed in the Reddit post Dynamically Scaled RoPE further increases performance of long context LLaMA with zero fine-tuning. Taking YaRN as an example, the length-dependent scaling factor in it is $s=\frac{L_{test}}{L_{train}}$; Dynamic Scaling replaces this with a dynamic $s(pos)=\frac{\max(L_{train}, pos+1)}{L_{train}}$, where $pos$ is the position index of the current token (counting from zero). This change means Dynamic Scaling tries to find, for every position, the smallest scale factor that in theory has the least impact on model performance (or equivalently, assigns a different $\theta_i(pos)$ to every position), thereby achieving the effect of refusing to pay the tax.
However, actually implementing a different $\theta_i(pos)$ for every single position is quite difficult. For the same reason that Leaky ReRoPE and ReRoPE need to compute attention repeatedly, since RoPE implements relative position through an absolute-position mechanism, a single computation can only realize one fixed $\theta_i$. To have different positions use different values of $\theta_i$, the K in the KV cache would have to be stored before applying RoPE, and different positions would each need to be computed separately multiple times — turning it into an RNN-like recursive process. We know that when an LLM responds in a dialogue turn, this can be split into two phases: prefill and generation, where prefill refers to the computation over the input, and generation is the token-by-token generation phase. Prefill is naturally parallelizable, but if it too were changed into a recursive process like generation, it would obviously slow computation down noticeably when the input is very long (e.g., feeding in an entire paper), making this impractical.
So a compromise is "locally static": during prefill we can compute how many tokens the input contains, and during generation we also set a max_gen_tokens; we add these two numbers together and use that as the $L_{test}$ for the current dialogue turn to compute the corresponding $\theta_i$. After this turn finishes, the next turn updates $L_{test}$ and $\theta_i$ in the same way. This avoids introducing an overly complex implementation or sacrificing efficiency, making it a fairly practical solution — especially when the input is very long, since max_gen_tokens is then much smaller than the number of prefill tokens, and the scale within a single turn is already close to constant.
The idea of Dynamic Scaling could be said to have been pushed to its extreme by CLEX, proposed in CLEX: Continuous Length Extrapolation for Large Language Models: CLEX likewise assigns a different $\theta_i(pos)$ to every position, treating $\theta_i(pos)$ as a continuous function of $pos$ and modeling it with a neural ODE, fitting the ODE's parameters via fine-tuning. This achieves results even better than YaRN, and the experimental results show that by continually applying Dynamic Scaling, one can obtain what amounts to nearly unlimited length extrapolation ability.
Starting Over
Besides Dynamic Scaling, another way to "refuse to pay taxes" is to "start over from scratch" — redesigning the model architecture used for pretraining so that, once training is finished, the model has the built-in potential for length extrapolation without any further modification. In this series, I have two related pieces of exploration: HWFA (Hybrid Window-Full Attention), mentioned in "Transformer Upgrade Path: 9, A New Idea for Global Length Extrapolation", and Key Norm, validated in "Transformer Upgrade Path: 15, Key Normalization Helps with Length Extrapolation".
Here, HWFA replaces attention in the first $L-1$ layers of the model with RoPE + Window Attention using a very small window, while the last attention layer is replaced with NoPE + Full Attention. A model trained this way exhibits a degree of length extrapolation ability without any further modification. A related idea can be found in Focused Transformer: Contrastive Training for Context Scaling, although that work is not aimed at length extrapolation but rather at extending an LLM's context length via simple fine-tuning. The problem with HWFA is that its training performance falls somewhat short of a standard attention model, and to address this the author later proposed an improved version, HWFA2 (i.e., HWFA + ReRoPE), in Transformer Upgrade Roadmap: 14. When HWFA Meets ReRoPE.
Compared with HWFA, HWFA2's Window Attention uses a larger window size, restores the RoPE of the Full Attention layers, and allows more than one layer of Full Attention to be interspersed among the Window Attention layers (rather than putting just one such layer at the end). These changes close the gap with standard Attention in terms of training performance (and occasionally even surpass it), but the drawback is that length extrapolation is no longer available out of the box (RoPE would need to be replaced with ReRoPE) — so it's a bit of a trade-off. Of course, if we set aside extrapolation performance altogether, we can simply treat HWFA2 as an acceleration scheme that reduces model complexity without sacrificing quality. Incidentally, a paper posted on Arxiv last month, Zebra: Extending Context Window with Layerwise Grouped Local-Global Attention, proposes a method called Zebra, which likewise combines several Full Attention layers interspersed among Window Attention layers, similar to HWFA2.
As for Key Norm, it originated from the "accidental discovery" that after applying L2 normalization to the Key in Attention, the model's length extrapolation ability turned out to improve markedly. Further reflection on this finding deepened the author's understanding of length extrapolation. For the standard Attention based on the inner product of Q and K, we can express it as
I need to see the actual text you'd like translated. It looks like only a placeholder (\begin{equation}s(n|m) = \boldsymbol{q}_m\cdot \boldsymbol{k}_n = \Vert\boldsymbol{q}_m\Vert \Vert\boldsymbol{k}_n\Vert \cos(\boldsymbol{q}_m,\boldsymbol{k}_n),\quad p(j|i) = \frac{\exp\left(\frac{s(n|m)}{\sqrt{d}}\right)}{\sum\limits_{j=1}^i \exp\left(\frac{s(n|m)}{\sqrt{d}}\right)}\end{equation}) was included, without any surrounding content or context to translate.
Could you please share the full blog post text (with its placeholders) that you'd like translated into English?
It's clear that to increase the relative attention of $n$ toward a particular $m$, the model has two choices: increase $\Vert\boldsymbol{k}_n\Vert$, or increase $\cos(\boldsymbol{q}_m,\boldsymbol{k}_n)$. Because of the curse of dimensionality, increasing $\Vert\boldsymbol{k}_n\Vert$ is much easier than increasing $\cos(\boldsymbol{q}_m,\boldsymbol{k}_n)$, so whenever possible, the model will preferentially increase $\Vert\boldsymbol{k}_n\Vert$, while $\Vert\boldsymbol{k}_n\Vert$ has nothing to do with $i$ and instead describes absolute importance — this may be one of the reasons behind the attention distribution pattern described by Scissorhands. On the other hand, the model's tendency to increase $\Vert\boldsymbol{k}_n\Vert$ implies that the training for $\cos(\boldsymbol{q}_m,\boldsymbol{k}_n)$ may be insufficient, which is probably the more fundamental reason why attention fails to generalize to longer lengths.
From this, it becomes immediately clear why Key Norm improves length-extrapolation ability. Key Norm normalizes all $\Vert\boldsymbol{k}_n\Vert$ to 1, so the model no longer has the option of "increasing $\Vert\boldsymbol{k}_n\Vert$" — it can only focus on adjusting $\cos(\boldsymbol{q}_m,\boldsymbol{k}_n)$, which makes the training of $\cos(\boldsymbol{q}_m,\boldsymbol{k}_n)$ more thorough. I have also run comparison experiments confirming that Key Norm only exhibits length-extrapolation ability when combined with RoPE; neither Key Norm + NoPE nor plain NoPE shows any extrapolation effect. This is presumably because RoPE's own rotational action enriches the diversity of the angle between $\boldsymbol{q}_m,\boldsymbol{k}_n$ (somewhat like data augmentation), which in turn makes the training of $\cos(\boldsymbol{q}_m,\boldsymbol{k}_n)$ more thorough.
There's also a paper called CoCA: Fusing Position Embedding with Collinear Constrained Attention for Fine-Tuning Free Context Window Extending that proposes a solution from a different angle: it modifies how attention is computed so that every group $\boldsymbol{q}_m^{(i)},\boldsymbol{k}_m^{(i)}$ has $\cos(\boldsymbol{q}_m^{(i)},\boldsymbol{k}_m^{(i)})=1$, where the grouping $i$ here refers to the pairwise grouping of the $\boldsymbol{q},\boldsymbol{k}$ components of RoPE mentioned earlier. This design ensures that the relatively large $\cos(\boldsymbol{q}_m,\boldsymbol{q}_n)$ values are, as much as possible, all ones the model has actually seen during training (the maximum of $\cos$ is just 1), while the parts that are undertrained are confined to the small values (whose post-Softmax probabilities will be small and won't noticeably interfere with the attention distribution), thereby achieving a certain degree of length extrapolation ability. However, CoCA's modification to attention risks lowering the capacity ceiling of each attention head — that is, with the same number of parameters, it may only have the fitting capacity of a standard attention head with head_size/2.
Other approaches
We're getting close to the end of this introduction to length extrapolation. Even though we've already covered quite a lot of ground, it's still hard to give a detailed account of every piece of work on length extrapolation. Below I'll just list, somewhat scattershot, a few other related pieces of work that come to mind.
A first hypothesis was that Attention fails to extrapolate simply because positions "go out of range" at inference time, and a natural fix for this is to perturb the position encodings during training — a kind of data augmentation — so that the model gets used in advance to the positions it will encounter at inference. Two posts in this series, Transformer Upgrade Path: 8. Length Extrapolation and Positional Robustness and Transformer Upgrade Path: 13. Using Leaky ReRoPE in Reverse, fall into this category, as does PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training from a few months back. In my own experiments this class of method has proven not particularly stable, and it introduces extra complexity or randomness, so it's hard to guarantee that it won't affect the model's original scaling law.
有些读者曾提出过疑问：在YaRN的分析中，要做插值的是低频部分，那如果干脆直接去掉低频部分会怎么样？或者类似地，将base调小，让高频部分的比例增加？笔者曾在预训练中尝试过调小RoPE的base，结果是最终效果更差了，并且也没表现出长度外推能力。不过《Scaling Laws of RoPE-based Extrapolation》（知乎有中文版《RoPE外推的缩放法则 —— 尝试外推RoPE至1M上下文》）尝试过另一种方案，是在微调阶段才调小Base，配合短文本的微调后能体现出长文本的外推能力。
From the author's perspective, the practice of shrinking the Base or even discarding low-frequency components isn't really principled — even though it may achieve length extrapolation in some cases, it risks sacrificing the model's inherent capabilities. As Bowen Peng, an author of NTK-RoPE and YaRN, once put it, high frequencies encode local relative distances while low frequencies encode long-range absolute distances, and both are important — the relationship between them is more like a hierarchy. Using the numeral-system framing from Transformer Upgrade Roadmap: 10, RoPE is a Kind of β-ary Encoding, the low-frequency components correspond to the high-order digits; if you keep only the low-order digits and discard the high-order ones, the result is effectively taking a modulus (remainder), which can no longer accurately express positional information. What's more, "high frequency" and "low frequency" are inherently relative notions — a frequency that counts as low for a 10K-length text may well count as high for a 100K-length text.
Here's an interesting recent paper, Fortify the Shortest Stave in Attention: Enhancing Context Awareness of Large Language Models for Effective Tool Use, which found that using different bases for the same model, taking the outputs, and averaging them can improve overall model performance. This suggests that different base values each have their own strengths, and one shouldn't simply reduce it just for the sake of extrapolation.
All in all, despite substantial progress on length extrapolation, it remains a rather mysterious business. For instance, swapping RoPE for ReRoPE at inference time does show some length extrapolation benefit — so shouldn't swapping in ReRoPE during pretraining make the extrapolation even better? Quite the opposite, actually. I've experimented with switching to ReRoPE already during training, and the resulting model showed essentially zero length extrapolation ability. This seems largely connected to the Key Norm analysis discussed earlier: switching to ReRoPE during training reduces the diversity of $\boldsymbol{q}_n,\boldsymbol{k}_m$ angles, which in turn leaves $\cos(\boldsymbol{q}_n,\boldsymbol{k}_m)$ under-trained and ultimately weakens the model's length extrapolation ability. There are also many length extrapolation techniques that may be tied to specific architectures — some early positional encodings reputed to have length extrapolation ability, such as ALIBI, KERPLE, and XPOS, were all tested on Multi-Head Attention + Pre Norm setups, whereas in my own experiments with Single-Head GAU + Post Norm, I've never observed any length extrapolation ability from them. This suggests that our current understanding of length extrapolation is likely still missing an architecture-dependent piece of the puzzle.
Summary
In this article, drawing on my own learning journey, I have tried to lay out the progress made over the past year on the topic of length extrapolation, giving a concise account of the characteristics of the various methods and the ideas behind them, and attempting to string them together into a coherent narrative. I hope this piece helps readers gain a deeper and more systematic understanding of the length extrapolation problem. If there are any errors or omissions, please point them out and correct me.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.