Self-attention is order-blind: shuffle the tokens and the raw QKᵀ scores come out identical, so a Transformer has to be told where each token sits. The classic fix adds a positional embedding to every token vector — sinusoidal or learned. But those absolute codes are welded to the exact positions seen in training. Feed a model trained at length 1024 a sequence of 3000 and it meets position indices it has never represented, and quality collapses.
ALiBi (Press, Smith & Lewis, 2021 — "Train Short, Test Long") throws positional embeddings out entirely. Instead it adds a fixed, unlearned linear bias straight onto the attention scores, before the softmax.
One matrix: −m·(i−j)
The whole of ALiBi lives in the distance matrix (i − j) — how many steps a key is behind the query. Multiply by a slope and negate:
def distance(N):
i = np.arange(N)[:, None] # (N,1) query index
j = np.arange(N)[None, :] # (1,N) key index
return i - j # (N,N) >=0 in the causal region
def alibi_bias(N, m):
return -m * distance(N) # -m*(i-j) : linear recency penalty
It's 0 on the diagonal and grows more negative to the left — a penalty that grows linearly with how far back a key is. There is nothing to learn.
Per-head slopes form a geometric ladder
Each attention head gets its own fixed slope, so the model attends at many scales at once. For n heads the slopes are the geometric sequence m_h = 2^(-8h/n):
def head_slopes(n):
start = 2 ** (-8 / n) # = 2^(-8/n)
return np.array([start ** (h + 1) for h in range(n)])
head_slopes(8)
# [0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625, 0.0078125, 0.00390625]
# steep (local) -------------------------------> gentle (far)
Steep slopes give sharp recency — only the last few keys survive. Gentle slopes let a head gaze far back. Together they're a built-in, multi-scale recency prior.
The whole attention change is one added line
def alibi_attention(Q, K, V, m):
N, d = Q.shape
scores = (Q @ K.T) / np.sqrt(d) # (N,N) raw QK^T
scores = scores + alibi_bias(N, m) # <-- the ONE new line
scores[~causal_mask(N)] = -np.inf # can't see the future
A = np.stack([softmax(r) for r in scores])
return A @ V, A # output, attention weights
No positional embedding is added to Q, K, or the inputs anywhere. The softmax then turns that straight-line penalty into a clean exponential distance decay, giving each head a bounded effective window of about 1/m.
Why it extrapolates: distance, not position
The bias reads only i − j. Slide the whole window along the sequence and every entry is unchanged — it is translation-invariant, so the length N never appears as an absolute index the way a positional embedding would.
# translation invariance: bias(i,j) depends ONLY on (i-j)
b = alibi_bias(8, 0.25)
assert np.allclose(b[7,3], b[6,2]) # same distance 4 -> same bias
assert np.allclose(b[5,5], 0.0) # distance 0 -> no penalty
# softmax turns the LINEAR bias into an EXPONENTIAL decay in distance:
# weight(d) ∝ exp(-m * d) # bounded effective window ~ 1/m
Train at 1024, run at 3000: the model only ever sees relative distances, and far ones are simply taxed more. The attention-by-distance profile is the same shape at any length, which is exactly why a short-trained model extrapolates — and why ALiBi rides in BLOOM and MPT.
ALiBi vs sinusoidal vs RoPE
Sinusoidal embeddings are absolute and added to the inputs; they don't extrapolate past trained positions. RoPE is relative but applies its signal by rotating Q and K. ALiBi needs no embedding table and no rotation — just one added matrix of −m·(i−j), with a fixed per-head slope. It's the cheapest recipe of the three, and the one that most naturally handles test sequences longer than training.
Build the QKᵀ matrix, watch the bias ramp add on top, drag the per-head slope, and slide test length past train length to see the decay curves stay glued together, live at: https://dev48v.infy.uk/dl/day54-alibi.html
Top comments (0)