Attention Math Step by Step: From the Input Matrix X to Softmax and V

Whenever a language model produces a token, one computation runs over and over: attention. It isn't a figure of speech. On every forward pass, for every token, at every layer, the same arithmetic executes.

Most explanations stop at the intuition of "queries, keys and values", and that leaves the actual math as a black box. This piece opens the box. It covers what the input matrix really is, where Q, K and V come from, a numerical example small enough to run yourself, and the reason it's called self-attention. Afterwards, tracing a single attention head by hand should be doable.

The problem attention solves

Take this sentence: "The animal didn't cross the street because it was tired."

When the model reaches the token "it", it must figure out what the word refers to. The animal? The street? The word alone says nothing, so the answer has to come from the surrounding sequence.

Attention handles that by letting each token give every other token a weight, a number for how much it should look at that position while building its updated representation. For "it", the weights could come out like this:

"it" → "animal"  : 0.65
"it" → "street"  : 0.20
"it" → "tired"   : 0.10
"it" → others    : 0.05

Nobody writes those numbers by hand. They fall out of a learned computation built on three vectors, Query, Key and Value, and the entire formula fits on one line:

Attention(Q, K, V) = softmax( QKᵀ / √dₖ ) · V

Everything that follows just unpacks that line.

Step 0: the input matrix X

Before Q, K or V can exist, there has to be an input matrix, X.

Every token in "The cat sat" maps to a learned embedding vector. With an embedding size of d_model = 4, it looks like this:

The → [0.2, 0.5, 0.1, 0.7]
cat → [0.8, 0.3, 0.6, 0.2]
sat → [0.4, 0.9, 0.5, 0.1]

Stack the vectors row by row and you get X, whose shape is [n_tokens, d_model].

One detail is easy to miss: X isn't constant across layers. Each attention layer takes the output of the previous layer as its own X. At the first layer, X is just raw token plus position information. Further up the stack, say layer 12 in GPT-3, which has 96 layers, X already encodes deep contextual meaning built from many rounds of attention.

Step 1: computing Q, K and V

Nobody assigns queries, keys and values manually. All three are computed from X through three separate learned weight matrices:

Q = X · WQ
K = X · WK
V = X · WV

WQ, WK and WV are parameter matrices that stay frozen once training ends. The shapes line up as follows:

X.shape  = [n, d_model]
WQ.shape = [d_model, d_k]   → Q.shape = [n, d_k]
WK.shape = [d_model, d_k]   → K.shape = [n, d_k]
WV.shape = [d_model, d_v]   → V.shape = [n, d_v]

The distinction matters, so here it is in a table:

Symbol	What it is	Depends on the input?
WQ, WK, WV	Learned projection matrices	No, fixed after training
Q, K, V	Computed activations	Yes, they follow the current X

That's why "bank" in "river bank" and "bank" in "deposit at the bank" end up with entirely different attention patterns. The surrounding context produces a different X, which goes through the same weight matrices and comes out as different Q, K and V vectors.

A search engine makes a handy mental model:

• Q is what the token is looking for.
• K is what each token advertises about itself.
• V is the information handed back when a match occurs.

Q and K decide where to look. V decides what actually flows out.

Step 2: the dot product QKᵀ

Take Q and multiply it with K transposed. That gives a square [n, n] matrix of raw scores, where the entry at row i, column j says how strongly token i attends to token j. Before any numbers are filled in, it looks like this:

Keys
       cat   eats  fish
     ┌──────┬──────┬──────┐
Q cat│  ?   │  ?   │  ?   │
  eat│  ?   │  ?   │  ?   │
 fish│  ?   │  ?   │  ?   │
     └──────┴──────┴──────┘

A larger dot product means the two tokens are more relevant to one another.

Step 3: scaling by √dₖ

When d_k is large, raw dot products blow up. With d_k = 512, scores such as [180, 210, 95] are ordinary. Pass those into softmax and the output is [0.00, 1.00, 0.00], a completely collapsed distribution. Gradients then vanish and training stalls.

Dividing by √d_k keeps the scores in a sane range. That division is what earns the name scaled dot-product attention.

Step 4: softmax turns scores into weights

Softmax runs row by row, and each row becomes a probability distribution over all tokens. Here's the full calculation for the scores [3.2, 4.0, 1.0, 6.2, 8.0]:

exp(3.2) =     24.53
exp(4.0) =     54.60
exp(1.0) =      2.72
exp(6.2) =    492.75
exp(8.0) =  2,980.96
         ──────────
Σ        =  3,555.56

Token 1:   24.53 / 3555.56 →  0.69%
Token 2:   54.60 / 3555.56 →  1.54%
Token 3:    2.72 / 3555.56 →  0.08%
Token 4:  492.75 / 3555.56 → 13.86%
Token 5: 2980.96 / 3555.56 → 83.83%  ← winner

Token 5 takes 83.83% even though it leads Token 4 by only 1.8 points. That non-linear amplification is what lets the model commit firmly to the most relevant tokens.

Stability tip: subtract the row maximum first, as in np.exp(x - x.max()), and only then divide. A 64-bit float can't hold exp(800), but after the shift every exponent stays safe, and the resulting probabilities are mathematically identical.

Step 5: the weighted sum with V

The last move multiplies the weights into V and adds everything up. For "it", that means:

output("it") =
    0.65 × V("animal")
  + 0.20 × V("street")
  + 0.10 × V("tired")
  + 0.05 × V(rest)

What comes out is a fresh vector for "it" containing mostly information from "animal", a bit from "street" and very little from the rest. This contextualized representation is what goes to the next layer.

A complete example in NumPy

The sentence here is "cat eats fish" with d_model = 2:

import numpy as np

X  = np.array([[1, 2], [3, 1], [2, 4]], dtype=float)
WQ = np.array([[1, 0], [0, 1]], dtype=float)  # identity
WK = np.array([[1, 0], [0, 1]], dtype=float)  # identity
WV = np.array([[0, 1], [1, 0]], dtype=float)  # swaps dimensions

Q, K, V = X @ WQ, X @ WK, X @ WV

# Raw scores for "fish" (Q=[2,4]) against each key:
# fish→cat:  [2,4]·[1,2] = 10
# fish→eats: [2,4]·[3,1] = 10
# fish→fish: [2,4]·[2,4] = 20  ← self-attention dominates

scores = Q @ K.T / np.sqrt(2)
e = np.exp(scores - scores.max(-1, keepdims=True))
weights = e / e.sum(-1, keepdims=True)  # softmax row-wise

output = weights @ V
# output[fish] ≈ [3.997, 2.001] ≈ V(fish)
# fish self-attends at ~99.8% in this toy example

In a real trained model, WQ and WK would break this symmetry and give meaningful cross-token attention. Identity matrices are just for keeping the arithmetic readable.

Why it's called self-attention

Q, K and V all come from the same source matrix, X. Each token produces a query, a key and a value out of its own representation, and then every token attends to every other token in that one sequence.

That separates self-attention, as used in GPT and LLaMA, from cross-attention, as used in encoder-decoder models like T5, where Q comes from one sequence and K and V come from another.

Practical engineering notes

KV cache. Recomputing K and V for every earlier token at each generation step would be O(n²). A KV cache stores all the previous K and V vectors and reuses them, bringing the cost per step down to O(n) in exchange for memory that grows with sequence length.

Head dimension. The usual configuration is d_k = d_model / num_heads. GPT-2 has 768 / 12 = 64, and GPT-3 has 12288 / 96 = 128.

Grouped Query Attention (GQA). LLaMA-2, Mistral and Gemma let several Query heads share a single KV head. That cuts KV cache memory sharply, which is a big reason these models cope with long contexts on consumer hardware.

Attention in a sentence

There's no magic in it. Attention is a learned, differentiable, soft-weighted lookup table that gets recomputed from scratch for each input, yielding representations that carry the whole relational structure of the text.