The Surprising Role of the Bias Term: RoPE + Bias = Better Length Generalization

Bias项的神奇作用：RoPE + Bias = 更好的长度外推性 苏剑林 · 2023-04-03 · Machine Learning language model attention positional encoding extrapolation RoPE

Never would I have guessed that the bias term could be tied to the length generalization ability of Transformers!

Length generalization (extrapolation) is a desirable property we want Transformers to have; I have systematically discussed this issue before in Transformer Upgrade Path: 7. Length Generalization and Local Attention and Transformer Upgrade Path: 8. Length Generalization and Positional Robustness. As for the bias term, the current mainstream view is that once a model is large enough, the bias term doesn't contribute anything particularly useful, so many models choose to drop it — notable examples being Google's T5 and PaLM, and our own RoFormerV2 and GAU-α followed suit.

So how exactly are these two seemingly unrelated things connected? Can the bias term really enhance a Transformer's length generalization? Let me walk through it step by step. more

A Hidden Easter Egg

First, why did I even think to look into the connection between the bias term and length generalization? It's because a few days ago, while revisiting the GAU paper Transformer Quality in Linear Time, I noticed a "hidden Easter egg" that I hadn't paid attention to before — an additive relative positional encoding. Its pseudocode looks like this:

Pseudocode of GAU's additive relative positional encoding

Here we mainly care about the $n\geq 512$ part. Written as a formula, it's roughly:

\begin{equation}\boldsymbol{q}_m^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n \quad\to\quad \boldsymbol{q}_m^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n+ \boldsymbol{a}^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{b}\label{eq:rel-bias}\end{equation}

where $\boldsymbol{\mathcal{R}}_m,\boldsymbol{\mathcal{R}}_n$ is RoPE's rotation matrix and $\boldsymbol{a},\boldsymbol{b}$ are two learnable parameters.

I'd noticed this additive relative positional encoding before, but at the time my only reaction was "I don't understand why they're using several kinds of positional encoding at once." Recently, however, I've been thinking a lot about length generalization, so this form suddenly caught my attention. One can show that when $\boldsymbol{a}=\boldsymbol{b}=[\sqrt{\lambda},0,\sqrt{\lambda},0,\cdots,\sqrt{\lambda},0]^{\top}$, the result is exactly the Sandwich encoding described in Transformer Upgrade Path: 7. Length Generalization and Local Attention, which improves length generalization. Its underlying principle is that $\boldsymbol{a}^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{b}$ decreases with $|m-n|$; once added to the attention matrix, it has a localizing effect on attention, and according to Transformer Upgrade Path: 7. Length Generalization and Local Attention, localized attention is the key to a language model's extrapolation ability.

So I couldn't help but wonder: could this additive relative positional encoding in the original paper have been designed precisely to enhance length generalization? Did the authors of GAU really have such foresight, proposing an idea similar to Sandwich even before Sandwich itself came out?

Switching to a Bias Term

That said, to me, any scheme that improves length generalization by adding an extra term onto the attention matrix feels somewhat inelegant. So regardless of the original authors' intent or the actual effectiveness of their approach, I'm not inclined to go that route. Is there a similar but almost "unnoticeable" alternative? It occurred to me that if $\boldsymbol{a}$ and $\boldsymbol{b}$ are simply the bias terms of $\boldsymbol{q}_m,\boldsymbol{k}_n$, this might achieve a similar effect. That is, consider:

\begin{equation}\boldsymbol{q}_m^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n \quad\to\quad (\boldsymbol{q}_m + \boldsymbol{a})^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n(\boldsymbol{k}_n + \boldsymbol{b})\end{equation}

Clearly, simply adding a bias term is, in both form and computational cost, almost imperceptible. If this alone could enhance length generalization, it would be a remarkably elegant solution. Is it actually feasible? Let's first look at the expanded result:

\begin{equation}\boldsymbol{q}_m^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n + \boldsymbol{a}^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n + \boldsymbol{q}_m^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{b} + \boldsymbol{a}^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{b} \label{eq:bias}\end{equation}

The first and fourth terms here correspond exactly to formula $\eqref{eq:rel-bias}$, which is what we want. So we want to examine what role the second and third terms play — if they don't have much of a noticeable effect, then simply adding a bias term is, at the very least, "promising" for achieving an extrapolation effect similar to formula $\eqref{eq:rel-bias}$ or Sandwich.

Here's my reasoning: as the Query and Key of attention, $\boldsymbol{q}_m$ and $\boldsymbol{k}_n$ should be fairly "isotropic" — that is, their directions should be roughly uniform, close to uniform sampling on a sphere. Meanwhile $\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n=\boldsymbol{\mathcal{R}}_{n-m}$ is just an orthogonal transformation, which doesn't change the isotropy of $\boldsymbol{q}_m$ and $\boldsymbol{k}_n$. So the two terms $\boldsymbol{a}^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{k}_n $ and $\boldsymbol{q}_m^{\top}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n\boldsymbol{b}$ amount to inner products between vectors sampled from an isotropic distribution and a fixed vector. Based on our discussion in The Distribution of the Angle Between Two Random Vectors in n-Dimensional Space, the angle between such vectors should be very close to 90 degrees — in other words, the expected value of this inner product should be 0. So in theory, the second and third terms should have a much weaker effect than the remaining two.

Of course, this is just a conjecture; how things actually turn out during training can only be determined experimentally. So without further delay, I ran the experiment.

Experimental Results

This time I chose a language modeling task for the experiment, using the same GAU-α architecture as before, with a training length and batch size of 512, and the Tiger optimizer. The only difference between the two models is whether the bias for Q and K is enabled (all other biases are still removed).

Comparison of extrapolation performance:

$$\begin{array}{c} \text{LM accuracy at different test lengths} \\ {\begin{array}{c|cccc} \hline & 512 & 1024 & 2048 & 4096 \\ \hline \text{w/o Bias} & 52.37\% & 33.15\% & 22.85\% & 17.87\% \\ \text{w/ Bias} & 52.75\% & 50.99\% & 45.25\% & 39.55\% \\ \hline \end{array}} \end{array}$$

As you can see, the bias term barely affects performance at the training length (512), but it opens up a clear gap in length generalization. Who would have thought this seemingly negligible bias term could have such a magical effect! Of course, if we rerun the experiment several times, the extrapolation results might fluctuate noticeably — after all, length generalization is a kind of "bonus feature" rather than something we're actively optimizing for.

To verify whether the remaining mechanism works the way we suspected, I visualized how the four terms in formula $\eqref{eq:bias}$ change for a particular sample at a particular layer:

Comparison of the four inner-product terms after adding bias

As you can see, the 4th term does indeed show a decaying trend, and it dominates in magnitude. Summing up these four terms and comparing against the model without bias gives:

Comparison of attention matrices with and without bias

For the model without bias (blue), attention does show a decaying trend within the training length (512), but it rises again once the length increases beyond that, with no clear locality — this is exactly why its extrapolation isn't very good. In contrast, and consistent with our earlier conjecture, the model with the bias term (orange) shows a much more pronounced decaying trend in its attention matrix, meaning it has a stronger localization effect, and thus better extrapolation performance. It's worth noting that not every layer of the bias-augmented model shows such a pronounced decay — generally, earlier layers show a more obvious decaying trend, while later layers show a weaker one, suggesting that layers closer to the input focus more on local information. This is consistent with the conclusions of The Devil in Linear Transformer.

[Note: after repeated testing, we found that the length generalization results in this article are not very reproducible (possibly closely related to model architecture, hyperparameters, etc.). Please use this conclusion with caution.]

Further Thoughts

This raises a question: haven't previous works on length generalization already shown that RoPE doesn't extrapolate all that well? Could it be that none of them used bias terms? To check this, I went back and looked into it, and as expected: the "pioneering work" ALIBI and the more recent XPOS indeed do not include bias terms, whereas KERPLE and Sandwich do. When reading these papers earlier, I always had the feeling that RoPE's extrapolation performance in KERPLE and Sandwich seemed better than in ALIBI and XPOS — now I can confirm that this was probably not just my imagination. Since both KERPLE and Sandwich include bias terms, based on the conclusion of this article, RoPE is indeed capable of showing better length generalization.

Some readers might recall that I previously said the bias term for the attention Key can be removed — does that contradict what's said here? On this question, you can refer to the Zhihu discussion Why Don't the Keys in Some Vision Transformers Need a Bias?. In fact, the conclusion that "the Key's bias can be removed" applies specifically to attention without RoPE, where, thanks to softmax, an added bias can be cancelled out:

\begin{equation}\frac{e^{\boldsymbol{q}\cdot(\boldsymbol{k}_n + \boldsymbol{b})}}{\sum\limits_n e^{\boldsymbol{q}\cdot(\boldsymbol{k}_n + \boldsymbol{b})}} = \frac{e^{\boldsymbol{q}\cdot\boldsymbol{k}_n}e^{\boldsymbol{q}\cdot\boldsymbol{b}}}{\sum\limits_n e^{\boldsymbol{q}\cdot\boldsymbol{k}_n} e^{\boldsymbol{q}\cdot\boldsymbol{b}}}= \frac{e^{\boldsymbol{q}\cdot\boldsymbol{k}_n}}{\sum\limits_n e^{\boldsymbol{q}\cdot\boldsymbol{k}_n}}\end{equation}

However, this "cancelling out" relies on $\boldsymbol{b}$ being independent of $n$. But from formula $\eqref{eq:bias}$ we know that after applying RoPE, $\boldsymbol{b}$ actually becomes a function of $m,n$ as well, and in practice cannot be cancelled out. Therefore, for models that use RoPE, removing the bias term does lead to a different outcome than keeping it.

There's one more question: why bother going to such lengths to pursue length generalization? Wouldn't it be simpler to just fine-tune the model directly on longer samples? Actually, even for readers who hold this view, length generalization still has its benefits. Setting compute aside, better length generalization means a smaller gap between fine-tuning and pretraining, which makes fine-tuning less prone to catastrophic forgetting — something that matters even more for today's LLMs. And to take this a step further, the ideal outcome would be: a model trained on short texts that can seamlessly switch to long-text scenarios without any loss of performance — or even with improved performance.

Summary

In this post, I shared a rather surprising finding: the bias term can enhance the length generalization of RoPE-based models! A term that seems to have almost no presence at all turns out to be tied to a Transformer's length generalization ability — a reminder of just how important the details can be. Sometimes it's the smallest, most overlooked pieces that end up playing a decisive role.