Skip to content

← Flax step 29 of 100

Hard Primitives

ALiBi: Attention with Linear Biases

Jump to the task

Why this matters

Ordinary self-attention compares token content through query-key dot products. ALiBi adds a distance penalty to those scores before softmax. A nearby token gets a smaller penalty than a distant one, but a strong content match can still outweigh the penalty.

This exercise uses bidirectional attention: every token can attend to every position, including later positions. The bias is symmetric in distance. A causal language model would need a separate mask to exclude future tokens.

A penalty for each head

Let H = num_heads, D = qkv_features and d = D // H. For head h, use:

slope[h] = 2 ** (-8 * (h + 1) / H)
bias[h, i, j] = -slope[h] * abs(i - j)
score[h, i, j] = dot(query[i, h], key[j, h]) / sqrt(d) + bias[h, i, j]

Apply softmax over key positions j, then use those weights to combine the values. Only the diagonal has zero bias. Every other entry is negative. Heads with smaller slopes penalize distance less, but they are not forced to learn a particular role.

Use this slope formula for every allowed head count in this exercise. It matches the original ALiBi schedule for power-of-two head counts; the original implementation uses a different construction for other counts.

Unlike a learned position table, this formula does not need an extra row of parameters for each new position. That makes the bias easy to compute at another sequence length. It does not guarantee that a trained model will perform well on longer sequences.

A small example

For two heads the slopes are 1/16 and 1/256. The first head assigns biases [0, -1/16, -2/16, -3/16] to query position zero. The following example isolates the bias by setting all content scores to zero:

import jax
import jax.numpy as jnp

H, T = 2, 4
slopes = 2.0 ** (-8.0 * (jnp.arange(H) + 1) / H)
positions = jnp.arange(T)
distance = jnp.abs(positions[:, None] - positions[None, :])
bias = -slopes[:, None, None] * distance[None, :, :]
weights = jax.nn.softmax(bias, axis=-1)

assert bias.shape == (2, 4, 4)
assert jnp.allclose(bias, bias.swapaxes(-1, -2))
assert jnp.allclose(bias[0, 0], jnp.array([0., -.0625, -.125, -.1875]))
assert jnp.all(weights[:, 0, 1:] > 0)  # Later positions remain available.
assert weights[0, 0, 0] > weights[0, 0, 3]

The bias is symmetric, but the attention weights need not be: softmax normalizes each row separately, and actual query-key scores can differ in both directions.

Add the bias inside attention

Flax Linen's dot_product_attention accepts an additive bias. Define a callback that closes over your (H, T, T) bias and forwards the projected queries, keys, values, mask and other keyword arguments to that function. Pass the callback as attention_fn when constructing nn.MultiHeadDotProductAttention.

The module still creates its usual query, key, value and output projections. For this unbatched input the projected arrays have shape (T, H, d) and the scores have shape (H, T, T). The bias adds directly to those scores.

Do not pass the distance penalties as mask. A mask selects which scores are allowed; it does not add their numeric values. No mask is required here. Also, do not add the bias to the input embeddings or to the softmax output.

Problem

Implement mha_alibi(seed, x, num_heads, qkv_features):

  1. Build the slopes and symmetric distance bias defined above.
  2. Create a custom attention callback that adds the bias through flax.linen.attention.dot_product_attention.
  3. Construct nn.MultiHeadDotProductAttention(num_heads=H, qkv_features=D, attention_fn=your_callback). Keep its other defaults, including learned projection biases, zero dropout and output width x.shape[-1].
  4. Initialize the module directly with jax.random.PRNGKey(int(seed)) and x. Apply those variables to x without a mask or decoding cache.
  5. Return the (T, D_in) output flattened in row-major order.

x is a float32 array of shape (T, D_in). The dimensions and head count are positive, and D is divisible by H. D_in need not equal D or be divisible by H. The output has shape (T * D_in,).

If you JIT-compile the whole function, keep seed, num_heads and qkv_features static. Do not add normalization, residual connections, dropout or a learned position table.

The diagram uses this exercise's slopes and shows both directions.

Loading visualization…