Understanding Transformers · Part 7 of 16

Self-Attention, Step by Step

The central mechanism of the transformer, built up one operation at a time and run on a worked example.

Three roles every token plays

Each token's vector is projected through three learned matrices to produce three new vectors, each with a job:

• The Query (q) describes what this token is looking for.
• The Key (k) describes what this token can be matched on, how it advertises itself to others.
• The Value (v) is the information this token passes along if another token attends to it.

These names describe how the projections behave. The roles emerge from training, which looks for projections that make prediction work.

The five steps

1. Score. Each Query takes a dot product with every Key. For n tokens this gives an n×n score matrix; row i holds token i's scores against every token.
   scores = Q K^T
2. Scale. With many dimensions, dot products grow large because they sum more terms: if the entries are independent with zero mean and unit variance, the dot product's variance is dk. Large scores push softmax toward all-or-nothing, which starves the gradients (Vaswani et al., 2017). Dividing by √dk keeps the scores in a workable range.
   scaled = Q K^T / √d_k
3. Mask (when needed). A decoder-only model applies a causal mask: future positions are set to −∞ before softmax so they get zero weight. Token i may use tokens 1..i, never i+1 onward.
4. Normalize. Softmax turns each row of scores into positive weights that sum to one, an attention budget the token spends across the others.
   softmax(z_i) = exp(z_i) / Σ_j exp(z_j)
5. Mix. Multiply each Value by its weight and add. The result is a weighted blend of information from every allowed position.
   Attention(Q,K,V) = softmax(Q K^T / √d_k) V

A worked example: “Mouse eats cheese”

Three tokens, two features each, so we can follow every number. Embeddings:

Choose a Query projection that swaps the two features, and use identity matrices for Keys and Values so we can watch the attention mechanism alone:

Step 1 — scores (QK^T)

Every Query dotted with every Key gives a 3×3 matrix. Rows are “who is asking,” columns are “who is being asked about.”

Step 2 — scale by √d_k = √2 ≈ 1.414

Every score of 4 becomes about 2.83; the zeros stay zero.

Steps 3–4 — softmax each row into weights

(No mask here, so every token may attend to all three.)

Step 5 — mix the Values

Each row of weights blends the Value vectors [2,0], [1,2], [2,0]:

Reading the result: Mouse and Cheese start from the same embedding, so they ask the same question and both spend most of their attention budget on “Eats,” the one token whose Key aligns with their Query. Their new vectors move toward Eats' information. “Eats,” whose Query matches all three Keys equally, takes an even blend of the whole sentence. Every token runs this process at once.

import numpy as np

# Mouse, Eats, Cheese
X  = np.array([[2, 0], [1, 2], [2, 0]])
Wq = np.array([[0, 1], [1, 0]])
Wk = Wv = np.eye(2)

# Q=[[0,2],[2,1],[0,2]]  K=V=X
Q, K, V = X @ Wq, X @ Wk, X @ Wv

scores = Q @ K.T                    # step 1: score
scaled = scores / np.sqrt(2)        # step 2: scale by sqrt(d_k)

# step 4: normalize
weights = np.exp(scaled) / np.exp(scaled).sum(axis=1, keepdims=True)

# step 5: mix -- matches the table above
output = weights @ V

Keys and Values are different jobs

The example above used identity matrices for W_K and W_V to isolate the mechanism. In a trained model they are different learned projections, because they solve different problems: Keys are optimized for matching, Values for what gets passed along once a match is made. A second small example makes the difference visible. Embeddings:

Applying each projection to the same three vectors:

The same starting vector now has one set of coordinates for being found and a different set for what it delivers. A Key answers “should you attend to me?”; a Value answers “now that you have, what do you take away?”

# Chef, Cuts, Bread
X  = np.array([[2, 0], [1, 2], [2, 1]])
Wk = np.array([[1, 0], [-0.5, 1]])
Wv = np.array([[1, 1], [0.5, 1]])

K = X @ Wk   # Chef=[2,0]  Cuts=[0,2]  Bread=[1.5,1]
V = X @ Wv   # Chef=[2,2]  Cuts=[2,3]  Bread=[2.5,3]