The Path to the Transformer, Part 7: Length Extrapolation and Local Attention
For Transformer models, length extrapolation is a desirable property we've long been chasing: it asks whether a model trained on short sequences can be applied to long sequences without any fine-tuning while still maintaining good performance. We pursue length extrapolation for two reasons—one is theoretical completeness, since this feels like a property an ideal model ought to have; the other is practical training efficiency, since it would let us train (on shorter sequences, at lower cost) a model usable on long sequences.
Below we analyze the key ideas for strengthening the length extrapolation of Transformers, and from this derive a "super-strong baseline" scheme. We then use this baseline as a lens through which to examine a number of related research works.
A Common Misconception
The first work to explicitly study Transformer length extrapolation was probably ALIBI, which came out in mid-2021—not that long ago, actually. Why did it take so long (relative to the original Transformer paper in 2017) for someone to specifically tackle this problem? Presumably because, for a long time, we simply took it for granted that Transformer length extrapolation was a matter of positional encoding—find a better positional encoding, and the problem would be solved.
Indeed, comparing the extrapolation performance of various existing positional encodings does turn up some evidence in favor of this view. For instance, several of the experimental results shared later show that relative positional encodings tend to extrapolate better on average than absolute ones; and function-based relative positional encodings like RoPE extrapolate somewhat better than trainable relative positional encodings. So it might seem that if we just keep refining the form of positional encoding, we'll eventually give Transformers better length extrapolation and solve the problem. However, things aren't so rosy. RoPE is considered one of the better-extrapolating positional encodings, and yet it can only extrapolate to roughly 10%-20% additional length before performance stays intact; beyond that, performance drops sharply. That ratio is far from what we'd hoped for—ideally, extrapolation worth the name should reach several times the trained length. So it's easy to imagine that relying solely on improving positional encoding to improve Transformer length extrapolation would take who-knows-how-long before we see extrapolation over substantially longer ranges.
Intuitively, many readers probably feel that function-based positional encodings like Sinusoidal or RoPE, having no trainable parameters, should extrapolate well. But that's not actually the case—this class of positional encodings shows no particular advantage in length extrapolation. Why is that? It's because when people assume such function-based positional encodings should extrapolate well, they forget the basic precondition for that to hold: "smoothness."
Extrapolation is, at its core, inferring the whole from the local. We should find this idea familiar—the Taylor series approximation is a classic example: knowing the values of several derivatives at a single point lets us effectively estimate the function's values in a neighborhood, and this relies on the function's high-order smoothness (higher derivatives existing and being bounded). But are Sinusoidal or RoPE functions of this kind? No, they are not. They are combinations of a series of sine and cosine functions, whose phase function is $k/10000^{2i/d}$; when $2i/d\approx 0$, the function is approximately $\sin k, \cos k$, which is a high-frequency oscillating function with respect to the position $k$, rather than something like a straight line or a function that asymptotically approaches one. So models based on it tend to have hard-to-predict extrapolation behavior. Could we design a non-oscillating positional encoding? That's difficult—if a positional encoding function doesn't oscillate, it usually lacks enough capacity to encode a sufficient amount of positional information. In a sense, the very complexity of the positional encoding function is itself a requirement for encoding position.
A Super-Strong Baseline
In fact, a more accurate characterization would be:
Length extrapolation is fundamentally a problem of mismatch between the training length and the inference length.
Specifically, there are two sources of mismatch:
1. At inference time, positional encodings that were never seen during training get used (whether absolute or relative);
2. At inference time, the attention mechanism has to handle far more tokens than it did during training.
Point 1 is probably easy enough to grasp—things that weren't trained for can't be guaranteed to work well; this is a very real phenomenon in deep learning, and it holds even for function-based positional encodings like Sinusoidal or RoPE. As for point 2, readers might be a bit puzzled: isn't attention, in theory, able to handle sequences of arbitrary length? How does a mismatch between training and inference length matter here? The answer is entropy. We already analyzed this in Viewing Attention's Scale Operation Through Entropy Invariance: the more tokens being averaged over by attention, the more "uniform" the resulting distribution becomes (higher entropy), i.e. the more diffuse the attention. A short training length, on the other hand, means lower attention entropy and more concentrated attention. This is another kind of train/inference mismatch, and it too affects performance.
As it turns out, for Transformer models with relative positional encoding, a remarkably simple attention mask can solve both problems at once, achieving performance close to SOTA:
Super-strong baseline model (bidirectional attention version)
Super-strong baseline model (unidirectional attention version)
It's not hard to see what's going on: this simply turns the attention used at inference time into a local attention, where each token can only see as many tokens as were seen during training. This way, the number of tokens visible to each token matches what it was during training, which resolves problem 2; and since the positional encoding is relative, with positions counted relative to the current token as the origin, this local attention also never requires more unseen position codes than during training, which resolves problem 1. So this single, simple attention mask solves both difficulties of length extrapolation at once, without requiring any retraining of the model. What's even more astonishing is that across various experimental results, when this is used as the baseline, the relative gains reported by comparable follow-up works turn out to be rather thin—meaning this baseline is already very close to SOTA on its own. It truly earns the title of a fast-and-effective "super-strong baseline."
Reviewing the Literature
Since ALIBI, quite a few works have indeed devoted themselves to studying Transformer length extrapolation. In this section, I've studied and organized some of the representative works, and from them we can see that they basically all share many commonalities with the baseline model above—one could even say they are, to some extent, variants of the baseline model. This will further deepen our appreciation for the profound connection between length extrapolation and the locality of attention.
ALIBI
As the "pioneering work," ALIBI is unavoidable. It comes from the paper Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. Looking back, the change ALIBI makes is quite simple: right before the softmax, it changes the attention computation from $\boldsymbol{q}_m^{\top}\boldsymbol{k}_n$ to
\begin{equation}\boldsymbol{q}_m^{\top}\boldsymbol{k}_n - \lambda|m - n|\label{eq:alibi}\end{equation}
where $\lambda > 0$ is a hyperparameter, with a different value set for each head. From this definition, we can already see ALIBI's resemblance to the baseline model: both subtract a non-negative matrix before the softmax, just a different non-negative matrix. ALIBI can be viewed as a "smoothed version" of the baseline model:
Matrix subtracted by the baseline model
Matrix subtracted by ALIBI
ALIBI is a rather simple (though certainly effective) technique for smooth local attention, but it wouldn't be quite right to understand it as a "positional encoding." If we try to generalize equation $\eqref{eq:alibi}$ to bidirectional attention exactly as written, then since $|m - n|=|n - m|$, the model would presumably be unable to distinguish "left" from "right" and could only recognize the magnitude of relative distance—clearly this can't accurately capture positional information. As for why ALIBI performs well on unidirectional language models, that's because unidirectional language models can already achieve non-trivial performance (clearly above random) even without any positional encoding at all; the local attention that ALIBI imposes simply strengthens locality, which fits naturally with the nature of the language modeling task itself.
KERPLE
KERPLE comes from the paper KERPLE: Kernelized Relative Positional Embedding for Length Extrapolation, and is essentially a straightforward generalization of ALIBI. It introduces two trainable parameters $r_1,r_2$ to generalize equation $\eqref{eq:alibi}$:
\begin{equation}\left\{\begin{aligned}&\boldsymbol{q}_m^{\top}\boldsymbol{k}_n - r_1|m - n|^{r_2} ,\qquad\qquad r_1 >0, 0 < r_2 \leq 2\\ &\boldsymbol{q}_m^{\top}\boldsymbol{k}_n - r_1\log(1+r_2|m - n|),\qquad\qquad r_1, r_2 > 0 \end{aligned}\right.\label{eq:kerple}\end{equation}
Given that it's both a generalization and comes with trainable parameters, it's no surprise that KERPLE achieves better results than ALIBI. That said, I do want to sharply criticize the KERPLE paper for its unnecessary obfuscation—based on its layout, Section 3 of the original paper is presented as the theoretical foundation of the work, but it's quite clearly just padding introduced to artificially inflate the mathematical "depth" of the paper, contributing essentially nothing to understanding KERPLE, and if anything dampening the reader's enthusiasm (to put it bluntly, it's there to please reviewers, not readers).
Sandwich
Sandwich is also from the KERPLE authors, from the paper Receptive Field Alignment Enables Transformer Length Extrapolation, posted to Arxiv just last month. It replaces equation $\eqref{eq:alibi}$ with
\begin{equation}\boldsymbol{q}_m^{\top}\boldsymbol{k}_n + \lambda\boldsymbol{p}_m^{\top}\boldsymbol{p}_n\label{eq:sandwich}\end{equation}
where $\boldsymbol{p}_m,\boldsymbol{p}_n$ is the Sinusoidal positional encoding and $\lambda > 0$ is a hyperparameter. From The Path to the Transformer, Part 1: Tracing the Origins of Sinusoidal Positional Encoding we know that $\boldsymbol{p}_m^{\top}\boldsymbol{p}_n$ is a scalar function of $m-n$, and on average is a monotonically increasing function of $|m-n|$, so its role is similar to that of $-\lambda|m-n|$. The reason for emphasizing "on average" is that $\boldsymbol{p}_m^{\top}\boldsymbol{p}_n$ isn't strictly monotonic overall but rather oscillates while decreasing, as shown in the figure below:
Plot of dot(p_m, p_n) (with d/2 subtracted)
If we wanted to, we could also convert Sandwich into the form used by RoPE, i.e. "implementing relative positional encoding via absolute positional encoding," simply by noting that
\begin{equation}\boldsymbol{q}_m^{\top}\boldsymbol{k}_n + \lambda\boldsymbol{p}_m^{\top}\boldsymbol{p}_n = \left[\boldsymbol{q}_m, \sqrt{\lambda}\boldsymbol{p}_m\right]^{\top}\left[\boldsymbol{k}_n, \sqrt{\lambda}\boldsymbol{p}_n\right]\end{equation}
In other words, Sandwich supplements absolute positional information via concatenation, and the resulting attention behaves as if it were relative positional encoding. However, for now this conversion seems to have only theoretical value, since concatenation increases the vector dimension and would actually further increase the computational cost of attention.
XPOS
XPOS comes from the paper A Length-Extrapolatable Transformer, which appeared on Arxiv the same day as Sandwich. It's a natural, continuous generalization of RoPE. Recall that the basic solution to RoPE is:
\begin{equation}\boldsymbol{q}_m\to \boldsymbol{\mathcal{R}}_m\boldsymbol{q}_m,\quad \boldsymbol{k}_n\to \boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n\end{equation}
where $\boldsymbol{\mathcal{R}}_n=\begin{pmatrix}\cos n\theta & -\sin n\theta\\ \sin n\theta & \cos n\theta\end{pmatrix}$. When I originally derived RoPE, I assumed that "$Q$ and $K$ undergo the same transformation"; but in fact, purely from the perspective of "implementing relative position via absolute position," there's no real need to constrain the two to have the same transformation form. For instance, XPOS instead considers
\begin{equation}\boldsymbol{q}_m\to \boldsymbol{\mathcal{R}}_m\boldsymbol{q}_m \xi^m,\quad \boldsymbol{k}_n\to \boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n \xi^{-n}\end{equation}
where $\xi$ is a scalar hyperparameter, and so
\begin{equation}(\boldsymbol{\mathcal{R}}_m\boldsymbol{q}_m \xi^m)^{\top}(\boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n \xi^{-n}) = \boldsymbol{q}_m^{\top}\boldsymbol{\mathcal{R}}_{n-m}\boldsymbol{k}_n \xi^{m-n}\end{equation}
The overall result still depends only on the relative position $m-n$. However, the issue now is that the exponent is $m-n$ rather than $|m-n|$, and as long as $\xi\neq 1$, one side or the other will always diverge. XPOS's clever move is that, like many related works, it chose a specific setting—focusing only on unidirectional language models—so that only the $m\geq n$ part of the attention ever gets used! In that case, simply choosing $\xi\in(0,1)$ achieves the effect of decay with increasing relative distance.
It's worth noting that this additional exponential decay term $\xi^m$ is not actually XPOS's invention. Among the papers I've read, the same term first appeared in PermuteFormer, though PermuteFormer was mainly concerned with the linear attention setting. In terms of details, XPOS assigns a different $\xi$ to each block, but when I discussed this privately with the authors, they ran a supplementary experiment sharing a single $\xi$ across blocks and found that the improvement from having distinct $\xi$ values was nearly negligible. We should also keep the value of $\xi$ appropriately controlled, to prevent $\xi^{-n}$ from overflowing when $n$ is large.
It's worth pointing out that here, the decay with relative distance is applied by direct multiplication onto the attention score before the softmax, so that scores at large relative distances become very close to zero, rather than tending to negative infinity as in the earlier designs. $e^0$ doesn't actually go to zero, so this design isn't really a variant of local attention, and correspondingly its performance falls short of SOTA. To make up this gap, XPOS designs a special form of local attention (called Blockwise Causal Attention, or BCA), and adding it closes the gap. When I discussed this with the authors, they mentioned that BCA was chosen for implementation convenience—in practice, the baseline model's local attention performs even better. So in the end, it seems that for extrapolation, local attention really is what matters.
The experiments in the original paper are rich and well worth consulting—I'd recommend reading it closely.
Summary
This post has summarized work aimed at enhancing the length extrapolation ability of Transformers, including a simple but powerful baseline scheme along with several works specifically focused on length extrapolation. From all of this, we can see that these works are, in essence, all variants of the baseline scheme—local attention—which turns out to be one of the key ingredients for length extrapolation.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.