Transformer Upgrade Path: 8. Length Extrapolation and Positional Robustness

Transformer升级之路：8、长度外推性与位置鲁棒性 苏剑林 · 2023-01-31 · Machine Learning language model attention positional encoding extrapolation

In the previous post, Transformer Upgrade Path: 7. Length Extrapolation and Local Attention], we discussed the length extrapolation of Transformers, concluding that length extrapolation is fundamentally a mismatch problem between training and inference, and that the main idea for resolving this mismatch is to localize attention — many improvements with good extrapolation are, in some sense, variants of local attention. Admittedly, judging from many metrics of current language models, the local-attention approach does seem to solve the length extrapolation problem. However, this kind of "forced truncation" may not appeal to every reader's taste, since it feels heavily hand-crafted and lacks a natural quality, which also raises doubts about its effectiveness on non-language-model tasks.

In this post, we revisit the length extrapolation problem from the angle of the model's robustness to positional encoding. This approach can improve the Transformer's length extrapolation performance essentially without modifying attention at all, and it is applicable to various positional encodings. Overall, the method is more elegant and natural, and it also applies to non-language-model tasks.

Problem Analysis

In the earlier post, we analyzed the causes of length extrapolation and arrived at the framing that "length extrapolation is a problem of length mismatch between training and inference," with two concrete sources of mismatch:

1. During inference, positional encodings that were never seen during training are used (whether absolute or relative).
2. During inference, the attention mechanism has to handle far more tokens than it did during training.

Point 2 refers to the fact that more tokens make attention more dispersed (i.e., the entropy of attention increases), causing a train/inference mismatch. We already discussed and partially addressed this in Viewing Attention's Scale Operation through Entropy Invariance]: the answer is to modify attention from

\begin{equation}Attention(Q,K,V) = softmax\left(\frac{QK^{\top}}{\sqrt{d}}\right)V\end{equation}

to

\begin{equation}Attention(Q,K,V) = softmax\left(\frac{\log_{m} n}{\sqrt{d}}QK^{\top}\right)V\end{equation}

where $m$ is the training length and $n$ is the inference length. With this modification (henceforth "$\log n$-scaled attention"), the entropy of attention varies much more smoothly with length, alleviating this mismatch. My own experimental results show that, at least on MLM tasks, "$\log n$-scaled attention" gives better length extrapolation performance.

So we can consider mismatch point 2 to be preliminarily resolved, and we should now focus on mismatch point 1.

Random Positions

Mismatch point 1 — "positional encodings unseen during training are used at inference time" — suggests that the fix should be: "during training, also train on the positional encodings that will be used at inference time." An ACL22 paper, still under anonymous review at the time, Randomized Positional Encodings Boost Length Generalization of Transformers], was the first to consider the problem from this angle and to propose a solution.

The idea of the paper is simple:

Random position training. Let $N$ be the training length (denoted $N=40$ in the paper) and $M$ be the inference length (denoted $M=500$ in the paper). Choose a large value $L > M$ (a hyperparameter, denoted $L=2048$ in the paper). During training, a sequence of length $N$ originally corresponds to the position sequence $[0,1,\cdots,N-2,N-1]$. Now, instead, we randomly sample $N$ distinct values without replacement from $\{0,1,\cdots,L-2,L-1\}$, sort them in increasing order, and use this as the position sequence for the current sequence.

Reference code using numpy:

def random_position_ids(N, L=2048):
    """从[0, L)中随机不重复挑N个整数，并从小到大排列
    """
    return np.sort(np.random.permutation(L)[:N])

At inference time, one can likewise sample the position sequence randomly, or simply take evenly spaced points from the interval (my own experiments suggest evenly spaced points tend to work slightly better). This resolves the issue of the inference-time positional encoding never having been trained. It is not hard to see that this is a fairly simple training trick (henceforth "random position training"), whose goal is to make the Transformer more robust to the choice of positions. But as we'll see below, it delivers a clear improvement in length extrapolation. I also ran experiments on MLM tasks, and the results show it is effective there too, with an even more pronounced improvement when combined with "$\log n$-scaled attention" (the original paper does not include this "$\log n$-scaled attention" step).

A New Benchmark

Many related works, including the various Local Attention variants mentioned in the previous post, build their evaluation metrics around language model tasks. But whether it's a unidirectional GPT or a bidirectional MLM, both are highly dependent on local information (locality). So it is quite possible that previous approaches only appeared to extrapolate well because language modeling itself is highly local — if we switched to a non-local task, performance might drop considerably. Perhaps for this very reason, the paper under discussion does not use a conventional language model task for evaluation, but instead adopts a length-generalization benchmark that Google specifically proposed last year in Neural Networks and the Chomsky Hierarchy] (hereafter referred to as the "CHE benchmark," i.e., the "Chomsky Hierarchy Evaluation Benchmark"). This gives us a fresh perspective for understanding length extrapolation.

This benchmark contains multiple tasks, divided into three difficulty tiers — R (Regular), DCF (Deterministic Context-Free), and CS (Context-Sensitive) — with difficulty increasing across tiers. A brief description of each task:

Even Pairs, difficulty R: Given a binary sequence, e.g. "aabba," determine whether the total count of "ab" and "ba" among its 2-grams is even. In this example, the 2-grams are aa, ab, bb, ba, of which ab and ba together number 2, so the output is "Yes." This task is equivalent to checking whether the first and last characters of the sequence are the same.
Modular Arithmetic (Simple), difficulty R: Compute the value of an expression built from the five numbers $\{0, 1, 2, 3, 4\}$ and the three operators $\{+,-,\times\}$, and output the result modulo 5. For example, given the input $1 + 2 − 4$, the value is $-1$, which modulo 5 is $4$, so the output is $4$.
Parity Check, difficulty R: Given a binary sequence, e.g. "aaabba," determine whether the number of b's is even. In this example there are 2 b's, so the output is "Yes."
Cycle Navigation, difficulty R: Given a ternary sequence where each element represents one of $+0$, $+1$, $-1$, output the result of applying the sequence of operations starting from 0, modulo 5. For example, if $0,1,2$ represent $+0,+1,-1$ respectively, then $010211$ represents $0 + 0 + 1 + 0 − 1 + 1 + 1 = 2$, and modulo 5 the output is $2$.
Modular Arithmetic, difficulty DCF: Compute the value of an expression built from the five numbers $\{0, 1, 2, 3, 4\}$, parentheses $(,)$, and the three operators $\{+,-,\times\}$, and output the result modulo 5. For example, given the input $−(1−2)\times(4−3\times(−2))$, the result is $10$, which modulo 5 is $0$, so the output is $0$. Compared with the Simple version, this task adds "parentheses," making the computation more complex.
Reverse String, difficulty DCF: Given a binary sequence, e.g. "aabba," output its reverse. In this example, the output should be "abbaa."
Solve Equation, difficulty DCF: Given an equation built from the five numbers $\{0, 1, 2, 3, 4\}$, parentheses $(,)$, the three operators $\{+,-,\times\}$, and an unknown $z$, find the value of $z$ that satisfies the equation modulo 5. For example, given $−(1−2)\times(4−z\times(−2))=0$, we get $z=3$. Although solving an equation looks harder, since the equation is constructed by taking a Modular Arithmetic expression and replacing one of the numbers with $z$, a solution is guaranteed to exist and to lie in $\{0, 1, 2, 3, 4\}$. Hence, in principle we could solve it by enumeration combined with Modular Arithmetic, so its difficulty is comparable to Modular Arithmetic.
Stack Manipulation, difficulty DCF: Given a binary sequence, e.g. "abbaa," and a sequence of stack operations composed of "POP / PUSH a / PUSH b," e.g. "POP / PUSH a / POP," output the final state of the stack. In this example, the output should be "abba."
Binary Addition, difficulty CS: Given two binary numbers, output the binary representation of their sum. For example, given the inputs $10010$ and $101$, the output is $10111$. Note that both numbers must be fed to the model at the character level rather than the numeric level for training and inference, and the two numbers are presented serially rather than aligned in parallel (i.e., the input can be thought of as the string $10010+101$).
Binary Multiplication, difficulty CS: Given two binary numbers, output the binary representation of their product. For example, given the inputs $100$ and $10110$, the output is $1011000$. As with Binary Addition, this requires character-level rather than numeric-level input for training and inference, with the two numbers presented serially rather than aligned in parallel (i.e., the input can be thought of as the string $100\times 10110$).
Compute Sqrt, difficulty CS: Given a binary number, output the floor of its square root in binary. For example, given the input $101001$, the output is $\lfloor\sqrt{101001}\rfloor=101$. This is comparable in difficulty to Binary Multiplication, since at minimum we could determine the result by enumerating from $0$ up to the given number combined with Binary Multiplication.
Duplicate String, difficulty CS: Given a binary sequence, e.g. "abaab," output the sequence repeated once. In this example, the output should be "abaababaab." This simple-looking task may appear to be difficulty R, but is actually CS — readers are invited to think about why.
Missing Duplicate, difficulty CS: Given a binary sequence with a missing value, e.g. "ab_aba," where it is known that the original complete sequence is a duplicated sequence (as in the previous task), predict the missing value. In this example, the output should be "a."
Odds First, difficulty CS: Given a binary sequence $t_1 t_2 t_3 \cdots t_n$, output $t_1 t_3 t_5 \cdots t_2 t_4 t_6 \cdots$. For example, given the input "aaabaa," the output would be "aaaaba."
Bucket Sort, difficulty CS: Given a sequence of $n$ numeric values (each value from a given set of $n$ possible numbers), return the sequence sorted in ascending order. For example, given the input $421302214$, the output should be $011222344$.

Notice that all these tasks share a common feature: their computations follow fixed, simple rules, and in principle the inputs are of unbounded length. This means we can train on short sequences and then test whether what was learned on short sequences generalizes to longer ones. In other words, this benchmark serves as a very strong test of length extrapolation.

Experimental Results

First, let's look at the results from the original paper Neural Networks and the Chomsky Hierarchy], which compares several RNN models against Transformer models (the evaluation metric is the average per-string accuracy, not the overall exact-match rate):

] Comparison of several models across several length-extrapolation test tasks

The result may come as a surprise: the currently "hot" Transformer has the worst length extrapolation performance (here the Transformer was tested with different positional encodings, and the best result on each task was taken). The best performer is Tape-RNN. The paper assigns them the following ratings:

$$\underbrace{\text{Transformer}}_{\text{R}^-} < \underbrace{\text{RNN}}_{\text{R}} < \underbrace{\text{LSTM}}_{\text{R}^+} < \underbrace{\text{Stack-RNN}}_{\text{DCF}} < \underbrace{\text{Tape-RNN}}_{\text{CS}}$$

The random position training method proposed in Randomized Positional Encodings Boost Length Generalization of Transformers], introduced above, recovers some of the Transformer's disadvantage:

] Comparison of length extrapolation performance for Transformers with different positional encodings, with and without random position training

We can see that, under random position training, Transformers with every kind of positional encoding show a marked improvement. This further confirms the conclusion from the previous post: length extrapolation performance has little to do with the specific design of positional encoding itself. Notably, random position training is the first method to achieve perfect accuracy on the Bucket Sort task. Although overall performance is still lacking, this represents substantial progress compared to previous results (I wonder whether combining it with "$\log n$-scaled attention" would improve things further?). It's also worth noting that ALIBI, which performs well on language modeling tasks, shows no particular advantage on the CHE benchmark — in fact, once random position training is added, its average score is even worse than RoPE's. This lends preliminary support to the earlier conjecture: the good performance of various Local Attention variants is likely due to the severe locality inherent in language-model-based evaluation tasks; on the non-local CHE benchmark, these methods show no advantage.

Rethinking the Mechanism

On closer reflection, "random position training" is rather puzzling. For simplicity, suppose $L=2048,N=64,M=512$; then the average position sequence used during training is roughly $[0, 32, 64, \cdots, 2016]$, while the position sequence used during inference is $[0, 4, 8, \cdots, 2044]$. The spacing between adjacent positions differs between training and inference — one could call this yet another kind of mismatch — and yet it still performs well. Why is that?

We can understand this from the perspective of "order." Since the position ids during training are randomly sampled, the gaps between adjacent positions are also random. So whether it's relative or absolute position, the model is unlikely to be able to rely on precise position ids to extract positional information; instead, it must rely on a fuzzier positional signal — more precisely, it encodes position through the ordering of the position sequence rather than through the position ids themselves. For instance, the position sequences [1,3,5] and [2,4,8] are treated as equivalent, because both are simply increasing sequences. Random position training "forces" the model to learn an equivalence class: all increasing position sequences are equivalent and mutually interchangeable. This is the real meaning of positional robustness.

However, my own experiments on MLM show that learning this "equivalence class" is still somewhat difficult for the model. A more ideal approach would be to keep using random positions during training (so that the positional encodings used at inference time are also trained on), but ensure that the initial portion of the inference-time position sequence matches the average result of the random positions. Returning to the earlier example: if the position sequence used at inference is $[0, 4, 8, \cdots, 2044]$, then we would want the average random position result during training to be $[0, 4, 8, \cdots, 252]$ (i.e., the first $N$ elements of the sequence $[0, 4, 8, \cdots, 2044]$), rather than $[0, 32, 64, \cdots, 2016]$. This would make the consistency between training and inference much tighter.

Extending the Idea

This led me to consider the following approach:

Equal-mean random position training. Let $n$ follow a distribution with mean $N$ and sample space $[0, \infty)$. During training, randomly sample a $n$, then uniformly select $N$ points from $[0, n]$ to use as the position sequence.

Reference code:

def random_position_ids(N):
    """先随机采样n，然后从[0, n]均匀取N个点
    """
    n = sample_from_xxx()
    return np.linspace(0, 1, N) * n

Note that the position sequences sampled this way are floating-point numbers, so this approach does not apply to discrete, trainable positional encodings — only to functional positional encodings such as Sinusoidal] or RoPE]. In what follows, we assume only functional positional encodings are used.

The biggest challenge with this idea is choosing an appropriate sampling distribution. My first instinct was the Poisson distribution], but since both the mean and variance of a Poisson distribution equal $n$, by the "3$\sigma$ rule" it can only extrapolate to a length of $n+3\sqrt{n}$, which is clearly too short. After some experimentation, I found two distributions that work well: one is the exponential distribution], whose mean and standard deviation are both $n$, so even by the "3$\sigma$ rule" it can extrapolate to a length of $4n$ — a fairly good range (in practice, even longer). The other is the beta distribution], defined on $[0,1]$; we can treat the test length as 1, so the training length is $N/M\in(0,1)$. The beta distribution has two parameters $\alpha,\beta$ with mean $\frac{\alpha}{\alpha+\beta}$, so once we fix the mean to equal $N/M$, we still have an extra degree of freedom to control the probability mass near $1$ — useful for cases where we want to push the extrapolation range even further.

My experiments show that combining "equal-mean random position training" with "$\log n$-scaled attention" achieves the best extrapolation performance on the MLM task (training length 64, test length 512, using the exponential distribution as the sampling distribution). Since I haven't previously run experiments on the CHE benchmark, I wasn't able to test this combination there yet, and will leave that for a future opportunity.

Summary

This post revisited the length extrapolation of Transformers from the angle of positional robustness, arriving at new schemes such as "random position training" for enhancing length extrapolation. We also introduced the new "CHE benchmark," which, compared to conventional language model tasks, exhibits stronger non-locality and thus provides a more effective way to evaluate work related to length extrapolation. Under this benchmark, the previous attention-localization methods do not show particularly outstanding performance; by comparison, "random position training" performs better. This reminds us that we should evaluate the effectiveness of such methods on a broader range of tasks, rather than restricting ourselves solely to language modeling tasks.