Transformer Upgrade Path: 9. A New Idea for Global Length Extrapolation

Transformer升级之路：9、一种全局长度外推的新思路 苏剑林 · 2023-05-12 · Machine Learning attention generalization extrapolation

When it comes to why Transformers can't handle extremely long sequences, most people's first reaction is usually the quadratic complexity of self-attention. But in fact, even setting aside compute constraints, regular Transformers still can't handle very long sequences, because their length extrapolation ability is poor — specifically, when the input sequence significantly exceeds the training length, model performance typically degrades severely.

Although some related work already exists, the length extrapolation problem is still quite far from being properly solved. This post introduces a reference scheme I've been thinking about, which may currently be the only length extrapolation method that can be used in generative models while retaining global dependency capability.

Review of Existing Methods

Length extrapolation, also known as length generalization, has been partially covered before in Transformer Upgrade Path: 7. Length Extrapolation and Local Attention and Transformer Upgrade Path: 8. Length Extrapolation and Positional Robustness. However, each of them has its own issues. more

The various schemes introduced in the first article all follow the idea of localizing attention. Although they show improvements on certain metrics, in essence this is just making the numbers look a bit better — it doesn't achieve extrapolation with global dependency, so it offers no real help for scenarios that truly require long-range dependency (such as in-context learning). The latter approach enhances robustness to positional signals through random positional perturbation, and could in theory preserve global dependency, but it only applies to encoder models and is not suitable for autoregressive generative models like GPT.

So the length extrapolation problem remains an urgent but still unsolved issue for Transformers. In fact, this problem isn't unique to Transformers — as we discussed before in Google's New Work Attempts to "Revive" RNNs: Can RNNs Shine Again?, the linear RNN models (including the very popular RWKV) also don't have good length extrapolation ability. In today's LLM era, length extrapolation capability is especially important, because we always want models to handle arbitrarily long text, yet we cannot possibly stretch the training samples to arbitrary lengths.

Translation Invariance

Next we'll discuss this in the context of autoregressive Transformers, though the method is also valid for bidirectional-attention encoders. Essentially, localized attention grants the whole model "translation invariance" by restricting attention's receptive range. A simple baseline for translation invariance is Window Attention, as shown below:

Window AttentionDiagram of stacked receptive fields

Suppose the model consists of $L$ stacked layers of Window Attention, with window size $w$; then for the last layer, the maximum receptive field for each token is $(w-1)L+1$. So, assuming the training length is $N$, under the constraint $(w-1)L+1 = \alpha N\,(0 < \alpha \leq 1)$, the model can gain a certain degree of translation invariance, because in that case the model's maximum receptive field never exceeds $N$, so the model's overall receptive field gets adequately trained. The smaller $\alpha$ is, the better the translation invariance usually is.

However, while this approach ensures translation invariance, it brings another problem: since the receptive field of each layer is restricted to within $w$, the capability of the attention mechanism is greatly weakened, leading to training performance worse than that of regular attention (referred to below as Full Attention). Moreover, what we actually want from length extrapolation is not just "translation invariance" but "translation-improves-things-ness" — that is, the further along we go, the better the performance should get (for instance, in an in-context learning scenario, the more examples given, the better the performance should be). So the model should also be able to capture global dependencies.

Global Dependency

To this end, I had the following thought: what Window Attention essentially produces is some kind of $n$-gram feature, except that after stacking multiple layers, this $n$ can become fairly large; whereas a single layer of Full Attention can be regarded as a form of "retrieval" (as suggested by the naming of query, key, value) and "fusion," whose behavior is relatively easy to analyze. Previously, in Viewing Attention's Scale Operation Through Entropy Invariance, we found that a single (full) attention layer can have its length extrapolation ability enhanced by adding a $\log n$ scaling factor.

So I came up with an idea:

If the preceding $L-1$ layers use Window Attention to obtain $n$-gram features, could the last layer be replaced with Full Attention carrying a $\log n$ factor, to retrieve and integrate these features — thereby making up for the performance gap and gaining the ability to capture global dependencies?

To this end, we propose the following combination of attention mechanisms (Hybrid Window-Full Attention, abbreviated HWFA):

1. The preceding $L-1$ layers use "Window Attention + RoPE" with window size $w$, satisfying the constraint $(w-1)(L-1)+1 = \alpha N$, where $N$ is the training length. To balance training performance and extrapolation performance, we recommend choosing $w$ as large as possible while respecting $\alpha\leq 3/4$;
2. The $L$-th layer uses Full Attention with a $\log n$ factor, but without RoPE.

The reason for using RoPE in the earlier layers is that numerous experimental results have shown RoPE helps improve model performance (at least at base/large model scales). The reason for not using RoPE in the last layer is that RoPE positions beyond the training length have never been trained on, which would hurt extrapolation performance. In fact, the RoPE in the preceding $L-1$ layers already provides sufficient positional information to the model, so omitting RoPE in the last layer barely affects training performance.

Experimental Results

Clearly, HWFA is a way of combining attention mechanisms — it can be used in standard multi-head attention, or in attention variants such as GAU. I ran experiments based on GAU_alpha: training length 512, 24 layers of GAU, with the first 23 layers using Window Attention with window size $w=16$. Per-token accuracy was measured, with the baseline being the case where all layers use Full Attention + RoPE (i.e., the conventional default setup).

The results are quite encouraging:

$$\begin{array}{c|cc} \hline \text{test length} & 512 & 4096 \\ \hline \text{Baseline} & 49.41\% & 24.17\% \\ \text{HFWA} & 48.70\% & 80.84\% \\ \hline \end{array}$$

512 represents training accuracy (which can also be called interpolation accuracy), and 4096 represents extrapolation accuracy. Why is the training accuracy only in the 40s while the extrapolation accuracy reaches over 80 — such an exaggerated gap? This is because, when constructing the test samples, I included some samples formed by repeated concatenation, i.e., the same piece of text no longer than 4096 in length, concatenated with itself repeatedly to reach length 4096. Since the latter part of these samples is a repetition of the earlier part, accuracy on that part is very high (i.e., the "correct answer" has essentially already been given earlier in the sequence). This confirms, as expected, that length extrapolation under this design does not sacrifice the ability to capture global dependency.

If we remove the repeated samples and keep only ordinary natural text samples, the results still look decent:

$$\begin{array}{c|cc} \hline \text{test length} & 512 & 4096 \\ \hline \text{Baseline} & 49.41\% & 23.16\% \\ \text{HFWA} & 48.70\% & 48.15\% \\ \hline \end{array}$$

To further verify the global dependency capability, I also ran the even pairs task from Transformer Upgrade Path: 8. Length Extrapolation and Positional Robustness (determining whether the first and last characters are the same). The method in this article achieves 100% extrapolation accuracy, which also shows that the model can learn global dependency (the attention needs to span the entire sequence in order to accurately determine whether the two characters match).

I also ran some ablation experiments, with the following results:

1. Window Attention without RoPE: both interpolation and extrapolation performance drop;
2. Full Attention with RoPE added: extrapolation performance drops;
3. Full Attention without the $\log n$ factor: extrapolation performance drops;
4. Using Window Attention throughout: both interpolation and extrapolation performance drop;
5. Changing to $L-2$ layers of Window Attention + 2 layers of Full Attention: extrapolation performance drops;
6. $w=32$ (in which case $(w-1)(L-1) > N$): extrapolation performance drops.

Comparative Analysis

Some readers might ask: why isn't there a comparison with other methods? The reason may surprise everyone — when I tried applying some of the methods from Transformer Upgrade Path: 7. Length Extrapolation and Local Attention to GAU, I found that they all failed (extrapolation ability was very poor for all of them)!

Why would that be? My first reaction was that those related works all experimented with standard multi-head attention, whereas I was experimenting with GAU. As an attention mechanism, GAU's biggest distinguishing feature is that it is single-headed (unlike the original GAU, the version I experimented with is also softmax-normalized), so I suspected the difference was due to single-head versus multi-head attention. Schemes like ALIBI, Sandwich, and XPOS indeed have parameter designs intended for multi-head attention, and their effectiveness in the single-head setting remains to be verified.

However, after further verification, I found that the difference between single-head and multi-head attention did not have as large an effect on length extrapolation ability as expected, which suggests there must be some other reason at play. It wasn't until a few days ago that I realized another important difference: I had consistently been using the Post-Norm architecture, while mainstream work now uses Pre-Norm. As we analyzed in Why Does Pre-Norm Underperform Post-Norm?, Pre-Norm's depth is actually somewhat "inflated," so when local restrictions are imposed on every attention layer, the features output by Pre-Norm at the end are actually somewhat more localized, which in turn leads to better extrapolation performance.

So, based on the current results, if I stick with the GAU + Post-Norm combination, then the method in this article appears to be the only scheme capable of achieving length extrapolation. This is guaranteed by "translation invariance" and "i.i.d.-ness": the preceding $L-1$ layers of Window Attention, whose total receptive field does not exceed the training length, produce "translation invariance," which yields a series of "i.i.d." features; the last layer of Full Attention then performs a weighted average over these i.i.d. features. From a statistical point of view, the average of i.i.d. variables can extrapolate stably.

In addition, I have also started trying to compare HWFA against other work under standard multi-head attention; I'll share further results as they become available.

Further Thoughts

From my experimental results, we can see that HWFA's combination performs slightly worse in terms of training performance compared to the baseline. So one natural worry is whether this gap will widen further as model scale increases. Or in other words, if the parameter count is scaled up to tens or even hundreds of billions, will such a design retain the same emergent capabilities as the standard design? This is indeed a concern many people have in the LLM era regarding various architectural modifications — the scaling law issue. Admittedly, until HWFA is actually scaled up to the tens-of-billions-of-parameters regime, there's no definitive answer to this question, but my preliminary guess is that there would be some capability bottleneck.

Of course, HWFA can currently only be regarded as a baseline for length extrapolation. Its main purpose is to achieve length extrapolation while retaining global dependency capability, and initial results suggest it has the potential to do so. The next step is to close the gap between HWFA's training performance and the baseline while still retaining global dependency capability. In addition, HWFA can only capture global dependency in the last Full Attention layer, which is likely to create a performance bottleneck as well; but using more such layers would in turn reduce length extrapolation ability — this too is an issue that urgently needs optimizing.

It's worth noting that, since the preceding $L-1$ layers of Window Attention only involve a finite receptive field, in principle it should also be possible to replace them with something like a CNN, as long as the total receptive field doesn't exceed the training length $N$. So it's also worth exploring how to combine the ideas behind HWFA with other base architectures.

Summary

This post introduces a length extrapolation scheme I've been thinking about. By combining Window Attention with Full Attention, it achieves length extrapolation capability while retaining global dependency capability, and should currently be the only length extrapolation method that can be used in generative models while possessing global dependency capability.