What writing attention by hand taught me

Before I built anything on top of a large language model, I spent five weeks writing a small one. I followed the route in Sebastian Raschka's Build a Large Language Model (From Scratch) and typed every piece myself, with a note beside each one saying why it was there. The part that took longest was attention, and I wrote it four times.

I had read about attention many times before that, and I could have drawn the diagram. This post is about what the diagram had not told me.

Version one: no weights at all

The first version had nothing to train. Take the vector for each token, dot it with the vector for every other token, run a softmax over the scores, and use the result as weights for a sum.

That was enough to see what attention is for. Each token's new vector is a blend of all the others, mixed in proportion to how similar they are. With no weights, "similar" just means the embeddings point the same way, which is not very useful. But the mechanism is all there, in three lines.

Version two: something to learn

The second version added the query, key and value projections. Now a token does not compare its raw vector with its neighbours. It asks a question with one projection, advertises itself with a second, and hands over content with a third, and all three are learned.

This is where I stopped thinking of attention as a similarity lookup. The model learns what to look for, what to be found by, and what to pass on, separately.

Version three: the mask

A model that predicts the next word must not see it. So every score above the diagonal is set to minus infinity before the softmax, and comes out as exactly zero after it.

Two details here came only from writing it. The first is the order. I did it the other way round first: softmax, then zero the weights above the diagonal. The rows no longer sum to one, so they have to be renormalised. Minus infinity before the softmax does both jobs in one step.

The second is where the mask lives. A plain tensor on the module stays on the CPU when the model moves to the GPU. A buffer does not: it is not a parameter and is not trained, and it travels with the model.

self.register_buffer(
    "mask",
    torch.triu(torch.ones(context_length, context_length), diagonal=1),
)

Version four: heads, and the reshape

My first multi-head attention was the obvious one: twelve small attention modules in a list, run one after another, outputs glued together. It was correct and slow.

The real version does one large projection and splits it afterwards. This is the hardest code in the project, and none of it is conceptually hard. It is four-dimensional bookkeeping:

keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
keys = keys.transpose(1, 2)
attn_scores = queries @ keys.transpose(2, 3)

For GPT-2 small, the shapes go like this:

Once the heads sit beside the batch, one matrix multiplication serves all twelve. What I took from it is that a head is not a separate module. It is a 64-number slice of one big projection, and the only thing that makes the slices independent is that the softmax runs over each one separately.

I got through it by writing the expected shape as a comment on every line and printing the real one to check.

The parameter count that was wrong

When the full model was assembled I counted its parameters and got 163.0 million. GPT-2 small is meant to have 124.

The difference is one matrix. GPT-2 uses the same 50,257 by 768 table to turn tokens into vectors on the way in and vectors into token scores on the way out. My model kept two copies. Count the table once and the total is 124.4 million. It was the first time the number in a model's name meant something to me.

The test I could not argue with

OpenAI released GPT-2's weights in 2019. If my architecture really was GPT-2, their weights should fit it, tensor for tensor.

They store the query, key and value projections as one fused matrix, and I had written three, so each block's matrix had to be cut into thirds. Their linear layers keep weights the other way round from PyTorch's, so every matrix is transposed on the way in. Their attention has bias terms that mine did not. Every copy went through a small function that refuses to assign a tensor whose shape does not match, so a mistake would stop the load and could not become a quietly wrong model.

Close to two hundred tensors went in with no shape left over. I gave it the prompt "The purpose of life is" and it replied "as important as the pursuit. What does the need for happiness make of living". That is not wisdom. It is fluent English from code I had written line by line, and it is the result from that summer I care about most.

What it changed

Everything I have built since sits on models far larger than this one, all through APIs. I do not write attention at work and I do not expect to.

But when a context window fills up, I know that it is a table of scores growing with the square of the length. When a temperature setting changes the answers, I know which division it is, because I wrote that too. When generation is slow, I know my version recomputed attention over the whole sequence for every new token, because I never wrote a key-value cache, and I know what that cache is for. None of that came from the diagram.