The Road to Upgrading Transformer: 15. Key Normalization Boosts Length Extrapolation

Transformer升级之路：15、Key归一化助力长度外推 苏剑林 · 2023-11-20 · Machine Learning attention positional encoding generalization extrapolation

Broadly speaking, current length-extrapolation techniques for Transformers can be divided into two categories: one is post-hoc modification, such as NTK-RoPE, YaRN, and ReRoPE, which directly modify the inference model and can achieve some degree of length extrapolation without any fine-tuning, but with the drawback that none of them preserve the model's identity behavior within the training length. The other category is naturally pre-hoc modification, such as ALIBI, KERPLE, XPOS, and HWFA, which can achieve a certain degree of length extrapolation without further changes, but the corresponding modifications need to be introduced before training, so they cannot be applied to off-the-shelf models without fine-tuning — and whether this class of methods can scale up has not yet gained wide acceptance.

In this post, I'll introduce a length-extrapolation scheme that I stumbled upon by accident — "KeyNorm" — which applies L2 normalization to the Key sequence in attention. Clearly it belongs to the pre-hoc modification category, but the change it makes to the attention mechanism is very small, which makes it look quite promising for scaling up.

Original Motivation

The reason I call this an "accidental discovery" is that the original motivation for this change had nothing to do with length extrapolation — it came from an attempt to replace the scaling scheme in Scaled Dot-Product Attention. As we know, the standard definition of attention is (this post mainly considers the causal case)

\begin{equation}\boldsymbol{o}_i = \frac{\sum_{j = 1}^i\exp\left(\frac{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}{\sqrt{d}}\right)\boldsymbol{v}_j}{\sum_{j = 1}^i\exp\left(\frac{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}{\sqrt{d}}\right)},\quad \boldsymbol{q}_i,\boldsymbol{k}_j\in\mathbb{R}^d\label{eq:sdpa}\end{equation}more

where the scale factor $\frac{1}{\sqrt{d}}$ has already been explained and even generalized multiple times, e.g., in A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers, Viewing Attention's Scale Operation Through Entropy Invariance, and Viewing Attention's Scale Operation Through Gradient Maximization. The standard derivation is carried out under the assumption that "$\boldsymbol{q}_i,\boldsymbol{k}_j$ are each independently sampled from a distribution with mean 0 and variance 1," and under this same assumption we also have

\begin{equation}\Vert\boldsymbol{q}_i\Vert\approx \sqrt{d},\quad \Vert\boldsymbol{k}_j\Vert\approx \sqrt{d}\end{equation}

This is because

\begin{equation}\Vert\boldsymbol{x}\Vert^2 = \sum_{i=1}^d x_i^2 = d\times\frac{1}{d}\sum_{i=1}^d x_i^2\approx d\,\mathbb{E}_{x\sim\mathcal{N}(0,1)}[x^2] = d\end{equation}

For related generalizations, see also The Amazing Johnson-Lindenstrauss Lemma: Theory. This approximation implies that, at the initial stage of training, Equation $\eqref{eq:sdpa}$ has the same effect as the following two variants:

\begin{align}\color{red}{\text{Q}}\text{uery}\color{red}{\text{N}}\text{orm:}\quad\boldsymbol{o}_i =&\, \frac{\sum_{j = 1}^i\exp\left(\tilde{\boldsymbol{q}}_i\cdot \boldsymbol{k}_j\right)\boldsymbol{v}_j}{\sum_{j = 1}^i\exp\left(\tilde{\boldsymbol{q}}_i\cdot \boldsymbol{k}_j\right)},\qquad \tilde{\boldsymbol{q}}_i = \frac{\boldsymbol{q}_i}{\Vert\boldsymbol{q}_i\Vert} \\[5pt] \color{red}{\text{K}}\text{ey}\color{red}{\text{N}}\text{orm:}\quad\boldsymbol{o}_i =&\, \frac{\sum_{j = 1}^i\exp\left(\boldsymbol{q}_i\cdot \tilde{\boldsymbol{k}}_j\right)\boldsymbol{v}_j}{\sum_{j = 1}^i\exp\left(\boldsymbol{q}_i\cdot \tilde{\boldsymbol{k}}_j\right)},\qquad \tilde{\boldsymbol{k}}_j = \frac{\boldsymbol{k}_j}{\Vert\boldsymbol{k}_j\Vert} \end{align}

This naturally raised the question of which of these two variants is better than the standard Equation $\eqref{eq:sdpa}$. For convenience, we refer to them respectively as "Query-Normalized Dot-Product Attention" and "Key-Normalized Dot-Product Attention", abbreviated "QNA" and "KNA".

Furthermore, since we can normalize either the Query or the Key, it's natural to also consider normalizing both at once. So we additionally experimented with the following "Scaled Cosine Attention (CosA)":

\begin{equation}\boldsymbol{o}_i = \frac{\sum_{j = 1}^i\exp\left(\lambda\,\tilde{\boldsymbol{q}}_i\cdot \tilde{\boldsymbol{k}}_j\right)\boldsymbol{v}_j}{\sum_{j = 1}^i\exp\left(\lambda\,\tilde{\boldsymbol{q}}_i\cdot \tilde{\boldsymbol{k}}_j\right)} = \frac{\sum_{j = 1}^i\exp\left(\lambda\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)\right)\boldsymbol{v}_j}{\sum_{j = 1}^i\exp\left(\lambda\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)\right)} \end{equation}

Here $\lambda$ uses the result from Viewing Attention's Scale Operation Through Gradient Maximization, namely $\lambda = 4\log n$ (the original post used 3.5, but since the training length here is relatively small, 4 is a more accurate choice), where $n$ is fixed at half the training length, or dynamically set to the position id plus one.

Results First

Following the previous experimental setup for length extrapolation, we use 100M-parameter small models with the GAU architecture, trained for the same number of steps (due to limited time, the models are actually not yet fully trained at this step count), with a training length of 512, extrapolated to a length of 4096. The results are shown in the table below. Here "Baseline" refers to Equation $\eqref{eq:sdpa}$, and $\text{-}\log n$ denotes adding the length-dependent scaling factor introduced in Viewing Attention's Scale Operation Through Entropy Invariance. The evaluation metric is per-token accuracy of the language model — higher is better.

$$\begin{array}{c|cc} \hline \text{test length} & 512(\text{training}) & 4096(\text{repeat}) & 4096(\text{no repeat}) \\ \hline \text{Baseline} & 49.41\% & 24.17\% & 23.16\% \\ \text{Baseline-}\log n & 49.40\% & 24.60\% & 24.02\% \\ \hline \text{QNA} & 49.55\% & 22.45\% & 22.18\% \\ \text{QNA-}\log n & 49.42\% & 19.55\% & 18.74\% \\ \text{KNA} & 49.60\% & 61.08\% & 47.69\% \\ \text{KNA-}\log n & 49.58\% & 63.17\% & 46.40\%\\ \text{CosA} & 49.73\% & 58.90\% & 46.98\% \\ \text{CosA-}\log n & 49.67\% & 64.74\% & 48.95\% \\ \hline \end{array}$$

From the table we can see: 1) Both QueryNorm and KeyNorm achieve better results within the training length, although this advantage is very slight and would likely become negligible with further training — still, this advantage is very consistent, hinting at the possibility of making training more stable; 2) KeyNorm brings a very noticeable improvement to length extrapolation — this is the "unexpected bonus" from the experimental results!

Note that, unlike NTK-RoPE and YaRN, which require modifying the model at inference time, the length extrapolation achieved here by KNA and CosA requires absolutely no modification at inference time. Some readers might therefore wonder: since KNA and CosA already extrapolate so well without any inference-time changes, would combining them with extrapolation tricks like NTK-RoPE or YaRN push things "to the next level"? I tested this too, and the results are shown below:

$$\begin{array}{c|cc} \hline \text{test length} & 512(\text{training}) & 4096(\text{repeat}) & 4096(\text{no repeat}) \\ \hline \text{Baseline} & 49.41\% & 24.17\% & 23.16\% \\ \text{Baseline-NTK} & 49.41\% & 60.57\% & 42.20\% \\ \text{Baseline-YaRN} & 49.41\% & 80.10\% & 47.45\% \\ \text{Baseline-ReRoPE} & 49.41\% & 76.11\% & 47.82\% \\ \hline \text{Baseline-}\log n & 49.40\% & 24.60\% & 24.02\% \\ \text{Baseline-}\log n\text{-NTK} & 49.40\% & 75.86\% & 47.06\% \\ \text{Baseline-}\log n\text{-YaRN} & 49.40\% & 82.57\% & 46.52\% \\ \text{Baseline-}\log n\text{-ReRoPE} & 49.40\% & 85.47\% & 48.87\% \\ \hline \text{QNA} & 49.55\% & 22.45\% & 22.18\% \\ \text{QNA-NTK} & 49.55\% & 52.28\% & 39.88\% \\ \text{QNA-YaRN} & 49.55\% & 82.53\% & 47.50\% \\ \text{QNA-ReRoPE} & 49.55\% & 78.22\% & 47.72\% \\ \hline \text{QNA-}\log n & 49.42\% & 19.55\% & 18.74\% \\ \text{QNA-}\log n\text{-NTK} & 49.42\% & 57.44\% & 41.56\% \\ \text{QNA-}\log n\text{-YaRN} & 49.42\% & 80.08\% & 45.16\% \\ \text{QNA-}\log n\text{-ReRoPE} & 49.42\% & 84.71\% & 48.31\% \\ \hline \text{KNA} & 49.60\% & 61.08\% & 47.69\% \\ \text{KNA-NTK} & 49.60\% & 64.44\% & 43.02\% \\ \text{KNA-YaRN} & 49.60\% & 84.19\% & 47.44\% \\ \text{KNA-ReRoPE} & 49.60\% & 77.76\% & 47.73\% \\ \hline \text{KNA-}\log n & 49.58\% & 63.17\% & 46.40\%\\ \text{KNA-}\log n\text{-NTK} & 49.58\% & 79.05\% & 47.43\%\\ \text{KNA-}\log n\text{-YaRN} & 49.58\% & 83.95\% & 47.16\%\\ \text{KNA-}\log n\text{-ReRoPE} & 49.58\% & 85.48\% & 48.78\%\\ \hline \text{CosA} & 49.73\% & 58.90\% & 46.98\% \\ \text{CosA-NTK} & 49.73\% & 62.50\% & 42.77\% \\ \text{CosA-YaRN} & 49.73\% & 83.40\% & 47.80\% \\ \text{CosA-ReRoPE} & 49.73\% & 77.82\% & 47.80\% \\ \hline \text{CosA-}\log n & 49.67\% & 64.74\% & 48.39\% \\ \text{CosA-}\log n\text{-NTK} & 49.67\% & 78.97\% & 47.46\% \\ \text{CosA-}\log n\text{-YaRN} & 49.67\% & 82.28\% & 45.72\% \\ \text{CosA-}\log n\text{-ReRoPE} & 49.67\% & 85.67\% & 48.39\% \\ \hline \end{array}$$

This table is a bit verbose, mainly intended to give readers a comprehensive sense of the differences among mainstream length-extrapolation techniques — feel free to compare along whichever dimension interests you. But note that if you're looking at length-extrapolation performance specifically, you should focus mainly on the "non-repeating" column, with the "repeating" column as secondary reference. From the table above, the results are honestly a bit surprising: KeyNorm seems to be "immune" to existing RoPE-based extrapolation tricks — stacking NTK, YaRN, etc. on top of it doesn't bring a clear improvement, and might even make things worse. That said, overall the "repeating" column does show a significant improvement, while the "non-repeating" column doesn't. These results suggest that KeyNorm still cannot effectively recognize positions beyond the training length (hence the mediocre "repeating" results), but it does effectively avoid the PPL-explosion problem (hence the decent "non-repeating" results).

This may be good news for those working on long-context modeling: on one hand, unlike ALIBI, KERPLE, and similar methods, KeyNorm's length extrapolation doesn't require adding a local constraint, and no modification whatsoever is needed after training — it's purely a "free lunch," and it even seems that adding KeyNorm actually improves training performance. On the other hand, precisely because it's non-local, one can continue training on longer texts, and when continuing training there's no longer any need to agonize over whether to choose PI or ABF — with KeyNorm, you don't need to change anything at all.

Analysis

Even though this was an accidental discovery, we should still try to explain it — otherwise it will forever remain just an accident. So in this section, let's try to think through why KeyNorm helps with length extrapolation.

Let's return to Equation $\eqref{eq:sdpa}$: the correlation score between the $i$-th token and the $j$-th token is computed via an inner product:

\begin{equation}s(j|i) = \boldsymbol{q}_i\cdot \boldsymbol{k}_j = \Vert\boldsymbol{q}_i\Vert \Vert\boldsymbol{k}_j\Vert \cos(\boldsymbol{q}_i,\boldsymbol{k}_j),\quad p(j|i) = \frac{\exp\left(\frac{s(j|i)}{\sqrt{d}}\right)}{\sum_{j=1}^i \exp\left(\frac{s(j|i)}{\sqrt{d}}\right)}\end{equation}

In the second equality, from a geometric point of view, we decompose it into the product of the respective norms and the cosine of the angle between them. Attention $p(j|i)$ is a conditional probability; $\Vert\boldsymbol{q}_i\Vert$ only depends on the current position $i$, so it doesn't change the relative magnitude of attention, only the sparsity. $\Vert\boldsymbol{k}_j\Vert$, on the other hand, is capable of changing the relative magnitude of $p(j|i)$, but doesn't involve interaction with $i,j$ — it can be used to express certain absolute signals; for example, Scissorhands shows that tokens at certain absolute positions consistently receive high attention, which could plausibly be expressed via $\Vert\boldsymbol{k}_j\Vert$. The remaining term, $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$, is what expresses the interaction between $i,j$, and it's the term with the greatest degree of freedom.

Clearly, to increase the relative importance of some position $j$, the model has two options: 1) increase the norm $\Vert\boldsymbol{k}_j\Vert$; or 2) increase $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$, i.e., shrink the angle between $\boldsymbol{q}_i,\boldsymbol{k}_j$. However, because of the "curse of dimensionality," significantly changing an angle in high-dimensional space is relatively difficult, so if the goal can be achieved simply by increasing the norm $\Vert\boldsymbol{k}_j\Vert$, the model will preferentially do so by increasing the norm $\Vert\boldsymbol{k}_j\Vert$. The direct consequence of this is that training of $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$ may end up being insufficient.

Here I'll make an assertion (a conjecture):

Insufficient training of $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$ is the main reason attention fails to extrapolate to longer lengths.

By "insufficient training" of $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$, I mean that the angles of $\boldsymbol{q}_i,\boldsymbol{k}_j$ that get trained only form a limited set, whereas during length extrapolation the model must face a much larger set, and thus fails to make correct predictions. If you carefully think through the derivation in the YaRN post, you'll notice that the reason NTK and YaRN work is precisely because they modify how RoPE is implemented at inference time, so that the angles of $\boldsymbol{q}_i,\boldsymbol{k}_j$ fall back into the limited set seen during training, avoiding the larger, unseen set — turning extrapolation into interpolation. ReRoPE is even more direct: it simply truncates relative positions beyond the window, so that none of the positional encodings encountered at inference time are "unfamiliar." These techniques all indirectly corroborate this assertion, to some extent.

Starting from this assertion, the cause of KeyNorm's length extrapolation becomes simple to understand. Whether it's KNA (which only applies KeyNorm) or CosA (which applies both QueryNorm and KeyNorm), both exclude $\Vert\boldsymbol{k}_j\Vert$ from the definition of attention. So in order to change the relative importance of $j$, the model has only one option left: "adjust $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$." This forces the model to train and utilize $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$ more fully, which in turn indirectly promotes length extrapolation. In addition, I also experimented with the combination "KeyNorm + NoPE," but found no length-extrapolation benefit there, which indicates that RoPE also plays an important role in KeyNorm's length extrapolation. This actually makes sense: RoPE rotates $\boldsymbol{q}_i,\boldsymbol{k}_j$, which helps expand the range of $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$ seen during training, thereby making the training of $\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$ more thorough.

Has anyone already tried QueryNorm and KeyNorm before? Yes. The 2020 paper Query-Key Normalization for Transformers experimented with CosA and also proposed a similar length-log-based scale factor, but didn't discuss the length-extrapolation issue. In addition, a paper from Google earlier this year, Scaling Vision Transformers to 22 Billion Parameters, also added normalization to the Query and Key — but it used LayerNorm, and both LayerNorm and RMSNorm come with a learnable gamma parameter, which means the norm of the normalized vector isn't necessarily constant. So it's not clear whether it would achieve the same length-extrapolation effect as in this post.

Summary

This post introduced a length-extrapolation scheme I discovered by accident, "KeyNorm" — applying L2 normalization to the Key sequence in attention — which achieves better performance within the training length and shows a marked improvement in length extrapolation. It belongs to the "pre-hoc modification" category, and compared to other pre-hoc methods like ALIBI and KERPLE, it has no local constraint, so it holds more promise for scaling up. Compared to "post-hoc modification" approaches like NTK-RoPE and YaRN, it doesn't sacrifice in-training-length performance when extrapolating.