ALiBi: Attention with Linear Biases

The Problem

Training Language Models on long context windows (e.g., 8K or 32K tokens) is extraordinarily expensive because standard attention scales quadratically ($O(N^2)$) with sequence length. A common desire was to "train short" (train on 2K tokens to save compute) and "test long" (deploy the model on 8K token inputs).

Unfortunately, models using standard positional encodings (like sinusoidal or learned absolute embeddings) utterly fail when presented with positions they have never seen during training. Even Rotary Position Embedding (RoPE) initially struggled to extrapolate significantly beyond its training length without complex fine-tuning tricks.

The Idea

The authors completely discarded the idea of adding positional information to the token embeddings at the bottom of the network. Instead, they proposed modifying the attention mechanism itself by biasing the attention scores right before the softmax layer.

Their method, ALiBi (Attention with Linear Biases), simply subtracts a penalty from the attention score based on the distance between the two tokens. The further apart the tokens are, the higher the penalty. The penalty grows linearly.

How It Works

In standard attention, you compute $\text{Softmax}(Q K^T)$.

With ALiBi, this becomes $\text{Softmax}(Q K^T - m \cdot D)$, where $D$ is a matrix of distances between tokens, and $m$ is a head-specific scalar (a slope).

For example, if token A is 5 positions away from token B, its attention score is penalized by $5m$. Different attention heads are assigned different, fixed slopes. Some heads receive steep slopes, strongly penalizing distant tokens and forcing the head to focus locally. Other heads receive gentle slopes, allowing them to attend to the broader context. Crucially, these slopes $m$ are mathematically fixed as a geometric sequence across heads and are not learned during training.

Because the penalty is based purely on relative distance, the model never sees an "unseen absolute position" at inference time. It simply applies relative distance penalties to tokens, allowing it to extrapolate gracefully.

Why It Mattered

ALiBi proved highly effective at length extrapolation out-of-the-box. A model trained on 1,024 tokens could successfully process 2,048 or even 4,096 tokens at inference time with minimal performance degradation.

It was also incredibly simple to implement and slightly faster than RoPE, as it required no complex rotations or modifications to the embeddings themselves. It became a popular choice for building open-weight models designed specifically for long-context tasks.

What Came After

ALiBi was widely adopted by prominent models, most notably MPT (MosaicML) and BloombergGPT.

Over time, despite its elegance for extrapolation, it has largely been superseded by RoPE. Researchers found that RoPE generally yields slightly lower perplexity on the training distribution, and the extrapolation problem for RoPE was subsequently solved using fine-tuning tricks like RoPE scaling (e.g., YaRN, position interpolation). This allowed modern models like Llama 3 to achieve 128K context windows while sticking to Rotary Position Embedding.