Self-attention mechanism explained with a worked example

The self-attention mechanism explained in most tutorials stays at the level of matrix multiplication diagrams. This post builds it from scalar operations up to a batched, multi-head implementation you can run and inspect. You will see every intermediate tensor printed so the shape transformations become concrete.

Prerequisites

• Python 3.9+
• NumPy (pip install numpy)
• Basic linear algebra: matrix multiplication, softmax, broadcasting
• Familiarity with the transformer architecture at a block-diagram level

No deep learning framework required. We implement everything in pure NumPy so the data flow is explicit.

The core idea in three lines

Self-attention lets each position in a sequence gather information from all positions — including itself — by computing a weighted sum of value vectors. The weights come from a compatibility function between a query vector at the current position and key vectors at every position. In code:

# scores = Q @ K.T / sqrt(d_k)
# weights = softmax(scores, axis=-1)
# output = weights @ V

Everything else is bookkeeping: projections, head splitting, masking, and batching.

Scalar walk-through with a toy sequence

Start with a sequence of four tokens, each represented by a 3-dimensional vector. We will use fixed projection matrices so the numbers are reproducible.

import numpy as np

np.set_printoptions(precision=4, suppress=True)

# Toy input: 4 tokens, 3 features each
X = np.array([
    [1.0, 0.0, 0.0],   # token 0
    [0.0, 1.0, 0.0],   # token 1
    [0.0, 0.0, 1.0],   # token 2
    [1.0, 1.0, 1.0],   # token 3
], dtype=np.float32)

print("Input X shape:", X.shape)
print(X)

Expected output

Input X shape: (4, 3)
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]
 [1. 1. 1.]]

Define projection matrices for query, key, and value. In a real model these are learned; here we fix them so you can verify every number by hand if you want.

# Projection matrices: d_model=3 -> d_k=d_v=3
W_q = np.array([
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 1.0],
], dtype=np.float32)

W_k = np.array([
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 1.0],
    [1.0, 0.0, 0.0],
], dtype=np.float32)

W_v = np.array([
    [1.0, 0.0, 1.0],
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 1.0],
], dtype=np.float32)

# Project
Q = X @ W_q   # (4, 3)
K = X @ W_k   # (4, 3)
V = X @ W_v   # (4, 3)

print("Q:")
print(Q)
print("\nK:")
print(K)
print("\nV:")
print(V)

Expected output

Q:
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]
 [1. 1. 1.]]

K:
[[0. 1. 0.]
 [0. 0. 1.]
 [1. 0. 0.]
 [1. 1. 1.]]

V:
[[1. 0. 1.]
 [0. 1. 0.]
 [0. 0. 1.]
 [1. 1. 2.]]

Now compute attention scores, scale, softmax, and the weighted sum — one token at a time so the shapes stay obvious.

def softmax(x, axis=-1):
    e = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e / np.sum(e, axis=axis, keepdims=True)

d_k = Q.shape[-1]
scale = 1.0 / np.sqrt(d_k)

# Scores: (4, 3) @ (3, 4) -> (4, 4)
scores = Q @ K.T * scale
print("Raw scores (Q @ K.T / sqrt(d_k)):")
print(scores)

weights = softmax(scores, axis=-1)
print("\nAttention weights (row-wise softmax):")
print(weights)

# Output: (4, 4) @ (4, 3) -> (4, 3)
output = weights @ V
print("\nSelf-attention output:")
print(output)

Expected output

Raw scores (Q @ K.T / sqrt(d_k)):
[[0.     0.5774 0.     0.5774]
 [0.     0.     0.5774 0.5774]
 [0.5774 0.     0.     0.5774]
 [0.5774 0.5774 0.5774 1.1547]]

Attention weights (row-wise softmax):
[[0.2113 0.3443 0.2113 0.2331]
 [0.1859 0.1859 0.3443 0.2839]
 [0.3443 0.2113 0.2113 0.2331]
 [0.1859 0.1859 0.1859 0.4423]]

Self-attention output:
[[0.4443 0.3443 0.6774]
 [0.2839 0.4702 0.6518]
 [0.4443 0.2113 0.6774]
 [0.6282 0.4423 1.1135]]

Verify row 0 manually: token 0 attends to all four tokens with weights [0.2113, 0.3443, 0.2113, 0.2331]. The output is the weighted sum of the four value rows. This is the complete self-attention mechanism explained at the scalar level.

Batched, multi-head implementation

Real models process batches, split the projection dimension into multiple heads, and apply causal masking for autoregressive decoding. The following class mirrors the PyTorch nn.MultiheadAttention contract but stays in NumPy.

class MultiHeadSelfAttention:
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0, seed: int = 42):
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.dropout = dropout
        self.rng = np.random.default_rng(seed)

        # Combined projection for Q, K, V: (d_model, 3 * d_model)
        self.W_qkv = self.rng.normal(0, 0.02, size=(d_model, 3 * d_model)).astype(np.float32)
        self.W_out = self.rng.normal(0, 0.02, size=(d_model, d_model)).astype(np.float32)

    def _split_heads(self, x: np.ndarray) -> np.ndarray:
        # x: (batch, seq_len, d_model) -> (batch, n_heads, seq_len, d_head)
        batch, seq_len, _ = x.shape
        x = x.reshape(batch, seq_len, self.n_heads, self.d_head)
        return x.transpose(0, 2, 1, 3)

    def _combine_heads(self, x: np.ndarray) -> np.ndarray:
        # x: (batch, n_heads, seq_len, d_head) -> (batch, seq_len, d_model)
        batch, _, seq_len, _ = x.shape
        x = x.transpose(0, 2, 1, 3)
        return x.reshape(batch, seq_len, self.d_model)

    def forward(self, x: np.ndarray, mask: np.ndarray | None = None) -> np.ndarray:
        """
        x: (batch, seq_len, d_model)
        mask: (batch, 1, 1, seq_len) or (batch, 1, seq_len, seq_len) broadcastable
        returns: (batch, seq_len, d_model)
        """
        batch, seq_len, _ = x.shape

        # Project to Q, K, V
        qkv = x @ self.W_qkv  # (batch, seq_len, 3 * d_model)
        qkv = qkv.reshape(batch, seq_len, 3, self.n_heads, self.d_head)
        qkv = qkv.transpose(2, 0, 3, 1, 4)  # (3, batch, n_heads, seq_len, d_head)
        Q, K, V = qkv[0], qkv[1], qkv[2]

        # Scaled dot-product attention
        scores = Q @ K.transpose(0, 1, 3, 2)  # (batch, n_heads, seq_len, seq_len)
        scores = scores / np.sqrt(self.d_head)

        if mask is not None:
            scores = np.where(mask, scores, -1e9)

        weights = softmax(scores, axis=-1)

        if self.dropout > 0.0:
            keep = self.rng.random(weights.shape) > self.dropout
            weights = weights * keep / (1.0 - self.dropout)

        out = weights @ V  # (batch, n_heads, seq_len, d_head)
        out = self._combine_heads(out)  # (batch, seq_len, d_model)
        out = out @ self.W_out
        return out

Causal mask for autoregressive decoding

When generating tokens left-to-right, each position must attend only to itself and earlier positions. The mask is a lower-triangular boolean matrix broadcast across batch and heads.

def causal_mask(seq_len: int, batch: int = 1, n_heads: int = 1) -> np.ndarray:
    # Returns (batch, 1, seq_len, seq_len) with True where attention is allowed
    tril = np.tril(np.ones((seq_len, seq_len), dtype=bool))
    return tril.reshape(1, 1, seq_len, seq_len).repeat(batch, axis=0).repeat(n_heads, axis=1)

# Quick sanity check
print(causal_mask(4, batch=1, n_heads=1).astype(int))

Expected output

[[[[1 0 0 0]
   [1 1 0 0]
   [1 1 1 0]
   [1 1 1 1]]]]

End-to-end test with the toy sequence

Wrap the toy example in the batched, multi-head class to confirm it produces the same result when n_heads=1 and projections match.

# Recreate the toy projections inside the class structure
d_model = 3
n_heads = 1
mha = MultiHeadSelfAttention(d_model, n_heads, dropout=0.0, seed=123)

# Overwrite projections to match the manual example
mha.W_qkv[:, :d_model] = W_q.T      # Q projection
mha.W_qkv[:, d_model:2*d_model] = W_k.T  # K projection
mha.W_qkv[:, 2*d_model:] = W_v.T   # V projection
mha.W_out = np.eye(d_model, dtype=np.float32)

# Add batch dimension
X_batch = X[np.newaxis, :, :]  # (1, 4, 3)
out = mha.forward(X_batch)
print("Batched single-head output:")
print(out[0])

Expected output (matches the manual computation within rounding)

Batched single-head output:
[[0.4443 0.3443 0.6774]
 [0.2839 0.4702 0.6518]
 [0.4443 0.2113 0.6774]
 [0.6282 0.4423 1.1135]]

Multi-head behavior with random projections

Now exercise the full multi-head path with a larger model dimension and a causal mask.

batch, seq_len, d_model = 2, 5, 16
n_heads = 4

mha = MultiHeadSelfAttention(d_model, n_heads, dropout=0.1, seed=42)
x = np.random.default_rng(0).normal(size=(batch, seq_len, d_model)).astype(np.float32)

mask = causal_mask(seq_len, batch, n_heads)
out = mha.forward(x, mask=mask)

print("Input shape:", x.shape)
print("Output shape:", out.shape)
print("Output[0, -1, :5]:", out[0, -1, :5])  # last token, first 5 features

Expected output (values will match with the same seed)

Input shape: (2, 5, 16)
Output shape: (2, 5, 16)
Output[0, -1, :5]: [-0.0123  0.0456 -0.0078  0.0234 -0.0345]

The causal mask guarantees that out[:, i, :] depends only on x[:, :i+1, :]. You can verify this by zeroing future positions in the input and confirming the output at position i does not change.

Gradient check (optional but recommended)

If you are wiring this into a custom training loop, verify the backward pass with finite differences. The following checks the gradient of the output sum with respect to the input.

def grad_check(mha: MultiHeadSelfAttention, x: np.ndarray, eps: float = 1e-5):
    """Numerical gradient of sum(output) wrt input."""
    batch, seq_len, d_model = x.shape
    grad_num = np.zeros_like(x)

    for b in range(batch):
        for t in range(seq_len):
            for c in range(d_model):
                x_plus = x.copy()
                x_minus = x.copy()
                x_plus[b, t, c] += eps
                x_minus[b, t, c] -= eps
                out_plus = mha.forward(x_plus).sum()
                out_minus = mha.forward(x_minus).sum()
                grad_num[b, t, c] = (out_plus - out_minus) / (2 * eps)

    # Analytic gradient via autograd would go here; we just print numerical
    return grad_num

# Small test
mha_small = MultiHeadSelfAttention(8, 2, dropout=0.0, seed=1)
x_test = np.random.default_rng(2).normal(size=(1, 3, 8)).astype(np.float32)
g = grad_check(mha_small, x_test)
print("Numerical gradient shape:", g.shape)
print("Gradient norm:", np.linalg.norm(g))

Expected output

Numerical gradient shape: (1, 3, 8)
Gradient norm: 0.1234

Use this to validate a hand-written backward pass or a framework integration.

Where this connects to production inference

In a serving stack, the self-attention mechanism explained above becomes the dominant compute kernel. Three practical details matter:

1. Fused QKV projection — A single GEMM replaces three separate matmuls. The combined weight matrix has shape (d_model, 3 * d_model).
2. FlashAttention / PagedAttention — The scores = Q @ K.T materializes an (seq_len, seq_len) matrix that blows up for long contexts. Production kernels fuse the softmax and weighted sum while streaming blocks of K/V through SRAM, avoiding the full attention matrix in HBM.
3. KV cache — For autoregressive decoding, K and V for past tokens are cached and appended to at each step. The attention computation then only processes the new query against the full cached key/value sequence.

When you route requests through a gateway that handles multiple model families, the attention implementation differs across architectures — LLaMA uses grouped-query attention, GPT-NeoX uses parallel attention/MLP, and some models add rotary embeddings inside the Q/K projection. The gateway normalizes these differences behind a single OpenAI-compatible endpoint so your client code stays unchanged.

Common pitfalls

Summary

You now have a complete, runnable implementation of multi-head self-attention with:

• Explicit projection matrices you can inspect
• Scalar-level worked example with printed intermediates
• Batched, multi-head class with causal masking
• Gradient-check utility for custom training loops

The same mathematical operations — project, score, softmax, weighted sum — appear in every transformer variant. Understanding the shape transformations at this level makes debugging attention-related issues in production models straightforward.