NNX ALiBi-Style Attention

Give nearby positions a smaller penalty

A query-key dot product measures how well two projected vectors match. It does not directly tell attention how far apart their positions are. This exercise adds that information by subtracting a distance penalty from each scaled attention score before softmax:

score[h, i, j] = dot(Q[h, i], K[h, j]) / sqrt(Dh) - slope[h] * abs(i - j)

Here i is the query position, j is the key position, and Dh is the head width. The diagonal has zero penalty. Keys farther from the query receive a more negative bias, in either direction.

This is an ALiBi-style symmetric bias. The original causal language model method also excludes future positions. This exercise has no causal mask: a distance penalty reduces a score but does not forbid attention.

One fixed slope per head

Use this exercise's schedule for every supported head count H:

slope[h] = 2 ** (-8 * (h + 1) / H), for h = 0, ..., H - 1

With four heads the slopes are [1/4, 1/16, 1/64, 1/256]. They are constants, not learned parameters. This matches the original schedule for power-of-two head counts; the original implementation has a separate construction for other head counts. Follow the formula above here even when H is not a power of two.

A larger slope gives a stronger distance preference, holding the content scores fixed. It does not force a head to attend locally: a sufficiently large query-key match can still outweigh the penalty. Likewise, a small slope is not position-free; its penalty accumulates over long distances.

For two keys with equal content scores, each additional unit of distance multiplies the ratio of their softmax weights by exp(-slope[h]). This gives a precise meaning to the distance preference.

A small example

Use two positions and four heads. Set all content scores to zero so the effect of the bias is visible. Values at the two positions are 0 and 1.

import jax
import jax.numpy as jnp

slopes = 2.0 ** (-8.0 * (jnp.arange(4) + 1) / 4)
positions = jnp.arange(2)
distance = jnp.abs(positions[:, None] - positions[None, :])
bias = -slopes[:, None, None] * distance[None, :, :]
weights = jax.nn.softmax(bias, axis=-1)
values = jnp.array([[0.0], [1.0]])
output = weights @ values
# At position 0, the output is the weight assigned to position 1.
expected = 1.0 / (1.0 + jnp.exp(slopes))
assert jnp.allclose(output[:, 0, 0], expected)
assert jnp.allclose(output[:, 1, 0], 1.0 - expected)
assert jnp.allclose(output[0, :, 0], jnp.array([0.4378235, 0.5621765]))

Even the first position assigns positive weight to the second. A causal mask would forbid that connection and produce zero for the first output in this example. The exercise also learns Q/K/V and output projections; the example isolates the position bias.

Longer sequences

The distance formula needs no learned table row for each position, so it can be evaluated at lengths beyond those used during training. That is a property of the formula, not a guarantee of model quality. Longer sequences introduce unfamiliar distances and more competing keys, and performance must be evaluated on the intended task.

This implementation still constructs dense (H, T, T) scores and weights. The position bias does not remove their quadratic storage cost.

For comparison, the diagram's ALiBi tab shows a causal penalty: future positions are blank. It also uses illustrative slopes 2 ** (-(h + 1)), rather than this exercise's H-dependent schedule. Use it to explore the difference between position embeddings and score biases; use the formula above for the solution.