T5 Relative Position Bucketing

Group relative distances into buckets

A relative-position bias describes where a key is compared with a query. T5 maps each query/key offset to a bucket ID, then uses that ID to look up a learned bias for each attention head. Nearby distances get separate buckets; more distant positions share buckets.

This exercise computes only the bucket IDs for bidirectional attention. It does not create the learned bias table, compute attention or apply a causal mask. The same distance rule works at different sequence lengths, but that alone does not guarantee good model predictions on longer inputs.

Keep direction and distance separate

Let N = num_buckets, K = N // 2 and E = K // 2. For query position i and key position j, define r = j - i. A positive offset means the key is later in the sequence.

• Use direction offset K when r > 0; otherwise use zero.
• Compute a bucket within that direction using distance a = abs(r).
• Add the direction offset to obtain the final bucket ID.

Zero distance belongs to bucket zero. The first bucket in the positive half, bucket K, is unused: positive offsets cannot have distance zero. Reversing a nonzero offset changes its bucket by K.

Use exact buckets nearby and wider buckets farther away

If a < E, the within-direction bucket is simply a. Otherwise use:

large = E + floor(log(a / E) / log(max_distance / E) * (K - E))
within_direction = min(large, K - 1)

The logarithm makes bucket widths grow with distance. At a = E, the logarithmic term is zero, giving bucket E. At a = max_distance, the expression before clamping reaches K, so clamp to K - 1. The last bucket starts before max_distance; it also contains every distance at or beyond that value.

Use an even N >= 4 and max_distance > E. When N is divisible by four, each direction has equal numbers of exact and logarithmic buckets. For other even counts, use the integer divisions above; do not assume those counts are equal.

Work through a small example

With N = 8 and max_distance = 16, K = 4 and E = 2. Distances 0 and 1 use exact buckets. Distances 2 through 5 share bucket 2; distances 6 and above use bucket 3. Add 4 for positive offsets.

For four positions, r = j - i is:

[[ 0,  1,  2,  3],
 [-1,  0,  1,  2],
 [-2, -1,  0,  1],
 [-3, -2, -1,  0]]

This complete example computes the corresponding IDs:

import jax.numpy as jnp

T, N, max_distance = 4, 8, 16
K, E = N // 2, (N // 2) // 2
positions = jnp.arange(T, dtype=jnp.int32)
r = positions[None, :] - positions[:, None]
distance = jnp.abs(r)
direction = (r > 0).astype(jnp.int32) * K
safe_distance = jnp.maximum(distance, E).astype(jnp.float32)
logarithmic = E + (
    jnp.log(safe_distance / E) / jnp.log(max_distance / E) * (K - E)
).astype(jnp.int32)
within_direction = jnp.where(
    distance < E, distance, jnp.minimum(logarithmic, K - 1)
)
buckets = direction + within_direction

expected = jnp.array([
    [0, 5, 6, 6],
    [1, 0, 5, 6],
    [2, 1, 0, 5],
    [2, 2, 1, 0],
], dtype=jnp.int32)
assert jnp.array_equal(buckets, expected)

Clamp the logarithm's input distance to at least E. That avoids computing log(0) for the diagonal, even though where later selects the exact-distance branch. It leaves the logarithmic branch unchanged wherever that branch is used. Its values are nonnegative, so converting to int32 truncates in the same direction as floor.