ALiBi

ALiBi adds nothing to the embeddings. It biases the attention logits directly with a penalty proportional to how far back the key is, using a slope $m_h$ fixed per head.

The slopes are a geometric sequence, not learned. For $H$ heads the ratio starts at $2^{-8/H}$, which spreads the heads from steeply local to nearly position-blind.

Why it extrapolates

A model trained with learned or sinusoidal absolute positions has no defined behaviour past its training length: it is being asked about embeddings it never saw. ALiBi has no such embeddings. Evaluated at four times the training length its perplexity keeps falling, because the only thing the extra distance changes is a larger penalty — and a large enough penalty is simply a soft local window.

That framing is also the limitation. What extrapolates is the recency bias, not the ability to retrieve from far away: an ALiBi model handed a fact 8k tokens back has been actively discouraged from attending to it.

import math
import torch
from torch import Tensor


def alibi_slopes(n_heads: int) -> Tensor:
    """Geometric slopes; the paper's fallback covers non-powers of two."""
    def powers(n: int) -> list[float]:
        start = 2 ** -(2 ** -(math.log2(n) - 3))
        return [start ** (i + 1) for i in range(n)]

    if math.log2(n_heads).is_integer():
        return torch.tensor(powers(n_heads))

    closest = 2 ** math.floor(math.log2(n_heads))
    extra = powers(2 * closest)[0::2][: n_heads - closest]
    return torch.tensor(powers(closest) + extra)


def alibi_bias(n_heads: int, seq_len: int) -> Tensor:
    """[1, H, N, N], causal. Build once, reuse for every layer."""
    pos = torch.arange(seq_len)
    distance = (pos[None, :] - pos[:, None]).tril()   # <= 0 below diagonal
    return (alibi_slopes(n_heads)[:, None, None] * distance)[None]

Note the sign: distance is negative or zero on and below the diagonal, so multiplying by a positive slope already gives the penalty — no separate negation. Add the result to the scores before the causal mask, or the masked positions will pick up a finite bias and stop being masked.