ALiBi 位置偏差：带有线性偏差的注意力机制

ALiBi, short for Attention with Linear Biases, is a method for giving a transformer language model a sense of word order. A transformer's attention layer compares every token with every other token, but the comparison itself ignores where tokens sit in the sequence, so models need some extra signal about position. ALiBi supplies that signal in the simplest possible way: when one token attends to another, the attention score is reduced by an amount that grows in a straight line with the distance between them. Nearby tokens are penalized a little, distant tokens a lot, and no position vectors are added to the token embeddings at all.

The mechanism works step by step as follows. First, the model computes the usual attention score for a query token at position i and a key token at position j, which is the dot product of the query and key vectors. Second, ALiBi adds a bias equal to a head-specific slope m multiplied by the negative distance, so the key one step back gets minus m, two steps back gets minus 2m, and so on, while the token attending to itself gets zero. Third, the softmax turns the adjusted scores into attention weights as normal. Because the slope differs between attention heads, some heads fall off steeply and focus on the last few tokens while others decay slowly and look much further back. The slopes are fixed before training rather than learned, and the bias is applied to query-key scores in every layer but never to the value vectors.

ALiBi was introduced by Ofir Press, Noah A. Smith and Mike Lewis in the paper Train Short, Test Long, first posted to arXiv on August 27, 2021 and revised in April 2022. The motivating question was extrapolation: can a model trained on sequences of one length make good predictions on longer sequences at inference time? The paper showed that the standard sinusoidal position embeddings from the original transformer break down almost immediately past the training length. On WikiText-103, a sinusoidal model trained on 512 tokens rose from 20.05 perplexity at 512 tokens to 132.41 at 2,512 tokens, while an ALiBi model trained on 512 tokens reached 18.40 at 3,072 tokens, better than a sinusoidal model trained at the full 3,072. Perplexity is a measure of how surprised a model is by the next token, so lower is better.

The practical selling point was cost. Training on short sequences is cheaper because attention cost grows with sequence length. The paper's larger experiment trained 1.3-billion-parameter models on 461 GB of CC100 and RoBERTa text using 128 V100 GPUs, and reports that an ALiBi model trained on 1,024-token inputs matched a sinusoidal model trained on 2,048 tokens while training 11% faster and using 11% less memory. The authors also report the method adds almost nothing at runtime: on one V100, training throughput and memory at 1,024 tokens were within about 1% of the sinusoidal baseline, and the per-head bias mask costs up to roughly 100 MB.

Adoption followed quickly. BigScience's BLOOM, the 176-billion-parameter open multilingual model whose training began on March 11, 2022 on 384 A100 GPUs, lists ALiBi as its positional encoding in its official model card. MosaicML's MPT-7B, released on May 5, 2023 and pretrained on one trillion tokens, also dropped position embeddings for ALiBi. Its StoryWriter variant was fine-tuned with a 65,000-token context, and MosaicML reported it handled about 68,000 tokens of The Great Gatsby and produced generations up to 84,000 tokens, which it attributed to ALiBi's extrapolation.

The obvious alternatives are absolute embeddings, relative biases and rotary embeddings. Absolute embeddings, sinusoidal or learned, attach a vector for each position number to the token, which ties the model to the positions it saw in training. The T5 relative bias learns a bias per bucket of distances; in the ALiBi paper it extrapolated better than sinusoidal and was the strongest baseline, but trained at least twice as slowly. Rotary Position Embedding (RoPE), proposed by Jianlin Su and colleagues in April 2021, rotates query and key vectors by an angle set by position so that scores depend on relative distance; it extrapolated better than sinusoidal in the ALiBi comparison but not as well as ALiBi. RoPE is the scheme used by models such as LLaMA, and methods such as Position Interpolation, posted by Chen and colleagues in June 2023, extend those RoPE models to 32,768 tokens with fewer than 1,000 fine-tuning steps.

ALiBi has real limits, several stated by its own authors. Their sliding-window analysis found that much of the gain at long evaluation lengths comes from reducing the early token curse, the poor predictions a model makes near the start of a chunk where it has little preceding context, rather than from actually using information beyond the training length. In other words, ALiBi lets a model run on longer inputs without collapsing, but it may not read the extra text in a useful way. On the 1.3B experiments, ALiBi's perplexity was best near twice the training length and only stayed usable to around 10,000 tokens. The slopes were picked from about ten manual trials on one dataset, learned slopes extrapolated poorly, and most experiments used a 247-million-parameter model. A 2023 NeurIPS study by Kazemnejad and colleagues found that on downstream reasoning-style tasks, ALiBi, rotary and absolute embeddings all generalized poorly to longer lengths, and a decoder with no positional encoding did better.

In practice ALiBi requires very little. It adds no trainable parameters, only a fixed matrix of distance penalties per head that can be precomputed, and the slope formula is a single line of code. The main engineering requirement is an attention kernel that accepts an additive bias, which matters because fused kernels such as FlashAttention need explicit support for it; MPT shipped with FlashAttention and ALiBi together. The method suits teams that want cheaper short-context pretraining with graceful behavior on longer inputs, while teams that need a model to retrieve facts from far back in a very long document should test long-context recall directly rather than rely on low perplexity alone.