One Line Changed: Where Every Symbol in the Attention Formula Comes From

This part ends at a formula:

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

One line in the paper, five in the code. This part starts from the crudest possible approach and earns each symbol in turn — and looking back from the end, the code going from “take an average” to “attention” changes by exactly one line.

Where things stand: hou, you and thou read the same row

Part 23’s bigram has a one-line forward pass:

logits = self.token_embedding_table(idx)

Each position’s prediction depends only on that position’s own character. Try three quite different contexts:

for s in ['hou', 'you', 'thou']:
    logits, _ = m(torch.tensor([encode(s)]))
    print(s, [itos[i.item()] for i in logits[0, -1].topk(3).indices])

All three lines print the same thing. The last number of all three idx is 59 (u), and the bigram sees nothing else:

torch.equal(logits_hou[0, -1], logits_you[0, -1])    # True, not one number different

Passing three characters rather than just u is deliberate — it shows that giving more does not help. The information is right there and the model cannot use it. It does not know whether it is spelling house, you or thou, which is why it writes no real words.

The hard constraint: position t may only see 0…t

What is wanted is “position 2’s output = a blend of h, o and u”, not just u. But there is one hard constraint.

Part 23 noted that a chunk of length 8 is training 8 questions at once: with xb[0] = "Let's he", position 3 sees Let' and must guess s. If position 3 could see position 4’s s, it would be looking straight at the answer.

🔑 Position t may only see positions 0…t. Every scheme below has to respect that. Part 23 closed by saying “the lower-triangular mask exists to stop those T cheating” — this part delivers on that.

First, separate the representation from the scores

The bigram’s nn.Embedding(65, 65) produces a vector that does two jobs: it is both the character’s representation and, directly, the prediction’s scores. That causes three problems: nothing can be inserted in between; the representation is locked into the shape of the answer; and the dimension is tied to the vocabulary (GPT-2’s is 50k, so the representation would be 50k-wide).

Split it:

self.token_embedding_table = nn.Embedding(vocab_size, n_embd)   # (65, 32)  in: character → representation
self.lm_head = nn.Linear(n_embd, vocab_size)                    # (32, 65)  out: representation → scores

The entrance handles “what it is”, the middle handles “look at the context”, the exit handles “guess what’s next”. As a network it goes from one 65 → 65 layer to two layers, 65 → 32 → 65.

The parameter count barely moves: 65×65 = 4225, and split it is 65×32 + 32×65 + 65 = 4225. Splitting buys no improvement by itself; it only makes room.

💡 lm_head is the industry’s common name: lm = language model, head = head. The backbone understands and is usually large; a head answers a specific task and is usually small. Open Qwen’s model structure and its last layer is called lm_head. Swapping in a classification head on the same backbone is the basic idea of fine-tuning.

This also explains why part 23’s “starting loss twice as bad as guessing” changes: the bigram’s logits came straight from an Embedding (standard normal), but now they pass through lm_head, and nn.Linear initialises with U(−1/√32, 1/√32) = U(−0.177, 0.177) — small weights, so the initial logits sit near 0 and the starting loss lands closer to ln 65 = 4.17. What part 14 did by hand with that factor of 0.01, nn.Linear’s default initialisation now does for you.

Scheme one: take the average

The most direct way to “blend every earlier position”:

x̄ₜ = (1/(t+1)) · Σᵢ₌₀ᵗ xᵢ

With three vectors:

x₀ = [2, 7]      ← h
x₁ = [6, 4]      ← o
x₂ = [6, 5]      ← u

position 0: x̄₀ = [2, 7]
position 1: x̄₁ = ([2,7] + [6,4]) / 2          = [4, 5.5]
position 2: x̄₂ = ([2,7] + [6,4] + [6,5]) / 3  = [4.67, 5.33]

Position 2’s new vector now carries information from h, o and u. This is a bag of words — the weakest form of “communication”: order is discarded and every earlier position counts equally.

The vectors are 2-dimensional so they can be checked by hand; in reality n_embd = 32. Averaging acts on each dimension independently, so changing the dimension changes not a word of the code — it is the C in (B, T, C).

The problem is speed: a double for-loop, which at B=64, T=256 means 16,384 Python iterations per step, none of it parallel on a GPU.

Scheme two: a matmul is a weighted sum

★ Derivation: each row of c blends the rows of b

Let c = a · b. By the definition of matrix multiplication, row i, column d of c is

c[i][d] = Σⱼ a[i][j] · b[j][d]

Fold the d dimension up and look at the whole row at once:

c[i] = Σⱼ a[i][j] · b[j]        ← b[j] is b's row j (a vector)

🔑 Row i of c is a weighted sum of b’s rows, and the weights are row i of a. So once the mixing proportions are written into a, a single matmul does the blending for every position at once.

Same three vectors as b, three choices of a:

① all ones = add everything (but breaks “no looking ahead”):

a =           b =         c = a @ b =
[[1, 1, 1],   [[2, 7],    [[14, 16],
 [1, 1, 1],    [6, 4],     [14, 16],
 [1, 1, 1]]    [6, 5]]     [14, 16]]

② lower triangular = add only the past:

a =           b =         c = a @ b =
[[1, 0, 0],   [[2, 7],    [[ 2,  7],     ← x₀ only
 [1, 1, 0],    [6, 4],     [ 8, 11],     ← x₀ + x₁
 [1, 1, 1]]    [6, 5]]     [14, 16]]     ← x₀ + x₁ + x₂

The zeros in the upper triangle set the future’s weights to zero. This follows directly from the derivation: a[i][j] = 0 means b[j] contributes nothing to c[i].

③ normalise each row = take the average:

a =                    c = a @ b =
[[1.00, 0,    0   ],   [[2.00, 7.00],    ← x₀
 [0.50, 0.50, 0   ],    [4.00, 5.50],    ← (x₀+x₁)/2
 [0.33, 0.33, 0.33]]    [4.67, 5.33]]    ← (x₀+x₁+x₂)/3

Identical to scheme one’s hand calculation. One matmul replaces all of scheme one’s loops, and on a GPU it is massively parallel.

wei = torch.tril(torch.ones(T, T))          # (8, 8) lower triangular
wei = wei / wei.sum(1, keepdim=True)        # normalise each row
xbow2 = wei @ x                              # (8,8) @ (4,8,2) → (4,8,2)
torch.allclose(xbow, xbow2)                  # True

Why the shapes fit: (8,8) and (4,8,2) differ in rank, so PyTorch broadcasts wei to (4,8,8) and does one (8,8) @ (8,2) per sequence — a batched matmul.

⚠️ keepdim=True cannot be dropped: sum(1) is (8,) and broadcasts along columns instead, giving a wrong result with no error. Same trap as part 19’s step ③.

The problem: only uniform weights

When predicting what follows hou, the h matters most (it is what makes it house rather than you). Uneven weights are needed, ideally computed from the data. But put arbitrary scores into scheme two:

scores [2, −1, 3]  → divide by row sum 4 → [0.5, −0.25, 0.75]    ❌ a negative weight
scores [1, −1, 0]  → row sum = 0         → division by zero      ❌ it blows up

And to block the future, arbitrary scores in the upper triangle are not necessarily 0 — there has to be a way to force them to 0.

Scheme three: softmax and −∞

raw scores:       [ 2,    −1,     3,     5  ]      ← position 3's 5 is the future
① fill future:    [ 2,    −1,     3,    −∞  ]
② exp:            [ 7.39,  0.37, 20.09,  0   ]
③ divide by sum:  [ 0.27,  0.01,  0.72,  0   ]    sum 27.84

All positive, summing to 1, the future exactly 0 — however large that future score was. And a higher score gets a larger weight.

When every score is 0:

[0, 0, 0, −∞]  → exp [1, 1, 1, 0]  → [0.33, 0.33, 0.33, 0]

🔑 Scheme two is scheme three in the special case where every score is 0. So scheme three both reproduces the uniform average and handles arbitrary scores.

tril = torch.tril(torch.ones(T, T))
wei = torch.zeros((T, T))                           # ① every score 0
wei = wei.masked_fill(tril == 0, float('-inf'))     # ② upper triangle → −∞
wei = F.softmax(wei, dim=-1)                        # ③ softmax each row
out = wei @ x                                       # ④ weighted sum

Softmax itself was part 7. What matters here is its shift invariance (add the same number to every score and eᶜ cancels top and bottom) — the reason part 19’s step ⑧ found dlogit_maxes exactly 0.

★ Normalising a matrix means saying which direction first

The same matrix [[1,1,0],[2,1,1],[3,3,2]] normalised by row (row sums 2, 4, 8) gives a third row of [0.38, 0.38, 0.25]; normalised by column it is a different matrix entirely; normalised as a whole, yet another. Without a direction, “normalise” means nothing.

Attention normalises by row: row t of wei is position t’s scores for every position, and each position is deciding “how do I share my attention out”, so its row sums to 1. By column it would become “the shares of position 0’s being looked at sum to 1” — not what is wanted.

🔑 The test: ask who is allocating, and to whom. The allocator is the row, the allocated-to are the columns → normalise by row. A non-negative matrix whose rows each sum to 1 is called row-stochastic, and wei is one.

The same question of direction runs through standardisation: BatchNorm goes by column (across samples, x.mean(0)), LayerNorm by row (across features, x.mean(-1)). BatchNorm’s going by column is why samples couple and batch size 1 breaks it — setting up why Transformers use LayerNorm.

Position embeddings: attention is a set operation

Attention is a weighted sum, and addition does not care about order. Position 3 of dogs and of gods sees the same pile of vectors {d, o, g, s} and produces exactly the same result. The causal mask only says which positions may be seen, not what order they are in.

self.position_embedding_table = nn.Embedding(block_size, n_embd)       # (8, 32)
tok_emb = self.token_embedding_table(idx)                               # (B,T,C)
pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
x = tok_emb + pos_emb                                                   # (B,T,C)

Exactly the same structure as the token table; only the lookup key is now a position. In miniature (d=[1,0], g=[2,1], pos0=[0.1,0], pos2=[0.3,0]):

dog:  d+pos0 = [1.1, 0.0]   …   g+pos2 = [2.3, 1.0]
god:  g+pos0 = [2.1, 1.0]   …   d+pos2 = [1.3, 0.0]

The same d is [1.1, 0] at position 0 and [1.3, 0] at position 2 — x now encodes both who and where.

Why add rather than concatenate: adding keeps the dimension, so nothing downstream changes; and 32 dimensions is room enough for training to put the two kinds of information along different directions.

This is a learned position embedding (GPT-2’s kind). The original Transformer used sinusoids (a fixed formula, not trained); Llama and Qwen use RoPE (rotary position encoding, applied to q and k). A learned table has only block_size rows and cannot look past them, which is the main reason modern models replaced it — and why generate() gains a line idx_cond = idx[:, -block_size:].

⚠️ After this change the loss not falling is correct (still about 2.49): there is no attention yet, and each position’s logits depend only on “this character + which position”. This step only builds the frame.

One line changed: q and k

Scheme three’s first line, wei = torch.zeros((T, T)), sets every score to 0. Now the scores are to come from content:

wei = q @ k.transpose(-2, -1)

That is the entire difference between attention and taking an average. The mask, the softmax and the multiplication — three lines — do not change by a word.

Position t’s score for position s is qₜ · kₛ. A dot product measures how aligned two directions are: large when aligned, 0 when perpendicular, negative when opposed.

q and k come from the same x, each multiplied by its own matrix:

key   = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)
k = key(x)      # (B,T,32) → (B,T,16)
q = query(x)

Each position computes independently with a shared W; there is no communication between positions yet.

Why two matrices rather than using x directly: with q = k = x the only thing measurable is whether two characters are alike. u is looking for “an initial consonant” while being a vowel itself — two independent matrices are what let “what I seek” differ from “what I have”.

bias=False is a simplification: a bias on k adds the same constant to every score in a row, and softmax is shift-invariant, so it does nothing; a bias on q is not wasted (it adds a fixed score for each position being looked at — worked out in part 26), but switching it off costs little.

★ Derivation: qₜ · kₛ can be written xₜ W_q W_kᵀ xₛᵀ

q and k are row vectors: qₜ = xₜ W_q and kₛ = xₛ W_k. A dot product of two row vectors can be written a · b = a bᵀ, so

qₜ · kₛ = (xₜ W_q)(xₛ W_k)ᵀ
        = xₜ W_q · W_kᵀ xₛᵀ             ← using (AB)ᵀ = BᵀAᵀ
        = xₜ · M · xₛᵀ                  where M = W_q W_kᵀ

Shapes (1×32)(32×16)(16×32)(32×1) = one number.

🔑 M = W_q W_kᵀ is a 32×32 “compatibility table”: M[i][j] is how strongly the asker’s dimension i and the asked’s dimension j attract. W_q or W_k alone means little; their product decides who attends to whom.

M is generally not symmetric, so attention is not symmetric. The numbers from the hand-worked demo below show it directly:

u looking at h: q_u · k_h = 4.0
h looking at u: q_h · k_u = 1.0

u very much wants to look at h, while h is not much interested in u. With W_q = W_k, M = WWᵀ would be symmetric and that one-way interest could never be learned.

✅ Checked both ways: computing q and k then the dot product, and computing M then xₜ M xₛᵀ, both give 4.0 for u looking at h.

Why not learn M directly: comparing after squeezing to 16 dimensions is cheaper, and rank(M) ≤ 16 keeps the degrees of freedom in check — later, with several heads, each judges relevance only in its own small subspace.

wei is not a parameter; it is computed

⚠️ The one variable name wei holds three stages: raw scores → after the mask → after softmax. When reading the code, check the line above it.

The triangular shape comes from the mask; the numbers inside come from content. q @ k.T first computes the full score matrix (future included — on a GPU, computing it all and then masking is faster than computing only the lower triangle), and the mask then sets the upper triangle to −∞.

k.transpose(-2, -1)

Swaps the last two dimensions. q (B,T,16) @ k.T (B,16,T) → wei (B,T,T), and cell [t][s] is exactly qₜ · kₛ.

(-2, -1) rather than .T: negative indices count from the end, so only the last two dimensions are ever swapped and the batch dimension is left alone. When k becomes 4-dimensional with several heads, this line does not change.

wei’s shape goes from scheme three’s shared (T,T) to (B,T,T) — one per sequence.

v: what is handed over once found

Scheme three’s last step was wei @ x: the attended-to position hands over its raw representation as is. A third projection replaces that:

value = nn.Linear(C, head_size, bias=False)
v = value(x)          # (B,T,32) → (B,T,16)
out = wei @ v         # what gets blended is v, not x

The library analogy: the label (k) is how the book is found; what goes home is the contents (v).

🔑 Change W_v and not one number of wei moves; only out does. q and k decide whom to look at and how much; v decides what is received.

Worked through by hand

x is one-hot (T=3, C=3, head_size=2) and the three weight matrices are set by hand to integers, so every step can be checked by hand:

x = torch.eye(3)                        # h o u
query.weight.copy_(torch.tensor([[1., 0., 2.], [0., 1., 0.5]]))
key.weight.copy_(  torch.tensor([[2., 0., 1.], [0., 1., 1. ]]))
value.weight.copy_(torch.tensor([[10., 0., 5.], [0., 10., 5.]]))

u’s output by hand: 0.80×[10,0] + 0.02×[0,10] + 0.18×[5,5] = [8.87, 1.13], nearly v_h itself — u has taken its information from h. Feed a different sequence (you) and the same weights give a different wei — which is what “computed from content” means.

✅ The whole table was recomputed from scratch in pure Python; every cell matches.

Why it is called self-attention: q, k and v all come from the same x — the sequence attends to itself.

Six annotations

With one head written, Karpathy stops to add six annotations. No new model — each answers a “why was it designed this way”.

① Attention is communication on a directed graph

Every node does the same thing: gather the v along the edges pointing at it, weighted by the edges’ thickness. The grey cells of the matrix do not exist in the graph at all.

🔑 The attention mechanism itself does one thing: move information along edges, by weights computed from content. What the graph looks like is decided by the mask line.

② No sense of space

Covered in the position-embedding section above. One comparison to add: convolution knows left from right by construction (one weight for the left neighbour, another for the right), so it needs no position encoding; attention computes weights from content without regard to side, so it must have one. ViT adds a position encoding to every image patch for the same reason.

③ No communication across the batch

wei is (B, T, T): each sequence has its own independent table, with no edges to any other sequence:

x2 = x.clone(); x2[1] = torch.randn(T, C)    # replace only sequence 1
torch.equal(out1[0], out2[0])                 # True — sequence 0 did not change by one number

“But with gradient updates every step, the sequences surely affect one another?” — yes, but that is a different kind of influence:

The second is learning itself; every network has it and must. The difference is only in the first. Part 16’s “the same sample in two batches gives −3.8143 and −4.3548” was the first kind: no parameter moved, only the batch-mates changed, and the output did too.

Why it matters: inference is forward-only. A server batches many users’ requests together — with BatchNorm your answer would depend on whom it happened to be computed alongside. That is one reason Transformers use LayerNorm rather than BatchNorm.

④ Encoder and decoder: also one line apart

The mask exists for one reason only, the task: we are predicting the next character, so position t must not see the answer. Delete that line and every position sees every other:

🔑 The names sound like two architectures; the code differs by one line. This is the part’s second “one line changed”.

Where the names come from: the original 2017 Transformer was a translation model. The encoder reads the whole Chinese sentence at once and “encodes” it — all of it given at once, every character seeing every other; the decoder “decodes” English word by word from that — seeing only what has been generated. The names describe roles in the translation pipeline, and the mask follows from the role. The two halves were later used separately: BERT kept only the encoder, GPT only the decoder.

When a mask is actually needed: in one case only — the answer is inside the input sequence, to the right of the current position.

Cloze prevents cheating not by masking attention but by removing the answer from the input — so attention may freely look both ways.

Most large models today are decoder-only, because understanding can be wrapped as generation (“This film was no good at all. The sentiment is:” → “negative”). But encoders have not gone away: ViT for images and embedding models for vector retrieval are both encoders. Qwen2.5-VL uses both: the image goes through a vision encoder, then to a language decoder.

A counterexample you can run: comment out the masked_fill line in a decoder and the training loss drops close to 0 (each position just copies the answer on its right), while generation is garbage — at generation time there is nothing to the right to copy. That is what “cheating” looks like.

⑤ Self-attention and cross-attention

Translating, to produce English love you need to look at the Chinese 爱 — the asking sequence is English, the one being consulted is Chinese:

q = query(x_english)               # what I am looking for
k = key(x_chinese)                 # what I have
v = value(x_chinese)
wei = q @ k.transpose(-2, -1)      # (T_en, T_zh), not necessarily square

The core problem of multimodality is “how does text look at an image”: q from text tokens, k/v from image features. Some models use explicit cross-attention layers; Qwen-VL-style models concatenate image and text tokens into one sequence and let self-attention have them look at each other.

⑥ Why divide by √dₖ

The real implementation is

wei = q @ k.transpose(-2, -1) * head_size ** -0.5

Measured: with q and k random of variance 1, at head_size = 4, 16, 64, 256 the variance of wei comes out at about 4, 16, 64, 255. The variance equals the dimension.

★ Derivation: Var(q · k) = d

A dot product is a sum of d products:

q · k = Σᵢ₌₁ᵈ qᵢ kᵢ

Step one: the terms are independent, so their variances add (part 15 derived that rule):

Var(q · k) = Σᵢ Var(qᵢ kᵢ)

Step two: each term. qᵢ and kᵢ are independent with mean 0 and variance 1:

Var(qᵢ kᵢ) = E[(qᵢ kᵢ)²] − (E[qᵢ kᵢ])²
           = E[qᵢ²] · E[kᵢ²] − (E[qᵢ] · E[kᵢ])²       ← independent, so the expectations split
           = 1 · 1 − (0 · 0)²
           = 1

Step three: add up d ones:

Var(q · k) = d

So dividing by √d brings the variance back to 1:

Var( (q · k) / √d ) = d / d = 1

✅ Checked by simulation in pure Python (20,000 draws per dimension): at d = 4, 16, 64, 256, Var(q·k) is 3.91, 15.98, 63.76, 258.20, and after dividing by √d it is 0.978, 0.999, 0.996, 1.009.

🔑 This is part 14’s formula used in a new place: there, the output variance of x @ W is fan_in × the input variance, so weights are scaled by 1/√fan_in. Here, the output variance of q @ kᵀ is head_size × 1, so scores are scaled by 1/√head_size. Exactly the same thing.

Without it: softmax turns into one-hot

The same scores, scaled up to the matching standard deviation:

s = [0.1, −0.2, 0.3, −0.2, 0.5]
softmax(s)      → [0.19, 0.14, 0.24, 0.14, 0.29]    fairly even
softmax(s × 4)  → [0.11, 0.03, 0.25, 0.03, 0.56]    ← head_size=16, standard deviation ×4
softmax(s × 8)  → [0.03, 0.00, 0.16, 0.00, 0.80]    ← head_size=64, looking at almost one position

Two consequences: information comes from a single node — what should blend the past just copies one position; and the gradient vanishes — near one-hot softmax saturates, and changing the scores barely changes the probabilities. The same kind of problem as part 14’s saturated tanh. And it is worst at initialisation: before learning whom to look at, the model is randomly locked onto one position.

Every symbol in the formula

Back to that opening line:

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

The “Scaled” in the paper’s Scaled Dot-Product Attention is that √dₖ. Every symbol in the formula is now accounted for.

Looking back: from scheme three’s average to attention, the code changed one line, wei = torch.zeros((T, T)) → wei = q @ k.transpose(-2, -1), plus x becoming v. The mask and softmax lines have not moved since scheme three.

Next the head goes into the model with its scaling, and the loss finally starts coming down from 2.49.