Notes on implementing Attention - Eli Bendersky's website

Attention from the code up

Attention blocks are the core building block of modern transformer models. This post walks through implementing them in pure Python with Numpy, with a focus on getting the array shapes and contractions right at every step. We'll assume familiarity with the motivation for attention; the goal here is the mechanics.

The implementations reference two key papers throughout: Attention Is All You Need (AIAYN) by Vaswani et al., and the GPT-3 paper by Brown et al.

Single-sequence scaled attention

We begin with the simplest case: scaled dot-product self-attention on one sequence of tokens, with no masking. The input is a 2D array of shape (N, D), where N is the sequence length (token count) and D is the embedding depth — the length of each token's embedding vector.

A self-attention module is parameterized by three weight matrices: Wk, Wq, and Wv. The AIAYN paper omits bias vectors, so we skip them too. Each weight matrix has shape (D, HS), where HS is the "head size," a fraction of D. The diagram below assumes N=6; @ denotes matrix multiplication in Python/Numpy syntax:

A minimal Numpy implementation follows:

# self_attention the way it happens in the Transformer model. No bias.
# D = model dimension/depth (length of embedding)
# N = input sequence length
# HS = head size
#
# x is the input (N, D), each token in a row.
# Each of W* is a weight matrix of shape (D, HS)
# The result is (N, HS)
def self_attention(x, Wk, Wq, Wv):
    # Each of these is (N, D) @ (D, HS) = (N, HS)
    q = x @ Wq
    k = x @ Wk
    v = x @ Wv

    # kq: (N, N) matrix of dot products between each pair of q and k vectors.
    # The division by sqrt(HS) is the scaling.
    kq = q @ k.T / np.sqrt(k.shape[1])

    # att: (N, N) attention matrix. The rows become the weights that sum
    # to 1 for each output vector.
    att = softmax_lastdim(kq)
    return att @ v  # (N, HS)

The scaling step divides kq by the square root of HS. This keeps dot-product magnitudes in check, preventing them from growing with the contracted dimension's size. The only dependency is a Softmax function that operates across the last axis of an input array:

def softmax_lastdim(x):
    """Compute softmax across last dimension of x.

    x is an arbitrary array with at least two dimensions. The returned array has
    the same shape as x, but its elements sum up to 1 across the last dimension.
    """
    # Subtract the max for numerical stability
    ex = np.exp(x - np.max(x, axis=-1, keepdims=True))
    # Divide by sums across last dimension
    return ex / np.sum(ex, axis=-1, keepdims=True)

For 2D input, the last axis is the columns, so Softmax runs independently over each row, yielding row values that sum to 1.

One dimension note: the second dimension of Wv can differ from that of Wq and Wk. The diagram shows why this works — after Softmax, the matrix is (N, N), and its product with V has the same second dimension as V. AIAYN denotes these dimensions as d_k and d_v, and most practical configurations set them equal. We simplify by making everything D; varying d_k and d_v is a minimal code change.

Batched attention and matrix mechanics

Real training runs process batches of sequences in parallel to exploit modern hardware. The batched input shape is (B, N, D), where B is the batch size. The weight matrices remain (D, HS); Numpy's matmul contracts the last axis of the input against the first axis of the weight matrix, yielding (B, N, HS):

# self_attention with inputs that have a batch dimension.
# x has shape (B, N, D)
# Each of W* has shape (D, D)
def self_attention_batched(x, Wk, Wq, Wv):
    q = x @ Wq  # (B, N, HS)
    k = x @ Wk  # (B, N, HS)
    v = x @ Wv  # (B, N, HS)

    kq = q @ k.swapaxes(-2, -1) / np.sqrt(k.shape[-1])  # (B, N, N)

    att = softmax_lastdim(kq)  # (B, N, N)
    return att @ v  # (B, N, HS)

The batched code differs from the single-sequence version in two ways:

• When k is 3D, a plain matrix transpose is ambiguous, so we explicitly swap the last and penultimate axes while leaving the batch axis intact.
• The scaling factor uses k.shape[-1] to grab the last dimension of k rather than k.shape[1], which only works for 2D arrays.

In fact, this function handles the unbatched case as well, with B=1. From here on, we assume batched inputs implicitly and drop the "batched" qualifier. The core operation is best understood as a reweighting of token embeddings within the sequence. The Softmax output — the attention matrix — is shape (N, N). Cell (R, C) holds a high value when token C strongly informs token R's output.

This example from AIAYN illustrates how the model learns coreference. In the purple attention head, the token "its" (index 8) assigns high weight to "Law" (index 1), as reflected in a near-1 value at attention matrix cell (8, 1). This intuition becomes essential when we introduce masking.

Multiple heads

Transformer models typically route attention through multiple independent "heads," each with its own K, Q, and V weights. With NH heads, each head outputs (N, HS); these are concatenated along the last dimension into (N, NH * HS) and then fed through a final linear projection layer. AIAYN often sets NH * HS = D (e.g., D=512 with NH=8, HS=64), though this equality isn't strictly required.

The combined implementation, with masking logic to be ignored for now, is:

# x has shape (B, N, D)
# In what follows:
#   NH = number of heads
#   HS = head size
# Each W*s is a list of NH weight matrices of shape (D, HS).
# Wp is a weight matrix for the final linear projection, of shape (NH * HS, D)
# The result is (B, N, D)
# If do_mask is True, each attention head is masked from attending to future
# tokens.
def multihead_attention_list(x, Wqs, Wks, Wvs, Wp, do_mask=False):
    # Check shapes.
    NH = len(Wks)
    HS = Wks[0].shape[1]
    assert len(Wks) == len(Wqs) == len(Wvs)
    for W in Wqs + Wks + Wvs:
        assert W.shape[1] == HS
    assert Wp.shape[0] == NH * HS

    # List of head outputs
    head_outs = []

    if do_mask:
        # mask is a lower-triangular (N, N) matrix, with zeros above
        # the diagonal and ones on the diagonal and below.
        N = x.shape[1]
        mask = np.tril(np.ones((N, N)))

    for Wk, Wq, Wv in zip(Wks, Wqs, Wvs):
        # Calculate self attention for each head separately
        q = x @ Wq  # (B, N, HS)
        k = x @ Wk  # (B, N, HS)
        v = x @ Wv  # (B, N, HS)

        kq = q @ k.swapaxes(-2, -1) / np.sqrt(k.shape[-1])  # (B, N, N)

        if do_mask:
            # Set the masked positions to -inf, to ensure that a token isn't
            # affected by tokens that come after it in the softmax.
            kq = np.where(mask == 0, -np.inf, kq)

        att = softmax_lastdim(kq)  # (B, N, N)
        head_outs.append(att @ v)  # (B, N, HS)

    # Concatenate the head outputs and apply the final linear projection
    all_heads = np.concatenate(head_outs, axis=-1)  # (B, N, NH * HS)
    return all_heads @ Wp  # (B, N, D)

One further optimization sometimes appears in production code: storing all heads in a separate fourth dimension instead of looping over a list. See the Vectorizing across the heads dimension section for details.

Why masked attention matters

An encoder block lets all tokens attend to all other tokens, suitable for understanding or translation. Generative models need a different behavior: during training, if a word could attend to future words, the model would simply copy the answer rather than learn to predict the next token from its predecessors. Decoder blocks enforce causality via masking.

Consider the sentence: "People like watching funny cat videos". Masking forces attention weights for future-token relations to zero:

Looking at the earlier multi-head code with do_mask=True, two steps happen:

1. A lower-triangular (N, N) mask is built with ones on and below the diagonal, zeros above.
2. Before Softmax, masked positions in the scaled kq matrix are set to large negative values, which drives the corresponding output probabilities to zero while preserving proper normalization over remaining positions.

This causal self-attention terminology borrows from causal systems in control theory.

The intuition: blending context

What does attention actually do? For a single input token x[i], the output out[i] blends x[i]'s embedding with contextual information from all preceding tokens x[:i]. The mechanism uses:

• Query (Wq): a representation of what attributes the current token seeks in its context.
• Key (Wk): attributes each context token exposes, which queries can match against.
• Value (Wv): the actual content carried by each context token.

The product q @ K.T computes, per context token, how much that token's value should contribute to the output. Multiplying by V weights those values accordingly. Our implementation vectorizes this across the entire sequence at once — we compute Q from all of x, then use matrix multiplications to process every token in parallel. It's worth noting this explanation is an abstraction. Anthropomorphizing attention — as "tokens caring about each other" — is a useful mental model, but real models stack dozens of transformer layers, each with multiple attention heads operating over and over on intermediate representations.

Cross-attention

Self-attention implies the input sequence attends to itself. Cross-attention, by contrast, lets elements of one sequence attend to another sequence. This shows up in the decoder block of AIAYN, where the decoder query sequence interacts with encoder output. With input sequences xq and xv, each of potentially different length, the former supplies queries while the latter provides keys and values. Output shape becomes (Nq, HS):

Multi-head cross-attention generally skips masking — it's acceptable for elements of xq to attend to all positions in xv:

# Cross attention between two input sequences that can have different lengths.
# xq has shape (B, Nq, D)
# xv has shape (B, Nv, D)
# In what follows:
#   NH = number of heads
#   HS = head size
# Each W*s is a list of NH weight matrices of shape (D, HS).
# Wp is a weight matrix for the final linear projection, of shape (NH * HS, D)
# The result is (B, Nq, D)
def multihead_cross_attention_list(xq, xv, Wqs, Wks, Wvs, Wp):
    # Check shapes.
    NH = len(Wks)
    HS = Wks[0].shape[1]
    assert len(Wks) == len(Wqs) == len(Wvs)
    for W in Wqs + Wks + Wvs:
        assert W.shape[1] == HS
    assert Wp.shape[0] == NH * HS

    # List of head outputs
    head_outs = []

    for Wk, Wq, Wv in zip(Wks, Wqs, Wvs):
        q = xq @ Wq  # (B, Nq, HS)
        k = xv @ Wk  # (B, Nv, HS)
        v = xv @ Wv  # (B, Nv, HS)

        kq = q @ k.swapaxes(-2, -1) / np.sqrt(k.shape[-1])  # (B, Nq, Nv)

        att = softmax_lastdim(kq)  # (B, Nq, Nv)
        head_outs.append(att @ v)  # (B, Nq, HS)

    # Concatenate the head outputs and apply the final linear projection
    all_heads = np.concatenate(head_outs, axis=-1)  # (B, Nq, NH * HS)
    return all_heads @ Wp  # (B, Nq, D)

Merging the Head-Space Matmuls

The list-of-weights version is easy to follow, but it’s awkward for GPU and TPU execution. The fix is to pack all the per-head weight matrices into a single, wider matrix and let the matmul do the looping for us.

To see why that works, look at a plain matmul of an (8, 6) input by a (6, 2) weight matrix:

If we want to multiply that same input by a second (6, 2) weight matrix, we can just glue the two weight matrices side-by-side along their columns:

The result is exactly the two separate matmuls assembled side by side. When the yellow blocks coincide, the green output block is identical in both cases; the violet output is simply the input times the red block. That’s a direct consequence of how matrix multiplication distributes over column concatenation.

Multi-head attention is a perfect fit for this trick. The input x is multiplied by three independent families of weight matrices—one per head for Q, K, and V. If we stack all those weight matrices into one giant matrix of shape (D, 3 * D) (assuming NH * HS = D), one matmul computes all queries, keys, and values at once. The vectorized code below then slices and reshapes the result:

# x has shape (B, N, D)
# In what follows:
#   NH = number of heads
#   HS = head size
#   NH * HS = D
# W is expected to have shape (D, 3 * D), with all the weight matrices for
# Qs, Ks, and Vs concatenated along the last dimension, in this order.
# Wp is a weight matrix for the final linear projection, of shape (D, D).
# The result is (B, N, D).
# If do_mask is True, each attention head is masked from attending to future
# tokens.
def multihead_attention_vec(x, W, NH, Wp, do_mask=False):
    B, N, D = x.shape
    assert W.shape == (D, 3 * D)
    qkv = x @ W  # (B, N, 3 * D)
    q, k, v = np.split(qkv, 3, axis=-1)  # (B, N, D) each

    if do_mask:
        # mask is a lower-triangular (N, N) matrix, with zeros above
        # the diagonal and ones on the diagonal and below.
        mask = np.tril(np.ones((N, N)))

    HS = D // NH
    q = q.reshape(B, N, NH, HS).transpose(0, 2, 1, 3)  # (B, NH, N, HS)
    k = k.reshape(B, N, NH, HS).transpose(0, 2, 1, 3)  # (B, NH, N, HS)
    v = v.reshape(B, N, NH, HS).transpose(0, 2, 1, 3)  # (B, NH, N, HS)

    kq = q @ k.swapaxes(-1, -2) / np.sqrt(k.shape[-1])  # (B, NH, N, N)

    if do_mask:
        # Set the masked positions to -inf, to ensure that a token isn't
        # affected by tokens that come after it in the softmax.
        kq = np.where(mask == 0, -np.inf, kq)

    att = softmax_lastdim(kq)  # (B, NH, N, N)
    out = att @ v  # (B, NH, N, HS)
    return out.transpose(0, 2, 1, 3).reshape(B, N, D) @ Wp  # (B, N, D)

The initial (B, N, D) output is split into the three tensors, and each is reshaped to (B, NH, N, HS). At that point, both B and NH act as independent batch dimensions, and Numpy broadcasts the score and context matmuls over all of them simultaneously. On accelerators, the slicing and transposing often costs nothing—it’s just a different view of the same memory.

Some papers prefer einsum notation for these operations. The score computation from the sample above, for instance, can be written as:

kq = np.einsum("bhqd,bhkd->bhqk", q, k) / np.sqrt(k.shape[-1])

(See these notes on einsum for a deeper look at the notation.)

Full Sample Code

Complete, tested implementations of both the list-based and vectorized versions are in this repository.