Recurrent Memory › Relative Position Encoding in Transformer-XL

We've established that segment-level recurrence breaks absolute position encodings. A token at position 5 in the cached segment and position 5 in the current segment have the same absolute encoding, yet they are 8 positions apart from the perspective of the current segment. How do we fix this? The answer lies in rethinking what position information attention needs.

Attention computes a score between every query and every key to decide how much information should flow from each key to each query. Position information enters this computation to reflect the fact that two tokens at different positions should interact differently than two tokens at the same position. But the original formulation bakes in a specific assumption: that absolute position matters. Transformer-XL's relative position encoding challenges this assumption and replaces it with a more fundamental one: what matters is not where each token is, but how far apart the tokens are.

This changes the role of position information. The same pair of query and key content at a distance of 3 positions should produce the same position-influenced attention score regardless of whether that query is at absolute position 10 or absolute position 1000. The distance is always 3, so the position contribution to attention should always be the same. This property is called translation equivariance, and it is exactly what relative position encodings provide.

Think of relative position encoding as teaching the model about spatial relationships rather than absolute coordinates. Just as you naturally say "the store is two blocks north of here" rather than "the store is at GPS coordinates 37.7749, -122.4194," relative encodings describe relationships, not locations. Two tokens that are adjacent will always have a relative distance of 1, regardless of where they appear in the document.

When you read a sentence like "The cat sat on the mat because it was tired," you understand that "it" refers to "cat" not because of their absolute positions in the document, but because of their relative proximity. The pronoun comes shortly after its antecedent. This observation is the key insight: attention cares about how far apart tokens are, not where they are in absolute terms.

Consider two identical queries, one at position 10 and one at position 100, both attending to keys that are 3 positions before them. If the tokens involved have the same content, shouldn't these attention computations behave similarly? With absolute position encodings, they don't, because positions 7 and 97 have completely different encodings. With relative position encodings, they do, because "3 positions back" always means the same thing.

The implications for generalization are significant. If a model trained on sequences of length 512 learns that pronouns tend to refer to nouns within the previous 20 tokens, that pattern should generalize to any position in any sequence. Absolute position encodings undermine this generalization by making the same pattern look different at different absolute positions. Relative position encodings make the pattern position-agnostic, which is exactly what we want.

To understand how Transformer-XL achieves relative position encoding, we need to first dissect how position information enters standard attention. In the original transformer, the attention score between a query at position $i$ and a key at position $j$ starts as a simple dot product:

$$\text{score}_{ij} = \mathbf{q}_i^\top \mathbf{k}_j$$

where:

• $\text{score}_{ij}$: the attention score determining how much position $i$ attends to position $j$
• $\mathbf{q}_i$: the query vector at position $i$
• $\mathbf{k}_j$: the key vector at position $j$

But where does position come in? The original transformer adds position embeddings to token embeddings before projecting to queries and keys. So the query at position $i$ is $W_q(\mathbf{x}_i + \mathbf{p}_i)$ and the key at position $j$ is $W_k(\mathbf{x}_j + \mathbf{p}_j)$. Substituting these into the dot product and using the distributive property of matrix multiplication:

$$\text{score}_{ij} = (\mathbf{x}_i + \mathbf{p}_i)^\top W_q^\top W_k (\mathbf{x}_j + \mathbf{p}_j)$$

where:

• $\mathbf{x}_i, \mathbf{x}_j$: token embeddings at positions $i$ and $j$
• $\mathbf{p}_i, \mathbf{p}_j$: absolute position embeddings for positions $i$ and $j$
• $W_q, W_k$: learnable query and key projection matrices

This is where the magic of algebra reveals hidden structure. When we expand this product using the distributive property, we get four distinct terms:

$$\text{score}_{ij} = \underbrace{\mathbf{x}_i^\top W_q^\top W_k \mathbf{x}_j}_{\text{content-content}} + \underbrace{\mathbf{x}_i^\top W_q^\top W_k \mathbf{p}_j}_{\text{content-position}} + \underbrace{\mathbf{p}_i^\top W_q^\top W_k \mathbf{x}_j}_{\text{position-content}} + \underbrace{\mathbf{p}_i^\top W_q^\top W_k \mathbf{p}_j}_{\text{position-position}}$$

Each term tells us something different about why one token might attend to another:

• Content-content: "Does this token's meaning relate to that token's meaning?" This is pure semantic matching, independent of where the tokens appear.
• Content-position: "Given what this token is looking for, does that position matter?" For example, a verb might preferentially attend to its subject, which typically precedes it.
• Position-content: "Given where this token is, does that token's content matter more?" Early positions might attend differently than late positions.
• Position-position: "Do these two positions have an inherent affinity?" Adjacent positions might naturally attend to each other.

Why does this decomposition matter? Because it shows us exactly which parts of the attention score use absolute positions. Terms (b), (c), and (d) all depend on $\mathbf{p}_i$ or $\mathbf{p}_j$ (or both), the absolute position embeddings. These are the terms we need to fix. Term (a) is already position-free and can stay as is.

The problem with this decomposition is that all position information uses absolute positions $\mathbf{p}_i$ and $\mathbf{p}_j$. Transformer-XL's insight is that we can rewrite these terms to use relative position instead. The redesign makes two key changes:

1. Replace the key's absolute position with relative distance: Instead of $\mathbf{p}_j$ (the absolute position of the key), use $\mathbf{r}_{i-j}$ (the relative distance from query to key).

2. Replace the query's position with a learned global bias: The query's absolute position $\mathbf{p}_i$ becomes learned vectors $\mathbf{u}$ and $\mathbf{v}$ that don't depend on position at all.

The first change makes sense because the key's absolute position is what we want to replace. We want to say "this key is 5 positions before this query" rather than "this key is at absolute position 47." The second change is more subtle: the query's absolute position would still be ambiguous across segment boundaries even after fixing the key's encoding, so we replace it with a global learned bias that captures the query's general tendency to prefer nearby or distant keys, independent of where in the document the query appears.

The resulting formula is:

$$\text{score}_{ij} = \underbrace{\mathbf{x}_i^\top W_q^\top W_{k,E} \mathbf{x}_j}_{(a)} + \underbrace{\mathbf{x}_i^\top W_q^\top W_{k,R} \mathbf{r}_{i-j}}_{(b)} + \underbrace{\mathbf{u}^\top W_{k,E} \mathbf{x}_j}_{(c)} + \underbrace{\mathbf{v}^\top W_{k,R} \mathbf{r}_{i-j}}_{(d)}$$

where:

• $\mathbf{x}_i, \mathbf{x}_j$: token embeddings at positions $i$ and $j$, unchanged from before
• $\mathbf{r}_{i-j}$: a sinusoidal encoding of the relative distance $i - j$, not the absolute position
• $W_{k,E}$: key projection matrix for content (the "E" stands for embeddings)
• $W_{k,R}$: key projection matrix for relative positions (the "R" stands for relative)
• $\mathbf{u}$: a learned global bias for content attention, shared across all query positions
• $\mathbf{v}$: a learned global bias for position attention, also shared across all positions

Why does this formula make sense? Notice that every term that previously depended on $\mathbf{p}_i$ or $\mathbf{p}_j$ now depends on either a relative distance encoding $\mathbf{r}_{i-j}$ or a global learned constant. The relative distance $i-j$ is the same for any pair of tokens separated by the same gap, regardless of their absolute positions. This means the formula produces the same attention contribution for identical content pairs at identical relative distances, regardless of where in the document those pairs appear.

Let's unpack each component:

1. Term (a): Pure content-based attention. This is unchanged from standard attention. The word "cat" attends to "feline" because of semantic similarity, regardless of position.

2. Term (b): Content-dependent distance preference. The query content determines how much the model cares about distance. A pronoun might strongly prefer nearby tokens, while a discourse marker might look further back.

3. Term (c): Global content importance. Some tokens are just important regardless of the query's position. The beginning-of-sentence token might receive attention from everywhere.

4. Term (d): Global distance preference. The model learns a general preference for certain distances. Typically, nearby tokens receive more attention than distant ones.

Terms (b) and (d) now depend on $i - j$ rather than on $i$ and $j$ separately. This means the position signal is the same whether we're at the start of the document or the end, whether we're attending within the current segment or reaching back into cached memory.

With the theory in place, let's implement relative position encoding step by step. We need two components: (1) a function to generate sinusoidal encodings for each possible relative distance, and (2) a function to compute attention scores using the four-term formula.

The sinusoidal encoding for relative positions works similarly to absolute position encodings, but instead of encoding absolute positions 0, 1, 2, ..., we encode relative distances ..., -2, -1, 0, 1, 2, .... Negative distances mean the key is before the query; positive distances mean the key is after the query (though in causal attention, we only see non-positive distances). The sinusoidal function ensures smooth, gradual variation between nearby distances, so encodings for similar distances are similar vectors.