Relative Position Embeddings

Absolute schemes answer “where is this token”. Almost nothing a language model does needs that answer. Agreement, coreference, bracket matching, the local syntax that most heads are doing — all of it is about distance between two tokens, and the absolute index is a detour.

Shaw et al. cut the detour out. Keep a learned table indexed by offset, and add the entry for $j - i$ to the key (and optionally the value) when computing the score for the pair $(i, j)$.

The offset is clipped to a window: every distance beyond $\pm k$ shares the outermost entry. Shaw found $k = 16$ sufficient for translation, which is a strong claim about how local the useful signal is — everything further away is told only “far”, and the model still works.

Transformer-XL’s decomposition

Dai et al. rederived it from the four-term expansion of an absolute score, and replaced each positional piece with something offset-indexed:

Terms (c) and (d) replace the query’s absolute position — which cannot appear, because the query has none — with learned global vectors $u$ and $w$. Term (b) is the one that does the work: a content-dependent read of the offset, so a head can learn “attend three back if you are a verb” rather than “attend three back”.

This is the form that stuck. T5’s bias drops (b) and keeps a scalar version of (d); ALiBi fixes (d) to a straight line and drops the rest; RoPE gets (b) for free by rotating $q$ and $k$ instead of adding anything.

The memory problem

Terms (b) and (d) need $r_{i-j}$ for every pair. Built directly that is an $[n, n, d]$ tensor, which is worse than the attention matrix it decorates — and it defeats FlashAttention, which exists precisely to avoid instantiating anything $n \times n$.

The escape is that $r_{i-j}$ is a Toeplitz matrix: it has only $2n - 1$ distinct rows, arranged in shifted diagonals. Compute the $[n, 2n-1]$ product once and shift each row into place with a pad-and-reshape.

import torch
from torch import Tensor


def skew(qr: Tensor) -> Tensor:
    """[B, H, n, 2n-1] of q·r_k  →  [B, H, n, n] indexed by (i, j).

    Row i needs r_{i-j} for j = 0..n-1, which is row i of the input shifted by i.
    A pad of one column per row turns the shift into a reshape.
    """
    b, h, n, _ = qr.shape
    qr = torch.nn.functional.pad(qr, (0, 1))                  # [B, H, n, 2n]
    qr = qr.reshape(b, h, 2 * n * n).narrow(2, n - 1, n * (2 * n - 1))
    return qr.reshape(b, h, n, 2 * n - 1).narrow(3, 0, n)


def relative_scores(q: Tensor, r: Tensor) -> Tensor:          # r: [2n-1, d]
    return skew(torch.einsum("bhnd,kd->bhnk", q, r))

Correct, $O(n^2)$ rather than $O(n^2 d)$ in memory, and still an explicit score matrix. That last point is why the scheme did not survive: it is structurally incompatible with fused attention kernels, whereas rotary embeddings apply to $q$ and $k$ before the kernel ever sees them and cost the kernel nothing at all.