Relative Position Bias Attention

Shaw, Uszkoreit & Vaswani, "Self-Attention with Relative Position Representations", Google, 2018 — Raffel et al., "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer" (T5), Google, 2020

Learning Objectives

After completing this chapter, you will be able to:

1. Explain why absolute positional embeddings fail to capture the notion of "how far apart are these two tokens?" and why this matters for generalization
2. Describe how relative position bias injects distance information directly into the attention score computation via an additive bias matrix $B[i,j]$
3. Compute the bias matrix $B[i,j] = -\lambda \cdot \log(1 + |i - j|)$ by hand and explain why logarithmic decay is preferred over linear decay
4. Perform a full relative position bias forward pass on the shared “The cat sat on mat” example, comparing results with and without position bias
5. Explain T5's logarithmic bucketing scheme and why it uses a fixed number of buckets instead of one parameter per distance
6. Analyze how position bias redistributes attention weight from distant tokens to nearby tokens
7. Implement relative position bias attention from scratch in both NumPy and PyTorch
8. Connect relative position bias to its successors: RoPE (Chapter 8) and ALiBi (Chapter 9)

The Problem: Absolute Positions Are Blind to Distance

Why Absolute Positional Embeddings Fail

In Chapter 1, we saw that standard scaled dot-product attention computes $\text{softmax}(QK^\top / \sqrt{d_k}) \cdot V$. This formula treats every token pair identically regardless of how far apart they are in the sequence. The token at position 1 and the token at position 100 get no distance signal at all — the attention mechanism operates purely on content similarity.

The original Transformer (Vaswani et al., 2017) addressed this by adding sinusoidal positional embeddings to the input. Each position $p$ gets a unique vector $PE(p)$ added to the token embedding before computing Q, K, V. This means Q and K implicitly contain position information because they were projected from position-enriched inputs.

But here is the critical problem: these are absolute positions. Position 3 always maps to the same embedding vector, regardless of sequence length or context. The model must learn, from training data alone, that position 3 and position 5 are "close" while position 3 and position 50 are "far." This relationship is never explicitly computed — it must be discovered indirectly from the dot products of position-enriched Q and K vectors.

The Missing Signal

Consider a concrete example. The phrase "the cat sat on the mat" might appear at the start of a document (positions 0-5) or in the middle (positions 500-505). With absolute positions, the model sees completely different positional embeddings in each case, even though the internal structure of the phrase is identical. The model must independently learn that Q[500]·K[502] should produce the same attention pattern as Q[0]·K[2].

The Core Limitation: Standard attention never explicitly computes "how far apart are tokens i and j?" It knows $Q[i] \cdot K[j]$ (content similarity), but not $|i - j|$ (relative distance). The position signal is buried inside the Q and K vectors, mixed with content information, and the model must disentangle them.

This leads to three practical failures:

1. Poor length generalization. A model trained on sequences of length 512 struggles with sequences of length 1024, because it has never seen the positional embeddings for positions 512-1023.
2. Wasted capacity. The model uses some of its attention heads just to discover relative position patterns, instead of having this information provided directly.
3. Inconsistent local patterns. The same local structure ("adjective noun") at different absolute positions produces different Q·K products, forcing the model to learn position-invariant patterns redundantly.

The Intuition Behind Relative Position Bias

The Classroom Analogy

Imagine a student sitting in a classroom lecture. When the professor says something, the student naturally pays more attention to the words just spoken (nearby in time) and less to words from five minutes ago (distant in time). This isn't because the recent words are intrinsically more important — it's because temporal proximity is a strong prior for relevance. Nearby words are more likely to be part of the same thought, sentence, or argument.

The student still can recall something from five minutes ago if it's highly relevant to the current topic. The distance penalty doesn't block long-range connections; it just means that distant information needs to be more relevant (higher content similarity) to compete with local information. This is exactly what relative position bias does to attention.

The Key Insight

The key insight of Shaw et al. (2018) was remarkably simple: add a distance-dependent term directly to the attention scores before softmax. Instead of relying on the model to discover distance from absolute position embeddings, we tell it explicitly. For each pair of tokens at positions $i$ and $j$, we add a bias $B[i,j]$ that depends only on $(i - j)$, the relative offset.

This is a single addition to the standard attention formula:

Standard: $\text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) \cdot V$

With bias: $\text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + B\right) \cdot V$

Everything else remains identical. The bias matrix $B$ is the only new component. It shifts the attention landscape so that nearby token pairs start with a natural advantage, while distant pairs must overcome a distance penalty through high content similarity alone.

Historical Development

The idea developed in two key stages:

1. Shaw et al. (2018) introduced relative position representations in "Self-Attention with Relative Position Representations." They proposed learning separate key and value embeddings for relative positions and clipping distances beyond a maximum offset. This was the first work to demonstrate that relative positions outperform absolute positions on machine translation.
2. Raffel et al. (2020) simplified the idea in the T5 paper. Instead of learning full embedding vectors per distance, T5 uses a single learned scalar bias per relative position bucket. They introduced logarithmic bucketing: short distances get exact representation, while long distances are grouped into shared buckets. This is more parameter-efficient and has become the standard approach.

Mathematical Formulation

Standard Attention (Recap)

From Chapter 1, standard scaled dot-product attention is:

$A = \text{softmax}\!\left(\frac{Q \cdot K^\top}{\sqrt{d_k}}\right) \cdot V$

where $Q \in \mathbb{R}^{N \times d_k}$ is the query matrix, $K \in \mathbb{R}^{N \times d_k}$ is the key matrix, $V \in \mathbb{R}^{N \times d_v}$ is the value matrix, and $N$ is the sequence length.

Adding the Relative Position Bias

Relative position bias modifies this by adding a bias matrix $B \in \mathbb{R}^{N \times N}$ to the scaled scores before softmax:

$A = \text{softmax}\!\left(\frac{Q \cdot K^\top}{\sqrt{d_k}} + B\right) \cdot V$

The key constraint is that $B[i,j]$ depends only on the relative offset $(i - j)$, not on the absolute positions $i$ and $j$ individually. This makes the bias translation-invariant: the same phrase at different positions in the sequence produces the same bias pattern.

How Is B[i,j] Computed?

There are two main approaches:

1. Learned bias (T5 style). Each relative position bucket has a learned scalar parameter. The model learns, through backpropagation, what bias to assign each distance. T5 uses 32 buckets with logarithmic spacing.

2. Fixed bias (our simulation). A mathematical function computes the bias from the distance. Our simulation uses logarithmic decay:

$B[i,j] = -\lambda \cdot \log(1 + |i - j|)$

where $\lambda = 0.3$ is the bias scale. This formula has several desirable properties:

• $B[i,i] = 0$ — no penalty for attending to yourself (distance 0)
• $B[i,j] < 0$ for all $i \neq j$ — all non-self attention gets penalized
• Logarithmic growth means the penalty increases slowly — distance 4 gets only $2.3\times$ the penalty of distance 1, not $4\times$
• Symmetric: $B[i,j] = B[j,i]$ — direction doesn't matter

Symbol-by-Symbol Breakdown

Interactive: Bias Matrix Heatmap

Explore how the bias matrix changes with different bias scales and decay functions. Hover over any cell to see the exact calculation:

Step-by-Step Calculation

Setup: Shared Example

We use the same Q, K, V matrices from every chapter: 5 tokens ("The cat sat on mat"), each with $d_k = 4$ dimensions. The key parameters are: $\lambda = 0.3$ (bias scale), $\sqrt{d_k} = 2.0$ (scaling factor).

Step 1: Raw Scores $Q \cdot K^\top$

First, compute the raw dot products between all query-key pairs. These are identical to Chapter 1 — no position information yet:

Notice: "cat" attending to "The" has the highest raw score (3.0), while "The" attending to "The" has the lowest (0.0). These scores reflect only content similarity.

Step 2: Scaling by $\sqrt{d_k}$

Divide every score by $\sqrt{d_k} = \sqrt{4} = 2.0$:

Step 3: Building the Bias Matrix B

For each distance $d = |i - j|$, compute $B = -0.3 \times \log(1 + d)$:

The complete 5×5 bias matrix (a Toeplitz matrix — each diagonal has the same value):

Step 4: Adding Bias to Scores

Element-wise addition: $\text{biased\_scores} = \text{scaled\_scores} + B$

Notice how "cat"'s score for "The" dropped from 1.500 to 1.2921 (distance 1, mild penalty), while "cat"'s score for "mat" dropped from 0.250 to -0.1659 (distance 3, strong enough to push it negative). Nearby tokens are penalized less.

Step 5: Softmax

Let's trace the softmax computation for "The" (row 0) in detail:

1. Biased scores: $[0.0000, 0.7921, 0.1704, 0.0841, 0.2672]$
2. Subtract max (0.7921): $[-0.7921, 0.0000, -0.6216, -0.7079, -0.5249]$
3. Exponentiate: $[0.4529, 1.0000, 0.5371, 0.4927, 0.5916]$
4. Sum: $3.0743$
5. Normalize: $[0.1473, 0.3253, 0.1747, 0.1603, 0.1924]$

For "sat" (row 2) — the middle token benefits most from relative bias because it has equal-distance neighbors on both sides:

1. Biased scores: $[0.1704, 0.7921, 1.0000, 0.2921, 0.4204]$
2. Subtract max (1.0000): $[-0.8296, -0.2079, 0.0000, -0.7079, -0.5796]$
3. Exponentiate: $[0.4362, 0.8123, 1.0000, 0.4927, 0.5601]$
4. Sum: $3.3013$
5. Normalize: $[0.1321, 0.2460, 0.3029, 0.1492, 0.1697]$

"sat" now gives 30.3% weight to itself (same position, zero penalty), compared to 25.1% in standard attention. The bias has concentrated attention on the center.

Step 6: Output Computation

The output for each token is the weighted sum of all value vectors. For "The":

$\text{Output}[\text{The}] = 0.1473 \times V[\text{The}] + 0.3253 \times V[\text{cat}] + 0.1747 \times V[\text{sat}] + 0.1603 \times V[\text{on}] + 0.1924 \times V[\text{mat}]$

$= [0.2435, 0.4215, 0.2709, 0.2565]$

Interactive: Standard vs Biased Attention

Compare attention weights with and without relative position bias. Select a query token and adjust the bias scale to see how distance penalties redistribute attention:

Attention Weight Analysis

Standard vs Biased: Weight Comparison

The difference matrix (biased weights minus standard weights) reveals the redistribution pattern:

What Changed and Why

The diagonal is always positive — every token gained self-attention weight because distance 0 has zero penalty. The off-diagonal pattern follows distance:

• "The" row: "cat" (d=1) gained +0.0276, while "mat" (d=4) lost -0.0394. The nearest neighbor benefits most.
• "sat" row (center): Self-attention gained the most (+0.0524) because "sat" has equal neighbors on both sides. Both "The" (d=2) and "mat" (d=2) lost equally (-0.0198 and -0.0254).
• "mat" row (edge): Self-attention gained +0.0708 (largest diagonal gain) because "mat" is at the sequence edge — it has the most distant tokens (d=4 to "The"), so the bias strongly penalizes those pairs and redistributes weight to self.

Full Attention Weight Matrices

Standard Attention Weights (no bias)

Relative Position Bias Attention Weights

Full Output Matrices

Standard Output (no bias)

Relative Position Bias Output

Notice that "mat"'s output was $[0.3108, 0.3108, 0.3108, 0.3108]$ (nearly uniform) in standard attention but became $[0.3077, 0.3181, 0.3326, 0.3554]$ with bias — now favoring dim-3 because "mat"'s nearest neighbor "on" has V = $[0, 0, 0, 1]$ (strong in dim-3) and received more weight (+0.0092).

T5's Logarithmic Bucketing Scheme

Why Use Buckets?

A naive implementation would assign one learnable parameter per relative distance. For a sequence of length $N$, this requires $2N - 1$ parameters (distances from $-(N-1)$ to $+(N-1)$). For $N = 4096$, that's 8,191 parameters per attention head — not terrible, but with diminishing returns. The difference between "97 tokens away" and "98 tokens away" is linguistically meaningless. We don't need separate parameters for these nearly identical distances.

T5 solves this with logarithmic bucketing: nearby distances get their own bucket (exact representation), while distant ranges share a bucket (coarse approximation). With 32 buckets, you get:

• Exact region (buckets 0-7): distances 0-7 each get their own bucket. Position 3 is distinguishable from position 4.
• Logarithmic region (buckets 8-15): distances 8-128 are logarithmically distributed across 8 buckets. Distances 8-11 share one bucket, 12-16 share another, etc.
• Bidirectional: the first 16 buckets handle "token j is to the left of token i" and the second 16 handle "token j is to the right."

The Bucket Algorithm

The T5 bucketing algorithm works as follows. Let $r = i - j$ be the signed relative position:

1. If bidirectional, separate positive and negative offsets into two halves of the bucket space
2. For $|r| < n_{\text{exact}}$ (half of half-buckets): assign bucket = $|r|$ directly
3. For $|r| \geq n_{\text{exact}}$: assign bucket = $n_{\text{exact}} + \lfloor \log(|r| / n_{\text{exact}}) / \log(d_{\max} / n_{\text{exact}}) \times (n_{\text{half}} - n_{\text{exact}}) \rfloor$

This is a logarithmic mapping: doubling the distance adds roughly one bucket. The model learns one scalar bias per bucket, shared across all distances that map to that bucket.

Interactive: T5 Bucket Explorer

Explore how T5's logarithmic bucketing maps distances to buckets. Adjust the number of buckets and maximum distance to see how the exact/logarithmic boundary shifts:

Applications Across Domains

Connection to Modern Systems

Complexity Analysis

The key insight is that relative position bias adds $O(N^2)$ work to build and add B, but this is dominated by the $O(N^2 d_k)$ cost of the dot-product computation. In practice, the overhead is negligible — less than 5% of total attention computation time.

Python Implementation

The complete NumPy implementation with full execution trace. Click any line to see the exact values flowing through the computation:

      The    cat    sat     on    mat
The  0.000  0.000  0.000  0.000  0.000
cat  0.000  0.000  0.000  0.000  0.000
sat  0.000  0.000  0.000  0.000  0.000
on   0.000  0.000  0.000  0.000  0.000
mat  0.000  0.000  0.000  0.000  0.000

       The      cat      sat       on      mat
The   0.0000  -0.2079  -0.3296  -0.4159  -0.4828
cat  -0.2079   0.0000  -0.2079  -0.3296  -0.4159
sat  -0.3296  -0.2079   0.0000  -0.2079  -0.3296
on   -0.4159  -0.3296  -0.2079   0.0000  -0.2079
mat  -0.4828  -0.4159  -0.3296  -0.2079   0.0000

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

      The   cat   sat    on   mat
The   0.0   2.0   1.0   1.0   1.5
cat   3.0   0.0   2.0   1.0   0.5
sat   1.0   2.0   2.0   1.0   1.5
on    1.0   1.0   0.0   2.0   1.0
mat   1.0   1.0   1.0   1.0   1.5

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

       The      cat      sat       on      mat
The   0.0000  -0.2079  -0.3296  -0.4159  -0.4828
cat  -0.2079   0.0000  -0.2079  -0.3296  -0.4159
sat  -0.3296  -0.2079   0.0000  -0.2079  -0.3296
on   -0.4159  -0.3296  -0.2079   0.0000  -0.2079
mat  -0.4828  -0.4159  -0.3296  -0.2079   0.0000

       The      cat      sat       on      mat
The  0.0000   0.7921   0.1704   0.0841   0.2672
cat  1.2921   0.0000   0.7921   0.1704  -0.1659
sat  0.1704   0.7921   1.0000   0.2921   0.4204
on   0.0841   0.1704  -0.2079   1.0000   0.2921
mat  0.0172   0.0841   0.1704   0.2921   0.7500

       The      cat      sat       on      mat
The  0.1473   0.3253   0.1747   0.1603   0.1924
cat  0.4099   0.1126   0.2486   0.1335   0.0954
sat  0.1321   0.2460   0.3029   0.1492   0.1697
on   0.1523   0.1660   0.1137   0.3805   0.1875
mat  0.1508   0.1612   0.1758   0.1985   0.3138

       d0       d1       d2       d3
The  0.2435   0.4215   0.2709   0.2565
cat  0.4576   0.1603   0.2963   0.1812
sat  0.2170   0.3309   0.3877   0.2341
on   0.2460   0.2597   0.2074   0.4743
mat  0.3077   0.3181   0.3326   0.3554

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

       The      cat      sat       on      mat
The  0.1473   0.3253   0.1747   0.1603   0.1924
cat  0.4099   0.1126   0.2486   0.1335   0.0954
sat  0.1321   0.2460   0.3029   0.1492   0.1697
on   0.1523   0.1660   0.1137   0.3805   0.1875
mat  0.1508   0.1612   0.1758   0.1985   0.3138

       d0       d1       d2       d3
The  0.2435   0.4215   0.2709   0.2565
cat  0.4576   0.1603   0.2963   0.1812
sat  0.2170   0.3309   0.3877   0.2341
on   0.2460   0.2597   0.2074   0.4743
mat  0.3077   0.3181   0.3326   0.3554

1import numpy as np
2import math
3
4class RelativePositionBiasAttention:
5    """
6    Relative Position Bias Attention (Shaw et al. 2018 / T5 2020)
7
8    Adds a learned or computed bias B[i,j] to the attention scores
9    before softmax. B depends only on the relative distance (i - j),
10    encouraging the model to prefer local context by default.
11
12    Formula: A = softmax(Q·K^T / sqrt(d_k) + B) · V
13    """
14
15    def __init__(self, d_k: int, bias_scale: float = 0.3):
16        """
17        Args:
18            d_k: Dimension of key vectors (for scaling)
19            bias_scale: Controls strength of distance penalty
20        """
21        self.d_k = d_k
22        self.bias_scale = bias_scale
23        self.scale = math.sqrt(d_k)
24
25    def _build_bias_matrix(self, N: int) -> np.ndarray:
26        """Build the N x N relative position bias matrix."""
27        B = np.zeros((N, N))
28        for i in range(N):
29            for j in range(N):
30                B[i, j] = -self.bias_scale * math.log(1 + abs(i - j))
31        return B
32
33    def _softmax(self, x: np.ndarray) -> np.ndarray:
34        """Numerically stable softmax along last axis."""
35        x_shifted = x - np.max(x, axis=-1, keepdims=True)
36        exp_x = np.exp(x_shifted)
37        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
38
39    def forward(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray):
40        """
41        Args:
42            Q: Query matrix  (N, d_k)
43            K: Key matrix    (N, d_k)
44            V: Value matrix  (N, d_v)
45        Returns:
46            weights: Attention weight matrix (N, N)
47            output:  Context vectors (N, d_v)
48        """
49        N = Q.shape[0]
50
51        raw_scores = Q @ K.T
52        scaled_scores = raw_scores / self.scale
53        B = self._build_bias_matrix(N)
54        biased_scores = scaled_scores + B
55        weights = self._softmax(biased_scores)
56        output = weights @ V
57
58        return weights, output
59
60    def explain(self, Q, K, V, tokens, query_idx=0):
61        """Print step-by-step trace for one token."""
62        weights, output = self.forward(Q, K, V)
63        token = tokens[query_idx]
64        B = self._build_bias_matrix(len(tokens))
65        raw = Q @ K.T
66        scaled = raw / self.scale
67
68        print(f"\n=== Relative Position Bias trace for '{token}' ===")
69        print(f"bias_scale = {self.bias_scale}, d_k = {self.d_k}")
70        print(f"\nRaw scores:    {np.round(raw[query_idx], 4)}")
71        print(f"Scaled scores: {np.round(scaled[query_idx], 4)}")
72        print(f"Bias values:   {np.round(B[query_idx], 4)}")
73        print(f"Biased scores: {np.round((scaled + B)[query_idx], 4)}")
74        print(f"\nAttention weights:")
75        for j, t in enumerate(tokens):
76            w = weights[query_idx, j]
77            bar = "#" * int(w * 50)
78            dist = abs(query_idx - j)
79            print(f"  A[{token},{t}] = {w:.4f} (d={dist}) |{bar}|")
80
81        print(f"\nOutput[{token}] = {np.round(output[query_idx], 4)}")
82
83
84# ── Shared Example (used in every chapter) ──
85tokens = ["The", "cat", "sat", "on", "mat"]
86
87Q = np.array([
88    [1.0, 0.0, 1.0, 0.0],   # The
89    [0.0, 2.0, 0.0, 1.0],   # cat
90    [1.0, 1.0, 1.0, 0.0],   # sat
91    [0.0, 0.0, 1.0, 1.0],   # on
92    [1.0, 0.0, 0.0, 1.0],   # mat
93])
94
95K = np.array([
96    [0.0, 1.0, 0.0, 1.0],   # The
97    [1.0, 0.0, 1.0, 0.0],   # cat
98    [1.0, 1.0, 0.0, 0.0],   # sat
99    [0.0, 0.0, 1.0, 1.0],   # on
100    [1.0, 0.0, 0.5, 0.5],   # mat
101])
102
103V = np.array([
104    [1.0, 0.0, 0.0, 0.0],   # The
105    [0.0, 1.0, 0.0, 0.0],   # cat
106    [0.0, 0.0, 1.0, 0.0],   # sat
107    [0.0, 0.0, 0.0, 1.0],   # on
108    [0.5, 0.5, 0.5, 0.5],   # mat
109])
110
111# ── Run Relative Position Bias Attention ──
112rpb = RelativePositionBiasAttention(d_k=4, bias_scale=0.3)
113weights, output = rpb.forward(Q, K, V)
114
115print("=== Relative Position Bias Attention ===")
116print("Weights:")
117print(np.round(weights, 4))
118print("\nOutput:")
119print(np.round(output, 4))
120
121# ── Step-by-step trace for "The" ──
122rpb.explain(Q, K, V, tokens, query_idx=0)
123
124# ── Compare with standard (no bias) ──
125rpb_none = RelativePositionBiasAttention(d_k=4, bias_scale=0.0)
126std_weights, std_output = rpb_none.forward(Q, K, V)
127print("\n=== Standard vs Biased (The) ===")
128print("Standard:", np.round(std_weights[0], 4))
129print("Biased:  ", np.round(weights[0], 4))
130print("Diff:    ", np.round(weights[0] - std_weights[0], 4))

PyTorch Implementation

The PyTorch version adds GPU support, automatic differentiation for training, and uses vectorized operations instead of Python loops. It also includes an $\texttt{nn.Parameter}$ bias table for the T5-style learned approach:

      The   cat   sat    on   mat
The    0    -1    -2    -3    -4
cat    1     0    -1    -2    -3
sat    2     1     0    -1    -2
on     3     2     1     0    -1
mat    4     3     2     1     0

      The   cat   sat    on   mat
The    0     1     2     3     4
cat    1     0     1     2     3
sat    2     1     0     1     2
on     3     2     1     0     1
mat    4     3     2     1     0

       The      cat      sat       on      mat
The   0.0000  -0.2079  -0.3296  -0.4159  -0.4828
cat  -0.2079   0.0000  -0.2079  -0.3296  -0.4159
sat  -0.3296  -0.2079   0.0000  -0.2079  -0.3296
on   -0.4159  -0.3296  -0.2079   0.0000  -0.2079
mat  -0.4828  -0.4159  -0.3296  -0.2079   0.0000

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

       The      cat      sat       on      mat
The  0.0000   0.7921   0.1704   0.0841   0.2672
cat  1.2921   0.0000   0.7921   0.1704  -0.1659
sat  0.1704   0.7921   1.0000   0.2921   0.4204
on   0.0841   0.1704  -0.2079   1.0000   0.2921
mat  0.0172   0.0841   0.1704   0.2921   0.7500

       The      cat      sat       on      mat
The  0.1473   0.3253   0.1747   0.1603   0.1924
cat  0.4099   0.1126   0.2486   0.1335   0.0954
sat  0.1321   0.2460   0.3029   0.1492   0.1697
on   0.1523   0.1660   0.1137   0.3805   0.1875
mat  0.1508   0.1612   0.1758   0.1985   0.3138

       d0       d1       d2       d3
The  0.2435   0.4215   0.2709   0.2565
cat  0.4576   0.1603   0.2963   0.1812
sat  0.2170   0.3309   0.3877   0.2341
on   0.2460   0.2597   0.2074   0.4743
mat  0.3077   0.3181   0.3326   0.3554

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6class RelativePositionBiasAttention(nn.Module):
7    """
8    Relative Position Bias Attention - PyTorch
9
10    Adds a learned bias table to attention scores before softmax.
11    In this version, biases are stored as learnable parameters
12    indexed by relative position, matching the T5 approach.
13    """
14
15    def __init__(self, d_k: int, max_seq_len: int = 512,
16                 num_buckets: int = 32, bias_scale: float = 0.3):
17        super().__init__()
18        self.d_k = d_k
19        self.scale = math.sqrt(d_k)
20        self.bias_scale = bias_scale
21        self.num_buckets = num_buckets
22        self.max_seq_len = max_seq_len
23
24        # Learnable bias table: one scalar per bucket
25        self.relative_bias_table = nn.Parameter(
26            torch.zeros(num_buckets)
27        )
28
29    def _compute_bias(self, N: int) -> torch.Tensor:
30        """Build N x N bias matrix using log-decay formula."""
31        positions = torch.arange(N)
32        # Compute relative distances: (N, N) matrix
33        relative_pos = positions.unsqueeze(0) - positions.unsqueeze(1)
34        distances = relative_pos.abs().float()
35        # Log-decay bias (matches our NumPy simulation)
36        B = -self.bias_scale * torch.log(1 + distances)
37        return B
38
39    def forward(
40        self,
41        Q: torch.Tensor,
42        K: torch.Tensor,
43        V: torch.Tensor,
44        mask: torch.Tensor | None = None,
45    ) -> tuple[torch.Tensor, torch.Tensor]:
46        """
47        Args:
48            Q: (N, d_k), K: (N, d_k), V: (N, d_v)
49            mask: Optional (N, N) boolean, True = masked
50        Returns:
51            weights: (N, N) attention weights
52            output:  (N, d_v) context vectors
53        """
54        N = Q.shape[0]
55
56        scores = torch.matmul(Q, K.T) / self.scale
57        B = self._compute_bias(N).to(Q.device)
58        scores = scores + B
59
60        if mask is not None:
61            scores = scores.masked_fill(mask, float("-inf"))
62
63        weights = F.softmax(scores, dim=-1)
64        output = torch.matmul(weights, V)
65
66        return weights, output
67
68
69# ── Shared Example ──
70tokens = ["The", "cat", "sat", "on", "mat"]
71
72Q = torch.tensor([
73    [1.0, 0.0, 1.0, 0.0],
74    [0.0, 2.0, 0.0, 1.0],
75    [1.0, 1.0, 1.0, 0.0],
76    [0.0, 0.0, 1.0, 1.0],
77    [1.0, 0.0, 0.0, 1.0],
78])
79K = torch.tensor([
80    [0.0, 1.0, 0.0, 1.0],
81    [1.0, 0.0, 1.0, 0.0],
82    [1.0, 1.0, 0.0, 0.0],
83    [0.0, 0.0, 1.0, 1.0],
84    [1.0, 0.0, 0.5, 0.5],
85])
86V = torch.tensor([
87    [1.0, 0.0, 0.0, 0.0],
88    [0.0, 1.0, 0.0, 0.0],
89    [0.0, 0.0, 1.0, 0.0],
90    [0.0, 0.0, 0.0, 1.0],
91    [0.5, 0.5, 0.5, 0.5],
92])
93
94# ── Run with relative position bias ──
95rpb = RelativePositionBiasAttention(d_k=4, bias_scale=0.3)
96with torch.no_grad():
97    weights, output = rpb(Q, K, V)
98
99print("Relative Position Bias Attention Output:")
100print(output.numpy().round(4))
101print("\nAttention Weights:")
102print(weights.numpy().round(4))

Key Takeaways

1. One-line change, big impact. Relative position bias adds a single matrix $B$ to the attention scores before softmax. Everything else — Q·K\u1D40 computation, scaling, softmax, weighted sum — remains identical to standard attention.
2. Distance as a prior. The bias matrix encodes the assumption that nearby tokens are more likely to be relevant. This doesn't prevent long-range attention; it just requires distant tokens to have higher content similarity to overcome the distance penalty.
3. Translation invariance. Because $B[i,j]$ depends only on $(i - j)$, the same phrase at any position in the sequence produces the same bias pattern. This enables better length generalization than absolute positional embeddings.
4. Logarithmic decay. The function $-\lambda \cdot \log(1 + |i-j|)$ grows slowly — quadrupling the distance only doubles the log. This means the model distinguishes fine-grained nearby positions while treating distant positions more coarsely.
5. T5's bucketing is efficient. Instead of learning one parameter per distance, T5 uses 32 logarithmic buckets. This captures fine-grained local position with few parameters and generalizes to sequences longer than those seen during training.
6. Foundation for successors. This chapter lays the groundwork for RoPE (Chapter 8, which uses rotation instead of addition) and ALiBi (Chapter 9, which uses linear decay instead of logarithmic).

Exercises

1. Bias Scale Sensitivity: Recompute the attention weights for "The" using $\lambda = 0.0$, $\lambda = 0.3$, and $\lambda = 1.0$. How does increasing $\lambda$ affect the balance between local and global attention? At what value of $\lambda$ does self-attention dominate?
2. Linear vs Logarithmic Decay: Replace the bias formula with $B[i,j] = -\lambda \cdot |i-j|$ (linear decay). Compute the attention weights and compare. Which decay function penalizes long-range attention more harshly? When might linear be preferred?
3. Asymmetric Bias: In causal (autoregressive) models, we might want asymmetric bias where looking backward (j < i) has a different penalty than looking forward. Design a bias formula that penalizes backward attention less than forward attention. Why might this be useful for language generation?
4. T5 Bucket Mapping: For a model with 32 buckets and max_distance=128, compute the bucket assignment for distances 1, 5, 10, 50, and 100. Verify your answers with the interactive explorer above.
5. Memory Analysis: A model has 32 attention heads, sequence length 4096, and uses T5-style bucketing with 32 buckets. How many parameters does the bias table require? Compare this to storing one parameter per relative distance.

References

1. Shaw, P., Uszkoreit, J., & Vaswani, A. (2018). Self-Attention with Relative Position Representations. NAACL 2018.
2. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., & Liu, P. J. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR, 21(140), 1-67.
3. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, \u0141., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017.
4. Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., & Guo, B. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. ICCV 2021.
5. Press, O., Smith, N. A., & Lewis, M. (2022). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR 2022.
6. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.