Relative Position Embeddings, in depth: Shaw, Transformer-XL, T5 buckets and DeBERTa, and what each costs in memory, kernels and the KV cache

Why relative distance, and what it buys

Take one head with query q_i at position i and key k_j at position j. With absolute embeddings, a position vector is added to each token before projection, so the logit mixes content and position through cross terms, and the model must learn from data that the pairs (5, 3) and (105, 103) both mean "two back". With a relative scheme that sameness is built in: whatever is added depends only on the offset, so a pattern learned in one place applies everywhere. This translation invariance is the main argument for the family.

Two consequences follow. Local patterns such as "the previous token" are learned once, not once per position. And behaviour beyond the training length is at least defined: an absolute table has no row for position 4,097 if it was trained to 4,096, while a relative scheme that clips or buckets distances gives every offset a value. Defined does not mean good; whether quality holds is an empirical question for each scheme.

The cost is that a relative signal depends on the pair (i, j), so it lives inside every attention layer and touches the L-by-L score matrix, which is precisely the matrix fused kernels try never to write to memory. Most of this article is about where that cost shows up.

Shaw et al.: learned vectors on keys and values

Shaw, Uszkoreit and Vaswani (2018) learn two tables indexed by clipped distance. For each offset j minus i, clipped to the range minus k to plus k, there is a key vector a_K and a value vector a_V of head dimension. The logit becomes q_i dotted with (k_j + a_K(j minus i)), scaled as usual, and the output becomes the weighted sum of (v_j + a_V(j minus i)). The tables are shared across the heads of a layer.

Clipping is the key decision. With a clip of k there are only 2k + 1 vectors, so near positions are told apart precisely and everything farther than k looks identical, a simple version of what T5 later did with logarithmic buckets.

The cost is in activations. Done naively you build an L by L by d tensor of a_K for every pair; at L = 4,096 and d = 64 that is about a billion numbers per layer. The workable form computes Q times the 2k + 1 relative key vectors, an L by (2k + 1) matrix, and gathers it into the L by L layout by clipped offset, so batched matrix multiply does the heavy lifting.

The skewing trick: relative logits without an L-by-L-by-d tensor

Without clipping the gather becomes awkward. Music Transformer (Huang et al., 2018) replaced it with a reshape. Compute Q times E-transpose, where E holds one embedding per relative distance from minus (L minus 1) to 0. Pad a zero column on the left, reshape to (L + 1) by L and drop the first row: entry (i, j) is then query i against distance j minus i, exactly the causal relative logit.

import torch
import torch.nn.functional as F

def skew(qe):
    """qe[..., i, r]: query i vs distance r - (L-1). Returns s[..., i, j] = query i vs distance j - i (j <= i)."""
    *lead, L, _ = qe.shape
    padded = F.pad(qe, (1, 0))                 # [..., L, L+1], zero column on the left
    padded = padded.reshape(*lead, L + 1, L)   # reinterpret the same memory
    return padded[..., 1:, :]                  # drop the first row: [..., L, L]

L, d = 6, 8
q, E = torch.randn(L, d), torch.randn(L, d)    # E[r] embeds distance r - (L - 1)
s = skew(q @ E.T)
i, j = 4, 2
assert torch.allclose(s[i, j], q[i] @ E[(j - i) + (L - 1)])

Entries above the diagonal are junk, which a decoder masks anyway; an encoder runs the trick twice with mirrored embeddings or uses the gather. Extra memory drops from L times L times d to L times L, the same order as the scores themselves.

Transformer-XL: four terms and two global biases

Transformer-XL (Dai et al., 2019) caches the previous segment's hidden states and attends to them, so absolute positions would collide between segments. It expands the absolute-position logit into four terms and makes each relative:

1. (a) content to content: q_i dotted with W_KE x_j, ordinary attention.
2. (b) content-dependent position: q_i dotted with W_KR R(i minus j), where R is a fixed sinusoidal encoding of the offset with its own key projection.
3. (c) global content bias: a learned vector u dotted with W_KE x_j: "some keys matter whoever asks".
4. (d) global position bias: a learned vector v dotted with W_KR R(i minus j): "some distances matter regardless of content".

Terms (a) and (c) combine as (q_i + u) against content keys; (b) and (d) combine as (q_i + v) against projected relative encodings, followed by the same skew. The lasting idea is separation: content keys and position keys get separate projections, and u and v store query-independent preferences.

T5: one scalar per head, looked up by bucket

T5 (Raffel et al., 2020) cut the idea to the minimum: each head has a learned scalar per distance bucket, added to the logit before the softmax. Small distances get a bucket each, larger ones share logarithmically wider buckets up to a maximum distance, and everything beyond lands in the last bucket. In bidirectional attention half the buckets are for keys after the query. The function below transcribes the Hugging Face T5 bucketing; note that the offset is key position minus query position.

import math

def t5_bucket(rel, bidirectional=True, num_buckets=32, max_distance=128):
    """rel = key_pos - query_pos. Returns a bucket id in [0, num_buckets)."""
    bucket = 0
    if bidirectional:
        num_buckets //= 2
        if rel > 0:                      # key is to the right of the query
            bucket += num_buckets
        rel = abs(rel)
    else:
        rel = max(-rel, 0)               # causal: only look backwards
    max_exact = num_buckets // 2
    if rel < max_exact:
        return bucket + rel              # one bucket per small distance
    scaled = math.log(rel / max_exact) / math.log(max_distance / max_exact)
    large = max_exact + int(scaled * (num_buckets - max_exact))
    return bucket + min(large, num_buckets - 1)

# bias = nn.Embedding(num_buckets, n_heads); logits += bias[bucket(j - i)] per head

Worked example. With the defaults (32 buckets, maximum distance 128, bidirectional) each direction gets 16 buckets and offsets 0 to 7 are exact. For a key d positions to the left: d = 8 to 11 share bucket 8, 12 to 15 bucket 9, 16 to 22 bucket 10, 23 to 31 bucket 11, d = 32 maps to 12, d = 64 to 14, and 100, 128 and 1,000 all map to 15. Keys to the right get the same numbers plus 16. In causal mode all 32 buckets look backwards: offsets 0 to 15 are exact, d = 64 lands in 26, and everything from about 127 on lands in 31.

So resolution is fine where syntax lives and coarse where only "far" matters, and beyond about 128 the bias carries no ordering at all; content must do that work. Because the bias ignores content, T5 computes it once in the first layer of each stack and reuses it in every later layer. The parameters are trivial; the cost is the materialised bias, heads by L by L, the size of one layer's scores.

DeBERTa: disentangled content and position

DeBERTa (He et al., 2021) represents each token by a content vector and uses relative position embeddings, and sums three terms: content-to-content, content-to-position and position-to-content. It drops position-to-position, which adds nothing about the pair beyond the offset. Offsets are clipped to a maximum distance k, the 2k relative embeddings are shared across layers with per-layer projections, and the scale becomes the square root of 3d because three terms are summed. DeBERTa also shows the limit of purely relative schemes: its Enhanced Mask Decoder adds absolute positions back near the output, because some predictions need to know where in the sentence a token is.

ALiBi and RoPE: the two endpoints

ALiBi is T5 with the table removed: a fixed slope per head times the distance, no parameters and no maximum distance. RoPE removes the additive term entirely, rotating query and key pairs by angles proportional to position so their dot product depends only on the offset; see Rotary Position Embeddings (RoPE). Because RoPE is applied before attention, it composes with every kernel and cache, which is why most current decoder-only models use it.

Kernels and the KV cache: where the choice bites

Fused kernels such as FlashAttention are fast because they never write the score matrix to high-bandwidth memory. Each scheme interacts with that differently:

For a bias your kernel does not support, PyTorch FlexAttention accepts a score modification function that receives the batch, head, query index and key index, so the bias is computed inside the tile loop and never materialised. A dense float mask passed to scaled dot product attention also works, but it costs an L-by-L tensor and usually keeps the dispatcher off the FlashAttention backend, so check which backend actually ran.

from torch.nn.attention.flex_attention import flex_attention

# bucket_ids: [L, L] long tensor of t5_bucket(j - i); rel_bias: nn.Parameter [n_heads, num_buckets]
def t5_score_mod(score, b, h, q_idx, kv_idx):
    return score + rel_bias[h, bucket_ids[q_idx, kv_idx]]

out = flex_attention(q, k, v, score_mod=t5_score_mod)   # q, k, v: [B, H, L, D]

For long-context serving, offset-only biases have a useful property: if the cache evicts middle tokens, nothing stored needs rewriting. RoPE keys carry their original rotation, so renumbering positions after eviction means re-rotating them or caching keys before rotation.

Failure modes

• Sign convention flipped. T5 code uses key minus query; many papers write i minus j. Mixing them swaps past and future buckets. Test with a two-token sequence.
• Bias per layer. Giving every T5 layer its own table changes the parameter count and breaks checkpoint loading; only each stack's first layer owns it.
• Clipping hides long-range order. A model trained on 512 tokens and run at 8,000 has no positional signal past the last bucket. Measure long-range retrieval before claiming length generalisation.
• Materialised bias blows memory. An fp32 bias of 16 heads at L = 16,384 is about 17 GB, the out-of-memory error that "only happens at long context".
• Silent kernel fallback. A dense bias can push attention onto the math backend, slower and quadratic in memory. Profile, or force the backend so it fails loudly.

Choosing a scheme

For a new decoder-only model RoPE is the default because it is free inside the kernel and every serving stack supports it; ALiBi is reasonable where kernels support slopes. Use T5-style buckets for a learned per-head distance prior when you can run FlexAttention or an equivalent, which is common in encoder and encoder-decoder models, and DeBERTa-style attention for encoders where accuracy matters more than throughput. Shaw and Transformer-XL are worth understanding because their ideas recur, but new code rarely implements them. Whatever you pick, test that shifting the input by a constant leaves the attention pattern unchanged, and evaluate beyond the training length on a task that needs far context, not just perplexity.

What to do next

1. Document your offset convention in the attention module and add the two-token sign test.
2. Print the bucket table for your configuration and check the exact range covers the distances your task needs.
3. Confirm the bias is computed once per stack and that loading a reference checkpoint reports no missing or unexpected bias keys.
4. Profile a forward pass at your longest context and record which attention backend ran; if it is the math backend, move the bias into a FlexAttention score function.
5. Evaluate at two and four times the training length on long-range retrieval before advertising a context length.