ALiBi: a positional bias baked into attention scores instead of the embeddings

The problem, from first principles

[04-seq-modeling/02-embeddings/03-sinusoidal-positional-encoding] and [06-rope] both inject positional information by modifying the QUERY/KEY VECTORS themselves, before attention ever runs. Press et al. (2021, "ALiBi," Attention with Linear Biases) proposed a different place entirely to encode position: leave the vectors untouched, and instead add a fixed, position-DEPENDENT PENALTY directly onto the raw attention SCORES, right where [04-flash-attention]'s and [05-sliding-window-attention]'s masks already get added, before softmax. The penalty is simple: for query position i and key position j, subtract a value proportional to the DISTANCE i - j, so keys further in the past get an increasingly large negative bias, and therefore, after softmax, an increasingly SMALL attention weight, a direct, built-in recency preference, without touching the content-based Q @ K^T similarity at all.

Different attention HEADS use different bias STRENGTHS (a range of slopes, from very mild to very aggressive recency preference), letting some heads stay close to ordinary, distance-agnostic attention while others behave almost like [05-sliding-window-attention]'s hard cutoff, all without any explicit windowing logic. A major practical payoff: because the bias is a simple, closed-form FUNCTION of distance rather than a fixed, learned per-position vector, ALiBi generalizes cleanly to sequence lengths LONGER than anything the model was trained on, a genuine weakness of [04-seq-modeling/02-embeddings/03-sinusoidal-positional-encoding]'s learned or fixed-table approaches.

From theory to code

Implement compute_alibi_slopes(num_heads) (a geometric sequence of per-head slopes), compute_alibi_bias(seq_len, num_heads) (-slope * (i - j) for every head and every query/key position pair), and alibi_causal_mask(seq_len, num_heads), combining that bias with [04-seq-modeling/04-attention/02-causal-mask]'s causal mask into one additive mask ready to hand directly to [01-scaled-dot-product-attention].

Constraints

• Slopes form a GEOMETRIC sequence: slopes[h] = ratio^(h+1) for h = 0, ..., num_heads-1, where ratio = 2^(-8/num_heads) (both the starting value AND the common ratio).
• bias[h, i, j] = -slopes[h] * (i - j): zero on the diagonal (i == j), increasingly NEGATIVE as j moves further into the past relative to i (larger i - j).
• alibi_causal_mask still enforces causality: j > i must remain -inf, exactly like [02-causal-mask], with the ALiBi bias added only where the causal mask itself doesn't already forbid attending.
• Output shape is (num_heads, seq_len, seq_len): unlike [02-causal-mask]'s single shared (seq_len, seq_len) mask, ALiBi's mask genuinely differs PER HEAD.

Hints