Fine-Tuning at Home Part 04

Fine-Tuning at Home, Part Four: Attention

Take the sentence "the server refused the connection because it was overloaded". Something has to work out that "it" means the server rather than the connection. Attention is how, and it is the only place in the entire architecture where one token can affect another.

Why This Is Hard

Start with why this is hard, because the difficulty is easy to miss.

After part three, each token is a vector, and that vector came out of a lookup table. The vector for "it" is the same vector every time "it" appears anywhere, in any sentence, about anything. It carries the general-purpose meaning of the word and nothing about the situation. If nothing further happened, the model would be trying to predict the rest of the sentence from a bag of context-free word meanings, which is roughly the position translation software was in for decades and is why it was bad.

What is needed is a way for the vector at each position to absorb information from the other positions, so that by the time it reaches the top of the stack, the vector sitting at "it" is no longer the generic pronoun. It has become something closer to "the server, which refused a connection". Attention is the operation that does that absorbing, and it is the only one in the architecture that does.

Three Vectors from One

Every token's vector produces three new vectors, by being multiplied through three separate learned tables. They are called the query, the key and the value, and the names are a database metaphor that half helps and half misleads.

ONE TOKEN'S VECTOR x [4096] the vector for "it", as it enters this layer W-query W-key W-value q [128] k [128] v [128] QUERY what this token is looking for KEY what it offers to anyone looking VALUE what it hands over if taken up The three tables are learned during pretraining, and each attention head has its own set of them.
Nothing is retrieved from anywhere. All three are the same input vector pushed through three different learned tables, which is worth holding onto when the database metaphor starts to feel like it is explaining more than it is. The query is what this token is looking for. The key is what it offers to anyone looking. The value is what it hands over when somebody takes it up.

The reason there are three rather than one is that a token plays three different roles simultaneously, and they call for different information. What "it" needs to go looking for is not the same as what "it" advertises to other tokens, and neither is the same as the content it would contribute if something selected it. Three separate learned projections let the model keep those roles apart, and the tables that produce them are learned like everything else, which means the model worked out for itself what each role should contain.

Every Query Against Every Earlier Key

Now the comparison. The model takes the query of each token and compares it against the keys of every token before it, one pair at a time, and each comparison produces a single number saying how good the match is.

The comparison itself is a dot product: multiply the two vectors together element by element and add up the results. Two vectors pointing in similar directions produce a large number; two pointing in unrelated directions produce something near zero. That is the entire matching operation, and its simplicity is the reason this scales at all. The result is then divided by a fixed factor related to the vector width, which sounds like a footnote and is not: without it the numbers grow large enough that the next step becomes numerically unstable, and the model trains badly or not at all.

Arrange every one of those comparisons and you get a grid.

ATTENTION SCORES: EVERY QUERY AGAINST EVERY EARLIER KEY the server ref... the conn... bec... it was keys, looked at queries, doing the looking the server refused the connection because it 0.61 0.12 was the row for "it" scores highest against "server", not "connection" dashed cells are masked out: a token cannot see the future Each row is turned into weights that sum to 1, and the token's new vector becomes a blend of the values, dominated by the best match. The grid is n by n, so doubling the input roughly quadruples this work. That is why long context is expensive, and why it took years to make it cheap.
Read a row as one token looking backwards. The row for "it" scores highest against "server", so the vector for "it" gets pulled toward carrying the server's information. The pronoun has been resolved, not by a rule anybody wrote, but by an arithmetic match between tables that were learned. Each row is normalized to sum to one, and the token's new vector becomes a blend of the values weighted that way.

Walking one row through slowly, because the mechanism is worth seeing once at full speed. The query vector for "it" is compared against the key vectors of every earlier token. The comparison with "server" comes out high. The comparison with "connection" comes out lower. Those raw numbers are turned into weights that sum to one, so "server" might end up at 0.61 and "connection" at 0.12, with the rest spread thinly. The new vector for "it" is then the sum of every token's value vector multiplied by its weight, which means it is dominated by the server's contribution.

Nothing selected the server. No branch was taken, no comparison was thresholded, no decision was made in any sense a programmer would recognize. The weights are continuous and every token contributes something. What happened is that one contribution was much larger than the others, and the result is a vector that carries mostly the server's information, blended with a little of everything else.

Two properties of that grid decide a great deal about how these models behave in practice, and both are visible in the figure rather than buried in the math.

The upper half is masked, so a token can only see backwards. That is not an optimization, it is the definition of the model. Before the weights are computed, every score for a future position is forced to negative infinity, so it comes out of the normalization as zero. This is why the whole thing runs left to right, why generation happens one token at a time, and why the model cannot revise something it already said: the text it produced is now part of what it is conditioning on, permanently. It is also why the same model can be trained efficiently on whole documents at once, since every position can be predicted in parallel while still only seeing its own past.

The grid is n by n, so cost grows with the square of the length. Every token compares against every previous token, so doubling the input roughly quadruples this work. This is why long context was hard for years, why it is still expensive, and why a long prompt is not linearly more expensive than a short one. Nearly all of the engineering effort in modern serving stacks is aimed at this one property, and most of the headline context-length improvements of the last few years are engineering around it rather than removing it.

Modern models share keys and values between heads to shrink it. Called grouped-query attention. Several query heads share a single set of keys and values rather than each having its own. It exists to reduce the memory described in part five rather than to improve quality, and it is why a model with thirty-two query heads may have only eight sets of keys. You will meet those two numbers disagreeing in a configuration file and wonder whether something is broken; that is the reason, and it was one of the practical changes that made long context affordable.

Many Heads, In Parallel

All of that happens in parallel, many times over, on different learned projections. Each copy is a head. A model of this size might have thirty-two of them in every layer, each with its own query, key and value tables, each producing its own grid and its own blended output.

Their outputs are concatenated and passed through one more learned table that mixes them back into a single vector of the original width. That mixing step matters: without it the heads would be thirty-two separate channels that never interact.

input [n, 4096] head 1 its own W-query, W-key, W-value, its own score grid head 2 may end up tracking which entity is being discussed head 3 may end up tracking grammatical agreement ... head 32 and many do things nobody has a name for concat and mix [n, 4096] Nobody assigns a head its job. They specialize because the training had no reason to make them all do the same thing.
Heads specialize without being told to. Nobody assigns head 3 the job of tracking agreement; it drifts into something like that because training had no reason to make every head do the same thing. Interpretability work has found heads doing jobs of this kind, and the labels here are illustrative rather than a claim about any particular model.

The reason to have many heads rather than one big one is that a single set of weights has to commit to one notion of relevance. Thirty-two sets can each commit to a different one, and a token can simultaneously attend to its grammatical subject, the entity under discussion, and the token three positions back, without those needs competing for the same weights.

Never Building the Grid at All

The grid in that figure has a practical problem that took years to solve properly, and knowing about it explains a setting you will be told to turn on in part nine.

Written naively, computing attention means building that whole grid in memory. For a thousand tokens that is a million numbers per head per layer, which is manageable. For a hundred thousand tokens it is ten billion per head per layer, which is not. The quadratic cost is a memory problem before it is ever a compute problem, and for a long time that memory ceiling, rather than arithmetic speed, was what capped context length.

The fix, which goes by the name FlashAttention, is to never build the grid at all. Instead of computing every score, storing them, normalizing and then blending the values, it works through the sequence in tiles that fit in the small fast memory attached to the processor, carrying running totals as it goes and keeping only the result. The full grid exists conceptually and never exists as data.

Two things about it are worth being precise about, because it gets described as an optimization in ways that suggest a tradeoff that is not there.

It computes the same answer, not an approximation. This is the part people assume otherwise. There are approximate attention methods that trade accuracy for speed and this is not one of them; the output is mathematically the same as the naive version, to floating-point tolerance. What changed is the order of operations and what gets written to memory, which is why it can be turned on with no quality argument to have.

The win is memory traffic rather than arithmetic. A graphics card can do far more arithmetic per second than it can move numbers between its main memory and its processing units, so a great many workloads are limited by movement rather than by calculation. Attention written naively spends most of its time moving that giant grid around. Avoiding the round trip is where the speedup comes from, and it is the same principle behind most serious performance work on these machines.

It is a kernel, so it has hardware and version constraints. Because it is hand-written code targeting specific hardware, it has to be built for your setup and it does not support every card or every configuration. In part nine it is the one package in the install that compiles from source and takes a long time, and it is the one that is genuinely optional: drop it, remove the line from the configuration, and everything still runs a little slower and produces the same result.

What Attention Does Not Do

It is worth being precise about what this mechanism does not do, because attention gets described in ways that imply far more than is happening.

It is not the model deciding what is important. The word invites that reading and it is misleading. There is no evaluation of importance and no goal being pursued. There is a dot product between two vectors produced by tables that were adjusted, over trillions of tokens, in whatever direction reduced prediction error. The result often looks like judgment. The mechanism contains none.

It does not happen once. The figures here show one layer. A model of this size does all of it about three dozen times over, and the vectors that enter layer two are the ones layer one already modified. Resolving "it" is not a single event at a single layer; it is something that emerges across many of them, and interpretability researchers can rarely point at one place where it happened.

It carries no notion of order by itself. A dot product between two vectors is the same regardless of which positions they came from. Order has to be injected separately, as part three mentioned, by schemes that modify the vectors according to position. Remove that and attention would treat your sentence as an unordered set, which is a useful thing to know when reading about models whose context lengths were extended: what was usually extended is the positional scheme, not the attention.

What This Means for Your Prompts

Three things follow from the score grid that are worth having if you never touch a model's internals again.

Position affects how much notice a token gets. Attention weights are not uniform across a long input and the pattern is not flat. Material at the very start and the very end of a long context tends to be attended to more reliably than material buried in the middle, an effect consistent enough to have been studied under that description. The practical version: in a long prompt, instructions buried halfway down are the ones most likely to be underweighted, and repeating a critical constraint at the end is not superstition.

A long prompt is not linearly more expensive. Doubling the input roughly quadruples the attention work, so the cost of a very long context is worse than intuition suggests. This is why stuffing an entire document into a prompt when a relevant section would do is expensive in a way that is easy to miss, and why retrieval, which part seven recommends for knowledge, is an efficiency argument as well as an accuracy one.

The model cannot look forward, so ordering is a real decision. A token can only attend backwards. Anything the model needs in order to interpret your request has to appear before the request, not after it. Putting a document after the question about it is a genuinely worse arrangement than putting it before, and this is one of the few prompt-structure claims that follows directly from the architecture rather than from folklore.

Why This Is the One to Know

Attention is the entire mechanism by which the start of your prompt can affect the end of the answer. Everything else in the architecture processes each position on its own, which is why a reader who understands only this figure still understands most of what matters about context.

It is also the source of nearly every practical constraint you will meet later. The cost of long inputs comes from the size of that grid. The cache that fills your graphics card, coming in the next part, exists because the keys and values in that grid are worth keeping rather than recomputing. And the fact that a model cannot take a sentence back is the mask, seen from the outside.

Next is what surrounds attention. It turns out to be half of a block, the block is repeated a few dozen times, and two pieces of plumbing that look like footnotes are the reason any of it trains at all.