Train Short, Test Long

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) struggles 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. Instead, they proposed modifying the attention scores directly after the Query-Key dot product, but before the softmax.

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)$.

In ALiBi, you compute $\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 (e.g., steep slopes for heads that focus on local context, gentle slopes for heads that look at the whole sequence).

Because the penalty is based purely on relative distance, the model never sees an "unseen absolute position" at inference time. It just sees relative distances.

Why It Mattered

ALiBi proved highly effective at length extrapolation. A model trained on 1024 tokens could successfully process 2048 or even 4096 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 models designed specifically for long-context tasks.

What Came After

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

However, despite its elegance for extrapolation, it has largely been superseded by RoPE combined with advanced techniques like YARN or position interpolation. Researchers found that RoPE generally yields slightly lower perplexity on the training distribution, and the extrapolation problem with RoPE can be solved by modifying the rotary base frequency during fine-tuning (e.g., RoPE scaling), allowing models like Llama 3 to achieve 128K context windows without needing ALiBi.