Relative Position Encoding › Implementation: Relative Position in Self-Attention
Language AI HandbookAll chapters Relative Position Encoding › Implementation: Relative Position in Self-AttentionOne section · Sections in this chapter7 / 16 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Now that we understand the mathematical formulation and the worked example, let's translate it into code. We'll build up the implementation in stages: first the relative position embeddings, then the modified attention score computation, then the modified value aggregation, and finally the complete layer.
Building the Relative Position Embedding Table
The foundation of Shaw-style attention is a table of learned embeddings indexed by relative position. Since relative positions range from to after clipping, we need embeddings. The key implementation detail is converting from relative positions (which can be negative) to array indices (which must be non-negative).
import numpy as np
class RelativePositionEmbedding:
"""
Learnable embeddings for relative positions.
"""
def __init__(self, max_distance, d_k, d_v, seed=None):
"""
Initialize relative position embeddings.
Args:
max_distance: Maximum relative position to distinguish (k)
d_k: Dimension for key-side embeddings
d_v: Dimension for value-side embeddings
seed: Random seed for reproducibility
"""
if seed is not None:
np.random.seed(seed)
# Number of unique relative positions: from -k to +k
n_positions = 2 * max_distance + 1
# Initialize embeddings with small random values
scale_k = np.sqrt(2.0 / (n_positions + d_k))
scale_v = np.sqrt(2.0 / (n_positions + d_v))
self.a_K = np.random.randn(n_positions, d_k) * scale_k
self.a_V = np.random.randn(n_positions, d_v) * scale_v
self.max_distance = max_distance
def get_key_embedding(self, rel_pos):
"""Get the key-side relative position embedding for an offset."""
clipped = max(-self.max_distance, min(self.max_distance, rel_pos))
# Convert from [-k, k] to [0, 2k] for indexing
idx = clipped + self.max_distance
return self.a_K[idx]
def get_value_embedding(self, rel_pos):
"""Get the value-side relative position embedding for an offset."""
clipped = max(-self.max_distance, min(self.max_distance, rel_pos))
idx = clipped + self.max_distance
return self.a_V[idx]The index conversion is the subtle but essential detail. Relative positions range from to , but array indices must be non-negative. Adding shifts the range to :
- Offset (maximum backward distance) maps to index 0
- Offset 0 (same position) maps to index
- Offset (maximum forward distance) maps to index
This simple arithmetic lets us use relative positions as array indices without any conditional logic. The same conversion applies uniformly across all offsets, making the lookup fast and predictable.
Computing Attention Scores with Relative Position
With the embedding table in place, we can implement the modified attention score computation. Recall the formula: . For each query-key pair, we compute both the content term and the position term, then sum them before scaling.
def relative_attention_scores(Q, K, rel_embed):
"""
Compute attention scores with relative position information.
Args:
Q: Query matrix of shape (seq_len, d_k)
K: Key matrix of shape (seq_len, d_k)
rel_embed: RelativePositionEmbedding instance
Returns:
Attention scores of shape (seq_len, seq_len)
"""
seq_len, d_k = Q.shape
scores = np.zeros((seq_len, seq_len))
for i in range(seq_len): # Query position
for j in range(seq_len): # Key position
# Content term: q_i · k_j
content_score = np.dot(Q[i], K[j])
# Position term: q_i · a^K_{j-i}
rel_pos = j - i
a_K = rel_embed.get_key_embedding(rel_pos)
position_score = np.dot(Q[i], a_K)
# Combined score
scores[i, j] = content_score + position_score
# Scale by sqrt(d_k)
scores = scores / np.sqrt(d_k)
return scoresThis implementation iterates over all position pairs, making the logic transparent. For each pair , we compute the content score (), look up the relative position embedding for offset , compute the position score (), and sum them. The final division by applies the same scaling as standard attention to prevent score magnitudes from growing with dimension.
Computing Output with Relative Position in Values
The value aggregation follows the same pattern. For each output position , we aggregate contributions from all positions , but we add the relative position embedding to each value before weighting:
def relative_attention_output(attention_weights, V, rel_embed):
"""
Compute attention output with relative position in values.
Args:
attention_weights: Softmax attention weights (seq_len, seq_len)
V: Value matrix of shape (seq_len, d_v)
rel_embed: RelativePositionEmbedding instance
Returns:
Output of shape (seq_len, d_v)
"""
seq_len, d_v = V.shape
output = np.zeros((seq_len, d_v))
for i in range(seq_len): # Output position
for j in range(seq_len): # Value position
# Get value with relative position embedding
rel_pos = j - i
a_V = rel_embed.get_value_embedding(rel_pos)
value_with_pos = V[j] + a_V
# Weighted contribution
output[i] += attention_weights[i, j] * value_with_pos
return outputThe value modification works similarly to the key modification. For each contributing position , we add its relative position embedding to its value vector before applying the attention weight. This means the output at position contains weighted content and weighted position information about where that content originated. The model can learn to tag information with its source distance, which helps downstream layers reason about the structure of the input.
The Complete Relative Self-Attention Layer
With both components in place, we can assemble the complete layer. This class wraps the QKV projections, relative position embeddings, and the modified attention computation into a single forward pass:
class RelativeSelfAttention:
"""
Self-attention with relative position representations (Shaw et al.).
"""
def __init__(self, embed_dim, d_k, d_v, max_distance=16, seed=None):
"""
Initialize relative self-attention layer.
Args:
embed_dim: Dimension of input embeddings
d_k: Query/key dimension
d_v: Value dimension
max_distance: Maximum relative position to distinguish
seed: Random seed
"""
if seed is not None:
np.random.seed(seed)
# QKV projection matrices
scale_qk = np.sqrt(2.0 / (embed_dim + d_k))
scale_v = np.sqrt(2.0 / (embed_dim + d_v))
self.W_q = np.random.randn(embed_dim, d_k) * scale_qk
self.W_k = np.random.randn(embed_dim, d_k) * scale_qk
self.W_v = np.random.randn(embed_dim, d_v) * scale_v
# Relative position embeddings
self.rel_embed = RelativePositionEmbedding(max_distance, d_k, d_v, seed)
self.d_k = d_k
def forward(self, X):
"""
Compute relative self-attention.
Args:
X: Input embeddings of shape (seq_len, embed_dim)
Returns:
output: Attention output of shape (seq_len, d_v)
attention_weights: Weights of shape (seq_len, seq_len)
"""
# Project to Q, K, V
Q = X @ self.W_q
K = X @ self.W_k
V = X @ self.W_v
# Compute attention scores with relative positions
scores = relative_attention_scores(Q, K, self.rel_embed)
# Softmax
scores_stable = scores - scores.max(axis=1, keepdims=True)
exp_scores = np.exp(scores_stable)
attention_weights = exp_scores / exp_scores.sum(axis=1, keepdims=True)
# Compute output with relative positions in values
output = relative_attention_output(attention_weights, V, self.rel_embed)
return output, attention_weightsThe forward pass follows the standard self-attention pattern with two modifications. After projecting input embeddings to QKV, we use our modified relative_attention_scores function instead of the standard dot product, and we use relative_attention_output instead of the standard weighted sum. The softmax step remains unchanged, as the relative position terms are already incorporated into the scores before normalization.
Testing the Implementation
Let's verify our implementation produces sensible results:
# Create a test sequence
np.random.seed(42)
seq_len = 6
embed_dim = 8
d_k = d_v = 4
max_dist = 3
# Random input embeddings
X = np.random.randn(seq_len, embed_dim)
# Create relative attention layer
rel_attn = RelativeSelfAttention(
embed_dim, d_k, d_v, max_distance=max_dist, seed=123
)
output, weights = rel_attn.forward(X)Relative Self-Attention Test ============================================= Sequence length: 6 Embedding dim: 8 Q/K dimension: 4 V dimension: 4 Max distance: 3 Input shape: (6, 8) Output shape: (6, 4) Weights shape: (6, 6) Attention weights (rows sum to 1): [[0.008 0.028 0.001 0.12 0.62 0.223] [0.26 0.098 0.35 0.157 0.052 0.083] [0.794 0.002 0.077 0.122 0.002 0.002] [0.016 0.394 0.025 0.108 0.356 0.101] [0.475 0.023 0.002 0.13 0.069 0.301] [0.002 0.227 0.001 0.014 0.66 0.097]]
The attention weights look similar to standard self-attention, but they now incorporate relative position information. Each query-key compatibility is biased by the relative distance between positions, even at random initialization. In a trained model, these biases become meaningful: heads specialize in different distance profiles, with some heads prioritizing local context and others focusing on longer-range dependencies.
The embedding matrix shows how each relative offset is encoded as a different vector. Offset 0 (same position) has a distinct pattern from offset -3 (three positions back) or offset +3 (three positions forward). These patterns are learned during training to capture which distances are relevant for the task. After training on a language modeling task, you would typically see the nearby offsets (, ) develop sharper, more structured patterns while the boundary offsets () develop softer patterns reflecting their shared "far away" role.