MPT’s ALiBi Context Extension: No Position Fine-Tuning Needed

Most models cannot be run past the sequence length they were trained on without either degrading badly or being fine-tuned to extend. MPT can, and the reason is a single design choice: it has no position embeddings at all. In their place is a fixed penalty on attention scores that grows with distance, and that penalty is defined for distances the model never saw.

Why a trained window is normally a hard edge

A transformer’s attention is, by itself, order-blind. Shuffle the tokens and the mechanism computes the same thing, because it is a weighted average over a set. Position information has to be injected, and the two dominant ways of injecting it both have a length baked in.

Learned absolute embeddings add a vector per position, looked up from a table with one row per position. There is no row 4,097 in a table of 4,096. The failure is total and immediate.

Rotary embeddings (RoPE) rotate the query and key vectors by an angle proportional to position. This is defined for any position — there is no table — so it does not crash. But the model has only ever seen the rotation frequencies that occur within its training length, and beyond that it is being asked to interpret angle patterns that never appeared in training. In practice output degrades rather than errors, which is why RoPE models are extended by explicit rescaling schemes rather than by simply asking for more, and why Llama 3.1’s RoPE scaling is a documented procedure with its own parameters rather than a configuration flag you can safely raise.

What ALiBi adds to the attention scores

ALiBi — Attention with Linear Biases — was introduced by Press, Smith and Lewis in the paper “Train Short, Test Long”, published on arXiv in 2021 and presented at ICLR 2022. It removes position embeddings entirely and modifies the attention score directly.

Ordinary causal attention computes a score for query at position i against key at position j, then softmaxes over j. ALiBi subtracts a penalty before the softmax:

score(i, j) = (q_i · k_j) / sqrt(d)  −  m * (i − j)

where  i − j  is how far back the key is,
       m      is a fixed per-head slope, not learned.

The slopes m form a geometric sequence across heads.
For 8 heads:  1/2, 1/4, 1/8, 1/16, 1/32, 1/64, 1/128, 1/256

Read what that does. The penalty is linear in distance, so a token twenty positions back is penalised twice as much as one ten positions back. It is applied per head with a different slope, so a head with a large slope is pushed hard toward recent tokens while a head with a tiny slope is barely constrained and can attend far back. The model gets a spread of effective receptive fields for free, and none of it is learned — the slopes are set by a formula, not by training.

Why that extrapolates

This is the whole argument, and it is short. The bias term −m * (i − j) is a function evaluated at run time. There is no table to index and no periodic pattern to leave the domain of. At a distance of 60,000 it computes a value in exactly the same way it does at a distance of 60, and that value is a smooth continuation of the ones the model trained on.

Contrast with the two alternatives. A learned table has no entry to return. RoPE returns something — the rotation is defined everywhere — but what it returns at long distances is a pattern of relative angles the model has no learned response to. ALiBi returns a number on the same monotone line the model has been fitted against throughout training. The model does not have to generalise to a new kind of input; it has to handle more of the same input.

There is a second, subtler part of the argument. The bias depends only on the distance i − j and never on absolute position, so the model is never told where in the sequence it is — only how far back something is. A relative-only signal has nothing to run out of. Absolute schemes encode “this is position 3,000”, and position 60,000 is a statement the model has never seen made. ALiBi makes no statement about position for the model to be wrong about.

The paper’s claim, and the reason for its title, is that a model trained at a short sequence length evaluates with lower perplexity at longer lengths than sinusoidal or rotary baselines do. That is a published result from the authors, measured on their setup; it is not a guarantee about any particular downstream task.

What MosaicML did with it in MPT

MosaicML released MPT-7B in May 2023, trained with a 2,048-token sequence length and using ALiBi rather than position embeddings. The release explicitly used the extrapolation property as a feature: because there is no position table, the same weights can be run with a larger configured maximum sequence length without changing the model.

The demonstration was MPT-7B-StoryWriter-65k+, a variant fine-tuned on long fiction with a 65,536-token sequence length, which MosaicML documented as able to run beyond even that on a single 80GB node — the “+” in the name is doing that work. MPT-30B, released in June 2023, was trained with an 8,192-token sequence length from the start. All the MPT models were released under Apache 2.0, with the instruct and chat variants carrying the terms of their fine-tuning datasets.

In practice extension is a runtime setting rather than a surgery on the weights:

from transformers import AutoConfig, AutoModelForCausalLM

name = "mosaicml/mpt-7b"
config = AutoConfig.from_pretrained(name, trust_remote_code=True)
config.max_seq_len = 16384          # trained at 2048

model = AutoModelForCausalLM.from_pretrained(
    name, config=config, trust_remote_code=True
)

There is no corresponding move for a RoPE model. Raising the equivalent field on a Llama or Yi checkpoint does not extend it, it just removes the guard rail — see why Yi ships separate 200K checkpoints rather than a config flag.

MosaicML was acquired by Databricks in 2023, and the MPT line was not continued — Databricks’ subsequent open release was DBRX, which uses rotary embeddings like everything else. That is a fair summary of where ALiBi ended up: a genuinely elegant answer to one problem, adopted by few of the models that followed, and worth understanding anyway because it is the clearest illustration of what a position encoding is actually for. Once you can say why ALiBi extrapolates and RoPE does not, the whole category of long-context extension techniques stops being a list of names.

Where the property stops holding

• Memory does not extrapolate. ALiBi solves the positional problem and nothing else. The KV cache still grows linearly with sequence length and attention is still quadratic in it. A model that can accept 65K tokens still needs the hardware to hold 65K tokens of cache, which is the constraint that hybrid stacks like Jamba’s attack instead.
• The bias is a recency prior, permanently. The penalty never switches off, so a token 50,000 positions back is heavily down-weighted in every head with a meaningful slope. For tasks that need precise recall of one distant fact, this is a real headwind, and it is the main reason ALiBi did not become the default despite the extrapolation property being genuinely useful.
• Quality is not flat across the extended range. Lower perplexity than a baseline at long lengths is not the same as unchanged capability. Treat an extended window as something to evaluate on your task at the depth you intend to use, not as capacity you have been granted.
• Tooling support is uneven. MPT ships custom modelling code, which is why the snippet above needs trust_remote_code=True. Optimised attention kernels have to implement the bias explicitly, and not every serving stack does.