Transformer Upgrade Road: 7. Length Extrapolation and Local Attention

ALiBi

As the "pioneering work," ALiBi cannot be ignored. It comes from the paper "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation". In hindsight, the modification made by ALiBi is very simple; it just changes the Attention calculation before Softmax 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 different values set for each head. From this definition, one can see the similarity between ALiBi and the baseline model. Both subtract a non-negative matrix before Softmax, but the subtracted non-negative matrices are different. ALiBi can be seen as a "smooth version" of the baseline model:

ALiBi is a very simple (and of course effective) smooth local attention trick. However, if it is understood as "position encoding," it might not be quite appropriate. If Equation [eq:alibi] is generalized to bidirectional attention, then because |m - n|=|n - m|, the model would theoretically be unable to distinguish "left" from "right" and could only recognize the relative distance, which is clearly insufficient for accurately identifying position information. As for its good performance on unidirectional language models, it is because unidirectional language models can achieve non-trivial results (significantly higher than random) even without any position encoding. The local attention imposed by ALiBi serves to enhance locality, fitting the characteristics of the language modeling task itself.

KERPLE

KERPLE comes from the paper "KERPLE: Kernelized Relative Positional Embedding for Length Extrapolation". It is essentially a simple generalization of ALiBi, introducing two trainable parameters r_1, r_2 to generalize Equation [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} With generalization and trainable parameters, it is not surprising that KERPLE achieves better results than ALiBi. However, I must severely criticize the KERPLE paper for its pretentiousness. According to the layout, the third section of the original paper is the theoretical support, but it is clearly an irrelevant section introduced just to forcibly increase the mathematical depth of the article. In substance, it does not help in understanding KERPLE at all and might even decrease the reader’s interest (to put it bluntly, it serves the reviewers, not the readers).

Sandwich

Sandwich was also developed by the authors of KERPLE, from "Receptive Field Alignment Enables Transformer Length Extrapolation", which was posted on Arxiv just last month. It replaces Equation [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 are Sinusoidal position encodings, and \lambda > 0 is a hyperparameter. From "Transformer Upgrade Road: 1. Tracing the Source of Sinusoidal Position Encoding", we know that \boldsymbol{p}_m^{\top}\boldsymbol{p}_n is a scalar function of m-n, and on average, it is a monotonically decreasing function of |m-n|. Therefore, its effect is similar to -\lambda|m-n|. The reason for emphasizing "on average" is that \boldsymbol{p}_m^{\top}\boldsymbol{p}_n as a whole is not strictly monotonic but oscillates downwards, as shown in the figure:

If necessary, we can also convert Sandwich into a form like RoPE where "absolute position encoding implements relative position encoding." This only requires 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 position information through concatenation, and the Attention result is equivalent to relative position encoding. However, it seems this conversion currently only has theoretical value because concatenation increases the vector dimension, which in turn increases the computational cost of Attention.

XPOS

XPOS comes from the paper "A Length-Extrapolatable Transformer", which appeared on Arxiv on the same day as Sandwich. It is a direct generalization of RoPE. We know that the basic solution for 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 derived RoPE, I assumed that "Q and K undergo the same transformation." In fact, from the perspective of "absolute position implementing relative position," there is no need to restrict both transformations to be identical. For example, XPOS 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. Thus: \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 total result still only depends on the relative position m-n. However, the problem now is that the exponent is m-n rather than |m-n|. As long as \xi \neq 1, one side will always diverge. The cleverness of XPOS is that it chooses the same scenario as many related works—focusing only on unidirectional language models—so we only use the attention part where m \geq n! In this case, simply choosing \xi \in (0,1) achieves the effect of decay with relative distance.

In fact, the additional exponential decay term \xi^m is not an original creation of XPOS. In the literature I have read, the same term first appeared in PermuteFormer, although PermuteFormer mainly focused on linear Attention scenarios. In terms of details, XPOS assigns different \xi to each block, but in private communication with the author, the author performed experiments sharing the same \xi and found that the improvement from setting different \xi was almost negligible. Additionally, we must appropriately control the value of \xi to prevent \xi^{-n} from overflowing when n is large.

It is worth noting that the decay with relative distance here is directly multiplied by the Attention Score before Softmax. The result is that Scores for larger relative distances become very close to 0, rather than tending toward negative infinity as in the previous designs. e^0 does not have the effect of tending toward zero. Therefore, such a design is not a variant of local Attention, and thus its performance does not reach SOTA. To make up for this gap, XPOS designed a special local Attention (called Blockwise Causal Attention, or BCA), which, when added, can fill this gap. During communication, the author stated that BCA was used because of implementation advantages; in reality, the local Attention of the baseline model works better. So, for extrapolation, one must look at local attention.

The experiments in the original paper are very rich and worth referencing; I recommend everyone read it carefully!