Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Relative Position Encoding › Efficient Matrix Implementation

Michael BrenndoerferUpdated June 2, 20253 min read

Language AI HandbookAll chapters

Part 14 of 67Positional Encoding

Chapter 99Relative Position Encoding › Efficient Matrix Implementation

Reading section 9 of 16 · Efficient Matrix Implementation · 3 minSection 9 of 16 · 3 min

Section 9 of 16 · 3 min read

Relative Position Encoding › Efficient Matrix ImplementationOne section · Sections in this chapter9 / 16 ↓

Sections in this chapter16

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

The loop-based implementation above is clear but slow. In practice, we compute relative attention efficiently using matrix operations. The key insight is that position biases form a Toeplitz-like structure: all entries on each diagonal share the same relative distance.

The naive loop runs in $O(n^2 d_k)$ time, with a Python-level loop over all $n^2$ position pairs. The efficient version replaces the inner position term computation with a single matrix multiplication, reducing the Python overhead from $O(n^2)$ loop iterations to $O(1)$ . The total asymptotic complexity remains $O(n^2 d_k)$ , but the constant factor drops dramatically because NumPy's matrix operations run in compiled code.

In[20]:

Code

def efficient_relative_scores(Q, K, rel_embed):
    """
    Compute relative attention scores using efficient matrix operations.

    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

    # Content scores: standard Q @ K.T
    content_scores = Q @ K.T

    # Build relative position bias matrix
    # Each entry [i, j] needs Q[i] · a^K_{j-i}

    # Step 1: Collect all relative position embeddings into a matrix
    # Shape: (2*max_dist + 1, d_k)
    all_a_K = rel_embed.a_K
    max_dist = rel_embed.max_distance

    # Step 2: Compute Q @ all_a_K.T to get all possible position scores
    # Shape: (seq_len, 2*max_dist + 1)
    all_position_scores = Q @ all_a_K.T

    # Step 3: Index to build the (seq_len, seq_len) position bias matrix
    position_biases = np.zeros((seq_len, seq_len))
    for i in range(seq_len):
        for j in range(seq_len):
            rel_pos = j - i
            clipped = max(-max_dist, min(max_dist, rel_pos))
            idx = clipped + max_dist
            position_biases[i, j] = all_position_scores[i, idx]

    # Combined scores
    scores = (content_scores + position_biases) / np.sqrt(d_k)

    return scores


# Verify the efficient implementation matches the naive one
scores_naive = relative_attention_scores(Q_test, K_test, rel_attn.rel_embed)
scores_efficient = efficient_relative_scores(Q_test, K_test, rel_attn.rel_embed)

Out[21]:

Console

Maximum difference between naive and efficient: 8.88e-16
Implementations are equivalent!

The efficient version first computes all possible query-position dot products in a single $(n, 2k+1)$ matrix multiply, then indexes into this precomputed matrix. This reduces the inner loop computation from $O(d_k)$ per entry to $O(1)$ , with the upfront cost of a single $(n, 2k+1)$ matrix multiplication. The Toeplitz structure of the position bias matrix is what makes this possible: once you know how every query interacts with every possible offset embedding, you only need to map indices to offsets, which is a pure indexing operation.

Real-world implementations in PyTorch or JAX go further. The indexing step can be replaced with a gather operation or einsum that exploits the diagonal structure, bringing the full bias computation down to a few vectorized calls with no Python loops at all. This is important for GPU utilization: Python loops are sequential, but GPU kernels can parallelize across the entire $(n, n)$ matrix simultaneously.