Medium Research

Relative Position Encoding

Relative position describes where a key sits compared with a query. For example, a key two tokens ahead has offset +2; a key one token behind has offset -1. Here, each offset selects a scalar bias that is added to an attention score before softmax.

Your task

Implement relative_position_attention(scores, rel_bias, max_dist).

scores has shape (seq_len, seq_len). Row i is a query position and column j is a key position. The scores have already been computed and scaled.
rel_bias has shape (2 * max_dist + 1,), with entries ordered from offset -max_dist through offset +max_dist.
max_dist is a nonnegative integer. Clip larger positive and negative offsets to the corresponding end of the table.

For every query/key pair, compute:

relative_offset = j - i
clipped_offset = max(-max_dist, min(max_dist, relative_offset))
output[i, j] = scores[i, j] + rel_bias[clipped_offset + max_dist]

Return a new tensor with the same shape and floating-point dtype as scores. Both tensors have the same dtype and device. Leave the inputs unchanged. Do not rescale the scores, apply softmax or add a causal mask. Both earlier and later key positions receive a bias.

Example

For a three-token sequence, max_dist=1 and rel_bias=[-2, 0.5, 3], the bias matrix is:

[[ 0.5,  3.0, 3.0],
 [-2.0,  0.5, 3.0],
 [-2.0, -2.0, 0.5]]

Add it elementwise to scores. The top-right offset is +2, which clips to +1. The diagonal uses the middle entry, which need not be zero. When max_dist=0, the table has one entry and that bias is added everywhere.

Scope

This is a clipped scalar-bias exercise. It is simpler than the vector-valued relative key/value representations in Shaw et al. (2018). It also does not use logarithmic distance buckets. The table is supplied; there is no parameter initialization or training step to implement.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Give attention information about token order

Self-attention based only on token content has no absolute position input. Position information can be added to token vectors, added to attention scores, or introduced by transforming queries and keys.

Learned absolute embeddings provide a vector for each position. Sinusoidal encodings use sine and cosine pairs at different frequencies. When adding either to token embeddings, the feature widths must match. In a sinusoidal encoding, position determines the angle while the feature-pair index determines the frequency.

Describe the distance between query and key

Relative-position methods represent the relationship between two positions. A learned score table, for example, can use their signed distance as an index. For query position 3 and key position 1, key - query is -2, while query - key is 2. Use the task's sign convention.

Clipping distances to a fixed range makes all more-distant pairs share an entry at the corresponding end of that range.

Rotate vectors or add a score bias

RoPE rotates pairs of query and key coordinates according to position. Rotations preserve vector lengths. Applying the same frequency scheme to queries and keys makes their dot product depend on the relative rotation between their positions. RoPE transforms the vectors rather than adding an embedding to them.

Some implementations pair adjacent features; others pair corresponding features from the two halves of a vector. Use the specified pairing, because it changes the result.

ALiBi adds a linear distance bias to attention scores, with a slope for each head. It affects preferences among keys without creating a position-embedding vector. Causal masking is still needed separately: a finite negative score bias does not give a future key exactly zero attention.

Check position zero, equal positions, and a small offset. With a cache during generation, use the new token's actual sequence position; restarting at zero on every step gives the wrong positional relationship to cached keys.

Detecting runtime…