T5 Relative Position Bias

Shaw’s scheme gives every offset a learned $d$-dimensional vector, and pays $O(n^2 d)$ for the privilege. T5 asks what happens if the vector is a scalar and the offsets are bucketed.

$b$ is a table of $H \times B$ learned scalars and $\beta$ maps an offset to one of $B$ buckets. For T5-base that is 384 parameters — for the entire model, at every layer, across all context lengths. It is difficult to overstate how little this costs.

Log-spaced buckets

The bucketing is the part worth studying. Near offsets get one bucket each; distant ones are grouped logarithmically, so resolution decays with distance.

with the sign of $\delta$ selecting the half of the table for a bidirectional model, and everything past $D$ clamped into the final bucket. The design encodes a belief that turns out to be right: the difference between offsets 3 and 4 matters, and the difference between 300 and 400 does not.

Extrapolation, and its limit

Because far offsets share buckets, a longer sequence introduces no unseen inputs — every distance past $D = 128$ lands in the bucket the model already trained on. T5 therefore runs at lengths it never saw without crashing, and degrades gently rather than sharply.

But it does not extend context, and it is worth being clear about why. Beyond 128 tokens the bias is a constant: the model is told only that something is far, not how far, and cannot rank two distant candidates by proximity. That is enough to keep perplexity respectable and not enough to retrieve from long context, which is why the long-context work went to rotary scaling instead.

import math

import torch
from torch import Tensor, nn


def bucket(offset: Tensor, buckets: int = 32, max_distance: int = 128) -> Tensor:
    """Signed offset → bucket index, exact when near and log-spaced when far."""
    half = buckets // 2
    idx = (offset > 0).long() * half                          # sign selects the half
    d = offset.abs()

    exact = half // 2
    is_far = d >= exact
    far = exact + (
        torch.log(d.float().clamp(min=1) / exact)
        / math.log(max_distance / exact)
        * (half - exact)
    ).long()
    return idx + torch.where(is_far, far.clamp(max=half - 1), d)


class RelativeBias(nn.Module):
    def __init__(self, heads: int, buckets: int = 32, max_distance: int = 128):
        super().__init__()
        self.table = nn.Embedding(buckets, heads)
        self.buckets, self.max_distance = buckets, max_distance

    def forward(self, n: int, device: torch.device) -> Tensor:          # [1, H, n, n]
        pos = torch.arange(n, device=device)
        b = bucket(pos[None, :] - pos[:, None], self.buckets, self.max_distance)
        return self.table(b).permute(2, 0, 1).unsqueeze(0)

The bias is a function of sequence length alone, so it is built once and reused — never recomputed per layer or per batch element. For a causal model, only the lower triangle is ever read and the sign branch collapses, halving the table.