Dr. Ibrar Ahmed

HomeAIArticle

AI Mechanics

Transformer Attention Explained: Query, Key, Value

Dr. Ibrar Ahmed12 min readFrom the lecture notes

What you will learn

  • Attention starts from embeddings, so it is worth a quick refresher.
  • So back to the equation, because there are really only four moving parts in it.
  • Before the matrices arrive it is worth pausing to check where we have got to.
  • So where do query, key and value actually come from?
  • Softmax takes a list of numbers and turns it into a set of shares, rather like cutting up a pie.
01

The problem

A model reads context. But how? Down by the river bank. By the time you reach bank, river has already settled which meaning you intended, and you never had to think about it. Same problem, new sentence. Apple raised its price target. Which Apple? You knew instantly, because the words beside it did the work. The model gets no such help: it starts with exactly the same vector either way.

You might reasonably think that capitalization gives it away, so here are two sentences where the token really is identical. Same spelling, same lowercase, same leading space, which means the same ID and the same row of the embedding table. Before position has been added, and before a single block has run, the fruit and the company are literally the same seven hundred and sixty eight numbers.

The operation that fixes that is the one on screen, and almost every large language model you have ever used is built around it. Every word looks at every other word, works out what actually matters, and then rewrites itself accordingly. Right now that line probably looks like noise, which is fine, because by the time we are finished every symbol in it will make sense.

So we are going to build it piece by piece, and one sentence carries the whole video. Every architecture number you see is a real GPT two figure, and wherever I show attention values I will mark them on screen as illustrative.

02

embeddings

Attention starts from embeddings, so it is worth a quick refresher. The table is one large matrix, with a row for every token in the vocabulary and a column for every embedding dimension. When token four thousand two hundred arrives, nothing gets multiplied and nothing gets computed. You jump to that row and read it.

What makes those rows useful is where they sit relative to one another. Apple, orange and banana end up close together because they appeared in similar contexts during training, and for the same reason cat sits near dog, and king sits near queen. Nobody labeled any of that. The model worked out on its own that words used in the same way ought to point in the same direction.

But here is the ceiling. One token gets one row, and that row was fixed long before your sentence existed, so a word carrying two meanings still only ever gets one vector out of the table. Whatever fixes that has to happen after the lookup, which tells us what we are aiming for. Take the vector for apple and move it, starting from the same row both times and letting the neighbors decide where it ends up.

03

qkv

So back to the equation, because there are really only four moving parts in it. Three matrices, one number you divide by, and softmax. Once you know what those four things are you can read the entire line.

Every token gets a query, a key and a value. All three are vectors, built from the token's own representation, but none of them is the embedding. Think of walking into a library. The query is what you are looking for. The key is the label on the spine. And the value is what is actually inside the book.

So let us watch two words talk to each other. The word its needs context. It takes its query vector and asks the rest of the sentence who has something relevant to offer. Apple answers with its key vector. This is what I am about. You compare those two vectors, and what comes back is a single number. Here it is thirty three point six. A high number means highly relevant to this token, right now.

That comparison is a dot product, which sounds rather more intimidating than it is. You multiply the two lists position by position, add up everything you get, and you are left with one number. The reason it measures relevance is that a dot product comes out large when two vectors point in a similar direction and small when they do not, so if the query for its and the key for Apple happen to be pointing the same way, the score comes out high. Alignment turns into a number.

04

recap

Before the matrices arrive it is worth pausing to check where we have got to. The query asks a question, the key describes what is on offer, and the dot product measures how well those two match up. Value has not done anything at all yet, and nothing has actually been mixed together, which is worth holding on to, because it becomes the whole point in a few minutes.

05

qkv

So where do query, key and value actually come from? Three weight matrices, W q, W k and W v, and you get them by taking whatever vector is coming into the layer and multiplying it by each one in turn. These matrices are trainable parameters, which means they start out random and gradient descent sets them, exactly like every other weight in the network. That is the part people tend to skip over. Nobody ever told the model what a query is. It discovered that asking in this particular way lowered the loss.

Putting real shapes on it makes it concrete. Five tokens, and GPT two small is seven hundred and sixty eight wide, so X is five by seven hundred and sixty eight. Multiply by W q, which is seven hundred and sixty eight by sixty four, and Q comes out five by sixty four. Sixty four numbers for every token, and the same for K and for V. The weight matrices do not depend on position, so every token gets the same W q. And sixty four next to seven hundred and sixty eight is not a typo, because this is only one head.

06

scores

Now you do every pair at once. Q times K transposed gives you a five by five matrix in which every single entry is one dot product, which means one matrix multiplication just performed all twenty five comparisons, and that is precisely why this runs so well on a GPU. Look carefully at what came out, though, because it is square but it is nowhere near symmetric. The word its scores Apple at thirty three point six, while Apple scores its at fourteen point four, since asking is not the same thing as answering. And the way to read it is one row at a time.

It is worth seeing what happens if you skip the next step entirely. Push the raw scores straight into softmax and one token takes essentially everything, leaving nothing for the rest, and once that happens the gradients flatten out and there is very little left to learn from. Now divide every score by the square root of d k, which here is sixty four, so you are dividing by eight. Nothing has been reordered and the ranking is identical, but what you have now is a distribution you can actually train on, and that is the entire reason the scaling factor exists.

07

softmax

Softmax takes a list of numbers and turns it into a set of shares, rather like cutting up a pie. You raise e to the power of each number and then divide by the total, and two useful things fall out of that. Everything comes out positive, and everything adds up to one, which means you can read the result directly as how much of my attention goes here, and that happens to be exactly what a weighted average needs.

Softmax does have a temper, though, because it is exponential and it reacts hard to gaps. Feed it one, two and three and you get a reasonably healthy split across all three. Feed it ten, twenty and thirty and the whole thing collapses onto a single token, which in practice means a word ignoring the entire rest of the sentence. Once the scores get large, softmax saturates, the gradients shrink toward nothing, and learning quietly stalls.

If there is one idea to hold on to from this whole section, it is that a row is a single token asking where it should look. Take the row belonging to its. Sixty eight percent of its attention lands on Apple, and the remainder gets divided between itself and the other words. That is genuinely worth looking at, but it is not proof of how the model arrived at its answer, and I would be careful with anyone who tells you otherwise.

08

interrupt

It is worth noticing that so far, nothing has actually moved. Every number we have computed has been about relevance and nothing else, because query and key exist purely to decide who matters and they never reach the output at all. Value is the thing that actually moves, and that is the step we are going to take next.

09

value

So V finally gets to do something. You multiply X by W v and you get one value vector per token. Query and key existed purely to work out the weights, whereas value is the content, the part that genuinely gets passed along to the next layer. That separation is deliberate, because it lets a token advertise itself one way through its key and then hand over something quite different through its value.

Now we put it together for a single token. The row for its says sixty eight percent Apple, twenty five percent itself, and a small amount from everything else, so you take that share of each token's value vector and add them all together. Watch the pieces arrive, because Apple's contribution is visibly the largest one. The sum of all of that becomes the new vector for its, and this is the exact moment where its stops being a generic pronoun and starts carrying Apple around with it.

All of that weighted summing turns out to be another matrix multiplication, attention weights times V, producing one new vector per token with the same number of rows that went in. Nothing about the shape tells you anything happened at all, and yet every row is now a blend of the others. Static vectors went in, and contextual vectors came out.

10

milestone

That is self attention, complete, and it comes down to five steps. You project X into query, key and value, multiply Q by K transposed, divide by the square root of d k, run softmax across each row, and multiply by V. It is worth stopping there for a moment, because that was genuinely the hard part and you are through it.

11

evidence

Everything up to here has been illustrative, so here is the real GPT two small running on the same sentence. One head does put most of its weight on Apple. But look at the last token, where every earlier word is equally available. Sixty one percent of the attention still lands on the very first token, when a head with no preference at all would give it sixteen. That is an artifact called an attention sink, and it is not the model understanding anything. The mechanism is what you just learned. The tidy numbers were the teaching aid.

12

causal

Now for the version language models actually use. When you train a decoder, the label is simply the same text shifted one step to the right, so position one predicts position two, and so on down the sequence. The problem is that the attention we just built lets position one look at every token, including position two, which means it can read the answer it is being graded on. Train it that way and it reads the answer straight out of a future position, so the training loss collapses without the model ever learning to predict anything. Trivial to train, useless the moment it has to generate.

The fix is a mask. Anywhere a token would be looking forward, above the diagonal, you put negative infinity, and then you add that to the scaled scores. The past comes through untouched because adding zero changes nothing, and it is worth noticing where in the sequence this happens, which is before softmax rather than after. As for why negative infinity rather than zero, it is because zero is a perfectly respectable score that would still take a share, whereas e to the power of minus infinity is zero, so nothing gets through at all.

Here is what comes out, row by row. The first token has nowhere to look except itself, the second divides its attention between the first and itself, and only the final row sees the entire window. Each row still adds up to one across whatever positions remain, which means the mask simply dissolves inside the softmax. Every token attends backward and to itself, and never forward.

There is a cost to all of this. The matrix is n by n, where n is your context length, so eight tokens gives you sixty four cells, sixteen tokens gives you two hundred and fifty six, and thirty two tokens gives you over a thousand. Double the context and you quadruple the number of scores. At GPT two's context length of one thousand and twenty four that works out to over a million scores, per head, per layer, and that is why long context gets expensive so quickly.

13

heads

One more upgrade, because a single attention operation only ever has one opinion. One set of projections has to decide relevance for everything at once, which is a great deal to ask of one softmax. So instead of one large attention you run several smaller ones side by side, each watching for something different, and then pool what they found.

This is where that sixty four finally pays off. GPT two small runs twelve heads, and each one projects down to sixty four, so every head runs the complete attention operation we just built. Twelve heads at sixty four each comes back to seven hundred and sixty eight in total, which means you are not spending any more than a single full width head would have cost. You are splitting the same budget twelve ways.

14

architecture

So where does this sit architecturally? The original transformer has two stacks. The encoder attends with no mask, while the decoder begins with masked attention, the causal version we just built, then a second block called cross attention that reads from the encoder. GPT style language models use the decoder only form, so there is no encoder stack and no cross attention at all.

Which brings us to where attention actually lives. Inside a single block you get layer norm, then masked multi head attention, then a residual, then layer norm again, then the feed forward network, and then another residual. The shape going into that block matches the shape coming out of it, and that is the entire reason you can stack twelve of them on top of each other.

After the final block, GPT two applies one last layer normalization, and then the model width is projected out to the whole vocabulary to produce the logits, which softmax turns into a probability distribution over the next token. That projection shares its weights with the embedding table, the same matrix used in both directions.

15

ending

So, Apple raised its price target. One static row goes in, attention runs, and what comes out the other side is a vector that knows what sentence it is in. The query asks, the key answers, and the dot product scores the match. Scaling keeps softmax sane, softmax produces the weights, and value carries the content. Read that equation again now.

16

next

Without a cache, every new token would recompute the keys and the values for the entire prefix. Compute them once. Keep them. Reuse them. That is the KV cache.