Self-attention from scratch

Self-attention is the single mechanism that defines every modern LLM. It’s also surprisingly simple math — a few matrix multiplications and a softmax. The hard part is keeping the shapes straight. Let’s build it from scratch.

The mental model

For each token in the sequence, attention computes a weighted average of all other tokens’ values, where the weights come from how well that token’s “query” matches the other tokens’ “keys.”

For token i:
  q_i = "what am I looking for?"
  k_j = "what do I have to offer?" (for every j)
  v_j = "what do I contribute?" (for every j)

  weights = softmax(q_i · k_j / √d_k) over all j
  output_i = Σ weights_j · v_j

Q, K, V are all derived from the same input tokens — that’s why it’s “self” attention. The token both asks and answers.

That softmax over all j is the whole game: every token spreads a fixed budget of attention across the sequence. Each row below is one token’s distribution — and the “it” row reaches back to its referent, “cat”:

Build it, end to end

We’ll attend over a sequence of 4 tokens, each represented by a 6-dim vector.

import numpy as np

rng = np.random.default_rng(0)

# A sequence of 4 tokens, each a 6-dim embedding
# Shape: (seq_len, d_model) = (4, 6)
x = rng.standard_normal((4, 6))

# Three projection matrices that turn embeddings into Q, K, V
# In a real model these are nn.Linear layers — here we use random matrices
d_model = 6
d_k = 8           # head dimension (in real models, d_model // n_heads)
W_q = rng.standard_normal((d_model, d_k))
W_k = rng.standard_normal((d_model, d_k))
W_v = rng.standard_normal((d_model, d_k))

# Step 1: project x to Q, K, V
Q = x @ W_q       # (4, 8)
K = x @ W_k       # (4, 8)
V = x @ W_v       # (4, 8)
print(f"Q shape: {Q.shape}, K shape: {K.shape}, V shape: {V.shape}")

Q shape: (4, 8), K shape: (4, 8), V shape: (4, 8)

So far we’ve just projected each token three times. Now the attention math:

import numpy as np

rng = np.random.default_rng(0)
x = rng.standard_normal((4, 6))
d_k = 8
Q = x @ rng.standard_normal((6, d_k))
K = x @ rng.standard_normal((6, d_k))
V = x @ rng.standard_normal((6, d_k))

# Step 2: scores = Q · K^T   shape: (4, 4)
# scores[i, j] = how much token i "queries" token j
scores = Q @ K.T
print("scores shape:", scores.shape)

# Step 3: scale by sqrt(d_k) to keep variance under control
scores = scores / np.sqrt(d_k)

# Step 4: softmax over the LAST axis (j) so each row sums to 1
def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

weights = softmax(scores, axis=-1)
print("weights row sums:", weights.sum(axis=-1))
print("weights[0]:", weights[0].round(3))   # how token 0 attends to each token

# Step 5: weighted sum of values
output = weights @ V       # (4, 4) @ (4, 8) = (4, 8)
print("output shape:", output.shape)

scores shape: (4, 4)
weights row sums: [1. 1. 1. 1.]
weights[0]: [0.137 0.859 0.001 0.003]
output shape: (4, 8)

That’s it. The “transformer revolution” is five lines of NumPy:

scores  = Q @ K.T / sqrt(d_k)        # (n, n) — pairwise similarity
weights = softmax(scores, axis=-1)   # (n, n) — each row is a distribution
output  = weights @ V                # (n, d) — weighted blend of values

Why scale by √d_k?

Q · K^T is a sum of d_k products — its variance grows with d_k, pushing values into softmax’s flat regions where gradients vanish. Dividing by √d_k keeps variance constant regardless of head dimension. Without it, large models don’t train.

Causal masking — making generation work

For autoregressive (GPT-style) models, token i must not attend to future tokens — otherwise training would just cheat by looking at the next token directly.

Fix: mask the upper triangle of scores with -inf before softmax. -inf becomes 0 after softmax — zero attention weight.

import numpy as np

seq_len = 4
# A causal mask: 1 where attention is allowed, 0 where forbidden
mask = np.tril(np.ones((seq_len, seq_len)))
print("causal mask:")
print(mask)

# Apply: set forbidden positions to -inf BEFORE softmax
scores = np.array([
    [1.0, 2.0, 0.5, 1.5],
    [0.3, 2.1, 0.8, 0.9],
    [1.2, 0.5, 1.8, 0.4],
    [0.7, 1.3, 1.0, 2.5],
])
masked = np.where(mask, scores, -np.inf)
print("\nmasked scores:")
print(masked)

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

print("\nweights after softmax (notice upper triangle = 0):")
print(softmax(masked, axis=-1).round(3))

causal mask:
[[1. 0. 0. 0.]
 [1. 1. 0. 0.]
 [1. 1. 1. 0.]
 [1. 1. 1. 1.]]

masked scores:
[[ 1.  -inf -inf -inf]
 [ 0.3  2.1 -inf -inf]
 [ 1.2  0.5  1.8 -inf]
 [ 0.7  1.3  1.   2.5]]

weights after softmax (notice upper triangle = 0):
[[1.    0.    0.    0.   ]
 [0.142 0.858 0.    0.   ]
 [0.301 0.15  0.549 0.   ]
 [0.098 0.178 0.132 0.592]]

Notice: the first row only attends to position 0. The second row attends to positions 0–1. By the last row, the token can see everything. This is how GPT learns to predict the next token from the previous ones.

Cost reality

Q @ K.T has shape (n, n) — quadratic in sequence length. Doubling the context window quadruples this matrix. That’s why “million-token context” needed years of engineering: FlashAttention, sparse patterns, KV-cache tricks. The math we just wrote is the bottleneck.

In real PyTorch

import torch.nn.functional as F

# Fused, FlashAttention-backed when available — 2–4x faster, less memory.
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

F.scaled_dot_product_attention automatically dispatches to FlashAttention or a memory-efficient kernel when available. Use it for any real training run — don’t roll your own.

In one breath

• Self-attention gives each token a weighted average of every token’s value — weights come from how well that token’s query matches the others’ keys (Q, K, V all from the same input).
• The whole thing is three steps: scores = Q·Kᵀ/√d_k, weights = softmax(scores), output = weights·V.
• The √d_k scaling keeps score variance constant across head sizes so softmax doesn’t saturate and kill the gradient.
• Causal masking sets future positions to −∞ before softmax (weight 0), so token i sees only tokens 0..i — what makes autoregressive generation honest.
• Scores are n×n — quadratic in sequence length — which is why long context took FlashAttention, sparse patterns, and KV-cache tricks; in real code use F.scaled_dot_product_attention.

Quick check