Skip to main content

ALiBi: Attention with Linear Biases

Follow a query-key distance through ALiBi head slopes, attention logits and normalized weights, then understand the evidence and limits of length extrapolation.

Attention with Linear Biases (ALiBi) adds a head-specific distance penalty to attention logits before softmax. It supplies position information without a learned position embedding table. The query, key and value projections remain learned.

The penalty favors nearby keys when content scores are equal. A distant key with a strong content score can still receive more attention than a nearby one. Start with one query and keep that distinction visible through the calculation.

Follow one query-key pair

The query starts at position 5, “mat”, and the selected key is position 0, “The”. Their distance is five tokens. Head 0 of a four-head example has slope 0.25, so the distance bias is −1.25. Adding that to the illustrative content logit of 2 gives a biased logit of 0.75.

One query, one head, every key
Heads

Teaching logits, held fixed across queries and heads. Both columns use the same causal mask. Select a key to follow its calculation.

Query 5 “mat” → key 0 “The”

Head slope m
0.25
Distance |i − j|
5 tokens
Content logit
2.000
Bias −m × distance
-1.250
Biased logit
2.000 + (-1.250) = 0.750
No-bias weight
35.9%
ALiBi weight
20.0%
Inspect all key logits and weights

Each weight is normalized over all keys allowed for this query. Teal bars show ALiBi; gray bars show the same logits without distance bias.

All six key calculations for the selected query and head
KeyContentDistanceBiasBiased logitNo biasALiBi
0 · The2.0005-1.2500.750
35.9%
20.0%
1 · cat0.5004-1.000-0.500
8.0%
5.7%
2 · sat1.5003-0.7500.750
21.8%
20.0%
3 · on1.0002-0.5000.500
13.2%
15.6%
4 · the0.0001-0.250-0.250
4.9%
7.4%
5 · mat1.20000.0001.200
16.2%
31.4%

Larger slopes penalize distance more strongly. Content still matters: in the opening example, key 0 is farther away than key 4 but receives more weight. These values are calculated from toy logits, not measured from a trained model.

Change the selected key to compare distances, or change the head to compare slopes. Change the query to see the causal mask exclude future keys. The content logits stay fixed so you can isolate the effect of the bias and the set of allowed keys; these are teaching values, not outputs from a trained model.

The selected pair alone does not determine its attention weight. Softmax also needs every other allowed key's biased logit. In the opening example, key 0 receives about 20.0% with ALiBi, compared with 35.9% without the distance bias. Those percentages use the same six allowed keys in both columns.

From distance to a normalized weight

For query position i, key position j and head h, first compute the scaled content score:

c_{ij}^{(h)} = \frac{q_i^{(h)} \cdot k_j^{(h)}}{\sqrt{d_k}}

ALiBi then subtracts a linear distance penalty. In causal self-attention, the allowed keys satisfy j \le i, so their distance is i-j:

b_{ij}^{(h)} = -m_h(i-j), \qquad s_{ij}^{(h)} = c_{ij}^{(h)} + b_{ij}^{(h)} \quad (j \le i)

Future keys get s_{ij}^{(h)} = -\infty from the causal mask, regardless of their content score or distance bias. Finally, normalize over all allowed keys:

a_{ij}^{(h)} = \frac{\exp(s_{ij}^{(h)})}{\sum_{\ell=0}^{i}\exp(s_{i\ell}^{(h)})}, \qquad o_i^{(h)} = \sum_{j=0}^{i} a_{ij}^{(h)}v_j^{(h)}

The weights sum to one for each query and head. They mix the value vectors exactly as in ordinary attention; ALiBi changes the logits used to form those weights.

For two allowed keys, the softmax denominator cancels in their ratio:

\frac{a_{ij}^{(h)}}{a_{ik}^{(h)}} = \exp\!\left(c_{ij}^{(h)}-c_{ik}^{(h)}-m_h\big[(i-j)-(i-k)\big]\right)

With equal content scores, each extra token of distance multiplies the relative weight by e^{-m_h}. With unequal scores, content and distance compete. That is why the explorer's weights do not necessarily decrease in order of distance.

Different heads, different distance penalties

The slopes are fixed at initialization. For a power-of-two head count H, the reference schedule is, with heads indexed from zero:

m_h = 2^{-8(h+1)/H}, \qquad h=0,\ldots,H-1

The ratio between consecutive slopes is 2^{-8/H}. It is one half for eight heads, giving 0.5, 0.25, …, 0.00390625. For four heads the ratio is one quarter, giving 0.25, 0.0625, 0.015625 and 0.00390625. Larger slopes impose a stronger preference for recent keys; smaller slopes leave distant keys less penalized. Neither slope creates a hard attention cutoff.

For non-power-of-two head counts, the official slope generator preserves the schedule for the largest smaller power of two, then appends alternating entries from the schedule with twice that many heads. Six heads therefore use:

code
head: 0 1 2 3 4 5 slope: 0.25 0.0625 0.015625 0.00390625 0.5 0.125

Select six heads in the explorer to inspect that ordering. It is not a new geometric progression over six entries, and head 4 has a steeper slope than head 0.

The mask and the bias have different jobs

The causal mask prevents access to future tokens. ALiBi expresses a soft preference among the remaining tokens. Its self-distance bias is zero, but the current token's final weight still depends on all other allowed logits.

The unmasked absolute-distance matrix is symmetric. A causally masked attention matrix is not: future entries are excluded. You cannot reuse its lower triangle as an equivalent upper triangle.

There is a useful implementation simplification for causal attention. Since -m_h(i-j)=m_hj-m_hi, the term -m_hi is constant across one query row. Softmax is unchanged by adding or removing a row-wide constant. The reference implementation uses this property to construct a key-position bias and combine it with the causal mask. This is row-shift invariance, not symmetry of causal attention.

The following full-sequence example makes both jobs explicit. It assumes unpadded Q, K and V tensors of shape [batch, heads, tokens, head_dim], with the head slopes already chosen:

code
# src/pages/concepts/transformers/alibi/index.mdx import torch def causal_alibi_attention(q, k, v, slopes): """Compute full-sequence causal ALiBi attention for unpadded Q/K/V.""" scores = (q @ k.transpose(-2, -1)) / q.size(-1) ** 0.5 positions = torch.arange(q.size(-2), device=q.device) distance = positions[:, None] - positions[None, :] bias = -slopes[None, :, None, None] * distance logits = (scores + bias).masked_fill(distance < 0, float("-inf")) weights = torch.softmax(logits, dim=-1) return weights @ v

This materializes the logits to show the math. A production attention kernel may generate biases inside its tiles instead. With a KV cache, query and key positions must refer to their actual sequence positions, rather than treating the newest query as position zero. Padding masks, dropout and kernel support also need to follow the model's attention implementation.

What length extrapolation means

The distance rule remains defined when inference uses positions beyond the training sequence length. For example, a slope of 0.00390625 gives a bias of −3.99609375 at distance 1023 and −7.99609375 at distance 2047. No new position-table entries are needed to compute either value.

That arithmetic property makes length extrapolation possible; it does not establish model quality at every longer length. Longer contexts change the content, the number of competing keys and the distances seen by the model. Dense attention still has quadratic computation in sequence length, and a KV cache still grows with the stored tokens.

The ALiBi paper, Train Short, Test Long, reports a specific result: a 1.3-billion-parameter language model trained on sequences of length 1024 and evaluated at 2048 matched the perplexity of a sinusoidal-position baseline trained at length 2048. In that comparison, the shorter ALiBi training used 11% less memory and ran 11% faster.

Those are measured results for the paper's experiment. They are not a universal 11% kernel speedup, a guarantee at 16× the training length, or a ranking of ALiBi against every RoPE-based model. The authors' repository provides experiment instructions and checkpoints, including WikiText-103 evaluations with explicit context-window settings.

For an existing checkpoint, keep the position method and slope schedule it was trained with unless you have a tested adaptation procedure. For a new model, evaluate the desired lengths and tasks using the same inference and attention implementation you intend to deploy. ALiBi adds no learned position parameters, while its learned attention projections and the practical costs of longer sequences remain.

If you found this explanation helpful, consider sharing it with others.

Mastodon