Skip to content

← Position Encodings step 4 of 5

Medium Research

ALiBi-Style Distance Bias

A distance bias lowers the attention score between positions that are far apart. Different heads can use different penalty strengths: a larger slope subtracts more for the same distance.

ALiBi applies a distance-based penalty to attention scores rather than adding position vectors to embeddings. This exercise uses a specified symmetric distance penalty and a simplified slope schedule. Implement the formulas below, not another implementation's slope generator or causal mask.

Your task

Implement alibi_attention(S, n_heads).

  • S is a floating-point array of shape (n_heads, seq_len, seq_len).
  • n_heads is a positive integer matching S.shape[0].
  • seq_len is positive. Rows index queries and columns index keys.

For each zero-based head index h, use:

slope = 2**(-8 * h / n_heads)
output[h, i, j] = S[h, i, j] - slope * abs(i - j)

Use ordinary division in the exponent. Return the same shape without modifying S. The first head has slope 1. No softmax or masking is part of this function: return biased scores, not attention weights.

Example

For two heads the slopes are 1 and 1/16. If each score matrix is all ones and the sequence length is two, the outputs are:

head 0: [[1, 0],      head 1: [[1, 0.9375],
         [0, 1]]              [0.9375, 1]]

Diagonal scores stay unchanged because their distance is zero. The penalty is symmetric, so a key two positions ahead receives the same penalty as one two positions behind. The output need not be symmetric if S is not symmetric.

A finite penalty does not exclude future positions. A caller that needs causal attention must apply a separate causal mask before softmax. This task only adds the distance bias.