ALiBi Bias Matrix

Penalize distant keys before softmax

ALiBi stands for Attention with Linear Biases. It adds a distance penalty to attention scores instead of adding a position vector to token features. Each head has a positive slope m_h: larger slopes penalize distant keys more strongly. The slopes in this exercise are fixed, not learned.

For a query at position i and an allowed key at position j <= i, the penalty is -m_h * (i - j). In scaled dot-product attention it is added to q_i dot k_j / sqrt(d_head) before softmax over keys. This exercise builds only the bias tensor, not the attention output.

Read the axes

With H heads and T positions, return a tensor of shape (H, T, T) before flattening. Its axes are head, query position, key position:

bias[h, i, j] = -m_h * (i - j)   when j <= i
bias[h, i, j] = -1e9            when j > i

The diagonal is zero. Entries below it become more negative as the distance grows. Entries above it contain the exercise's finite marker for future keys. The signed distance i - j is not symmetric: it is negative above the diagonal, where the marker replaces the penalty.

Use the specified slope schedule

For this exercise, use this formula for every positive head count H:

m_h = 2^(-8 * (h + 1) / H),    h = 0, ..., H - 1

^ denotes exponentiation; Python uses **. For H = 2, the slopes are 1/16 and 1/256. For H = 8, they are 1/2, 1/4, ..., 1/256. In this schedule the first head has the largest slope.

This matches the original ALiBi schedule when H is a power of two. The original implementation handles other head counts with an additional construction; this exercise deliberately uses the formula above instead. Do not substitute another schedule when matching the expected output.

A larger slope favors nearby keys when their content scores are equal. It does not force a head to attend locally: a large query/key score can outweigh the penalty. Likewise, a small slope can accumulate a substantial penalty over a long distance. Computing this bias for a longer sequence does not by itself guarantee useful predictions at that length.

Build a small matrix

For head zero with H = 2 and T = 3, the bias is:

[[ 0,       -1e9,    -1e9],
 [-1/16,     0,      -1e9],
 [-2/16,    -1/16,     0 ]]

This runnable example builds both heads using broadcasting:

import jax.numpy as jnp

H, T = 2, 3
h = jnp.arange(H, dtype=jnp.float32)
slopes = 2.0 ** (-8.0 * (h + 1.0) / H)
i = jnp.arange(T, dtype=jnp.float32)[:, None]
j = jnp.arange(T, dtype=jnp.float32)[None, :]
distance = i - j
bias = -slopes[:, None, None] * distance[None, :, :]
bias = jnp.where(j > i, -1e9, bias)

expected = jnp.array([
    [0., -1e9, -1e9],
    [-1/16, 0., -1e9],
    [-2/16, -1/16, 0.],
])
assert bias.shape == (2, 3, 3)
assert jnp.array_equal(bias[0], expected)
assert jnp.array_equal(slopes, jnp.array([1/16, 1/256]))

slopes[:, None, None] has shape (H, 1, 1). Multiplying it by the (1, T, T) distance array gives one matrix per head.

A finite marker is not a strict attention mask

Return exactly -1e9 at future positions to match this exercise's finite expected-output values. This is an output convention, not a claim that -1e9 and negative infinity behave identically inside attention.

Softmax depends on score differences. A finite penalty can fail to block a future key if its content score is large enough. For an actual causal attention calculation, apply an explicit boolean mask with negative infinity after adding the bias. With finite allowed scores and at least one allowed key per query, masked keys then receive zero probability.