Attention

The problem a feedforward layer can't solve

Take the sentence: “The trophy didn't fit in the suitcase because it was too big.” What does “it” refer to, the trophy or the suitcase? Any fluent English speaker says the trophy, instantly. Now change one word: “because it was too small,” and “it” flips to the suitcase. Same grammatical structure, same position in the sentence, opposite answer, because the correct resolution depends on meaning carried by a word all the way at the other end of the sentence.

The neural network layers from the previous article can't do this. A feedforward layer processes each position with the same fixed weights every time, it has no built-in way to look at a different position and ask “what did that word mean, and is it relevant to me.” Earlier architectures tried processing a sentence one word at a time in order, recurrent neural networks, carrying a running summary forward. That summary had to compress everything seen so far into one fixed-size vector, and by the time a long sentence reached its end, details from the beginning had faded, the same way a rumor loses detail after passing through ten people.

Attentionreplaces that compressed summary with direct access. Every token gets to look at every other token in the sequence, regardless of how far apart they are, and decide, numerically, how relevant each one is. “It” can reach directly back to “trophy” or “suitcase,” whichever one the rest of the sentence actually points to, with no fading and no fixed compression bottleneck in between. This one mechanism, more than any other single idea, is why transformers replaced recurrent networks for language.

Query, key, and value: three questions per token

Every token's embedding gets multiplied by three separate, learned weight matrices, exactly the matrix multiplication from the previous article, producing three different vectors per token: a query, a key, and a value. A useful, if imperfect, analogy: the query is what this token is looking for, the key is what each token, including itself, has to offer, and the value is what that token actually contributes once it's chosen as relevant.

Computing attention is three steps, and the first one should look familiar. Take the dot product of one token's query with every token's key, the exact same operation the Embeddings article used to measure how close two vectors are. A high dot product means “this key matches what I'm looking for.” That gives one raw compatibility score per pair of tokens.

Second, turn those raw scores into a proper probability distribution with softmax, the same reweighting function from the very first article in this series, the one that turned raw model scores into next-token probabilities. Here it turns raw compatibility scores into attention weights that sum to exactly 1 across every token in the sequence.

Third, use those weights to take a weighted sum of every token's value vector. A token that scored a high attention weight contributes most of its value to the result, a token that scored near zero contributes almost nothing. That weighted sum is the output, a new vector for the query token that now carries information pulled in from wherever in the sequence it turned out to be relevant.

1import math
2
3def attention(query, keys, values):
4    d_k = len(query)
5    scores = [dot(query, k) / math.sqrt(d_k) for k in keys]
6    weights = softmax(scores)
7    output = [0.0] * len(values[0])
8    for w, v in zip(weights, values):
9        for d in range(len(v)):
10            output[d] += w * v[d]
11    return output, weights

Tracing the arithmetic by hand

Below is a real, checkable run of the three steps above on the sentence “The cat sat,” using small 4-dimensional toy vectors and fixed weight matrices, chosen so every number is easy to verify by hand rather than learned from actual training data. Pick which token is doing the “looking,” the query, and watch the attention weights and output vector update.

Notice what happens with “sat” as the query. Its toy vector shares one dimension with “The” and a different dimension with “cat,” so it ends up splitting its attention almost evenly between them, roughly 27% each, while still keeping the largest share, about 45%, on itself. Nothing about that split was hand-scripted, it falls directly out of the dot products between the actual vectors, exactly the same arithmetic that would run on real learned embeddings at a much larger scale.

One head isn't enough, and the cost that creates

A single set of query, key, and value matrices can only learn one notion of “relevant.” Real transformers run several of these in parallel, each with its own independently learned Q, K, and V matrices, called attention heads. One head might end up specializing in pronoun resolution, another in tracking which verb goes with which subject across a long clause, another in something with no clean linguistic name at all. Their outputs get concatenated and combined afterward. Nobody assigns heads their specialty, it falls out of training the same way embedding clusters did.

This power isn't free. Computing attention means comparing every token's query against every token's key, a sequence of length n requires roughly n² comparisons. Double the length of the text a model reads at once, and the attention computation roughly quadruples. This is the direct, mechanical reason context windows have real limits and real cost, not a policy decision, a consequence of the arithmetic covered in this article, explored in full in the Context Windows and KV Cache article later in this series.

Attention alone is just a weighted average, though, a linear combination of value vectors. Recall from the neural network article that stacking linear operations without anything nonlinear in between collapses into one linear operation. Attention needs a feedforward layer, activation function included, sitting right after it in the same block to actually build something new. How those two pieces, attention and feedforward, get assembled together with residual connections and normalization into the repeating unit used throughout a real model is the transformer block, covered next.