Transformers & Attention: How Words Look at Each Other

Read these two sentences: I sat on the river bank. I put money in the bank.

You knew which bank was which before you reached the full stop, and you probably didn’t notice deciding. A language model has to decide too, and it starts at a disadvantage. When the word “bank” arrives, it arrives as a token ID (Part 2), and it’s the same ID in both sentences. Whatever turns that number into “the muddy edge of a river” in one sentence and “a place that keeps your money” in the other has to happen inside the model.

The trick is the one you just used without noticing, made precise enough for a GPU to do for thousands of words at once. It’s called attention, and the architecture built around it is the transformerTransformerNeural-network architecture built from attention and feed-forward layers. The "T" in GPT., the T in GPT (Generative Pre-trained Transformer). Nearly every LLM you’ve used is one. By the end you’ll have a tiny one whose every number you can check by hand.

Same word, same numbers

Start where the model starts. Its first step is to look each token ID up in a table with one row per vocabulary token. Each row is a list of numbers called an embeddingEmbeddingA vector of numbers representing the meaning of a token or a piece of text., a vector that stands for the token’s meaning; Llama 3 8B’s table has 128,256 rows of 4,096 numbers. The numbers are learned in training (Part 5), and tokens used in similar ways end up with similar vectors.

But a lookup is a lookup. The row for “bank” is the row for “bank”, whatever sentence it came from, so after this first step the two banks are identical. The best the table can do is store a compromise between the meanings.

To watch what happens next, I’ll use a toy with four numbers per token instead of 4,096, and give the numbers names so we can read them: how much the word is about nature, how much about finance, whether it’s a content word (as opposed to “the” or “on”), and whether it’s ambiguous. “river” is (1, 0, 1, 0). “money” is (0, 1, 1, 0). “bank” is (0.5, 0.5, 1, 1): half of each meaning, flagged as ambiguous. (Two more slots join these in Fig. 4.) Real embeddings don’t come with labels. Their meaning is smeared across thousands of numbers, in directions nobody chose.

Letting the words look at each other

So do what you did: look at the other words. If “river” is nearby, lean toward the riverside meaning; if “money” is nearby, lean the other way. Turned into arithmetic, it goes like this:

1. “bank” gives each word so far, itself included, a score for how relevant it is.
2. The scores become attention weights that add up to 1. (These change with every sentence, unlike the model’s learned weights from Part 3, which stay fixed.)
3. “bank” takes the weighted average of what those words have to offer, and adds it to its own vector.

That is attentionAttentionMechanism that lets each token weigh and mix in information from other (earlier) tokens.: each token looks at itself and the tokens before it, and mixes in information from the ones that matter most.

In the river sentence, “bank” gives 74% of its attention to “river” and spreads the rest thinly. The mix it pulls in is mostly river’s meaning, so its vector swings off the 45° diagonal down to 22°, into riverside territory. Switch to the money sentence and the same learned weights, doing the same arithmetic, hand 74% to “money” instead and send “bank” to 68°. The lookup was identical in both. The context did all the work.

The third sentence is the strange one. “river” is right there, and “bank” can’t see it. In a GPT-style model, each token may look only at itself and the tokens before it, so “bank” has just “The” and itself to work with, and it stays on the fence. The ambiguity gets resolved three words later instead, at the “river” position, which can look back at “bank”. Why build in this blind spot on purpose? Hold that thought. It’s what keeps next-token training honest, and lets it run in parallel.

Attention, one number at a time

How did “bank” decide that “river” deserves 74%? Each token’s vector gets multiplied by three learned matrices, producing three new vectors with three different jobs:

• a query (Q): what this token is looking for;
• a key (K): what this token can be found by;
• a value (V): what this token hands over, scaled by how much attention it gets.

The usual analogy is a search engine: your query is matched against every page’s keywords, and you get back the pages’ contents. Where it breaks: a search engine hands you a ranked list and you open one page, while attention opens every page at once and blends their contents, weighted by how well each matched. Every token also plays all three parts at once: searcher, page and contents. And nobody writes the keywords; the three matrices are learned in training, like every other weight.

Here’s one score, worked out. In the toy, bank’s query is (2.3, 1.0) and river’s key is (1.5, 0.8). Multiply matching numbers and add them up: 2.3 × 1.5 + 1.0 × 0.8 = 4.25. That’s a dot product, and a big one means a good match. “the” has a key of (0, 0), so it scores 0. Divide both by √2 ≈ 1.41 (why comes shortly) to get 3.0 and 0, and softmax exponentiates: e^3.0 ≈ 20 against e⁰ = 1, so “river” gets about twenty times the attention “the” gets. Do that for every pair at once and the whole mechanism fits in one line:

Read it from the inside out. Stack every token’s query as one row of a table Q, and every key as one row of K. Q · K^T (K turned on its side, so the keys become columns) is a grid with one dot product for every query–key pair. Dividing by √d_k keeps those scores in a stable range. The mask blocks the future. SoftmaxSoftmaxFunction that turns scores into probabilities that add up to 1., the same one that turned logits into probabilities in Part 1, turns each row of scores into weights that add up to 1. And multiplying by V gives every token its weighted mix of values.

Here is every one of those steps on four words, with the same toy weights as Fig. 1. In this head, the query asks “how badly do I need context?” (it’s driven by the ambiguous number) and the key answers “do I carry a clear topic?” (nature or finance, minus ambiguity). I chose those weights so they’d be readable; a trained model learns its own, and they rarely are.

The first row is a party of one. “on” can see only itself, so it gets 100% of its own attention. The second row is a small surprise. A query of zeros scores 0 against every key, softmax of equal scores is an even split, and “the” divides its attention 50/50. That’s the arithmetic of a query with nothing to look for. Trained heads often learn a different habit: they park their spare attention on the very first token, a harmless spot (an attention sink; watch for it in Fig. 3). And “bank” gives “river” 84% here but 74% in Fig. 1, because softmax weights are relative. With “I” and “sat” gone, fewer words share the attention.

Now the mask. Switch it off and “river” sends 28% of its attention to “bank”. That looks harmless until you recall how the model is trained (Part 1; Part 5 has the details): at every position, it learns to predict the next token. The position holding “river” is supposed to predict “bank”. If it can see “bank”, it can copy the answer instead of learning anything. The causal mask adds −∞ to the score of every later token, softmax turns e^−∞ into a weight of exactly 0, and the exam stays honest. The mask pays off twice more before this post is over.

Why divide by √d_k?

In the toy, d_k = 2, so the division is a gentle ÷1.41: skip it and “river” would take 94% of bank’s attention instead of 84%. In a real model it’s essential. Guess first: Llama 3 8B’s heads have d_k = 128, so a score is a sum of 128 products of numbers around ±1. Is it typically about 1, 11 or 128? About 11. Random-ish products mostly cancel, and a sum of d_k of them typically lands around ±√d_k. Feed scores like 11 and 0 into softmax and the 11 gets e¹¹ ≈ 60,000 times the weight: one token takes everything. Worse, a saturated softmax barely moves when its inputs change a little, so training gets almost no signal through it (Part 6). Dividing by √d_k brings typical scores back to around ±1, whatever the size of the head.

Many heads, many questions

One head asks one kind of question; ours asks “which earlier word has a clear topic?” But a sentence holds many relationships at once: which noun a verb belongs to, what a pronoun points at, which word came just before. So a real layer runs many heads side by side, each with its own query matrix, free to learn its own question. This is multi-head attention. In the original design every head also had its own keys and values; Llama 3 8B shares each set among four of its 32 query heads, which matters a lot for memory (more at the end). Each head works with vectors of 128 numbers; their outputs (32 × 128 = 4,096 numbers) are joined back together and mixed by one more learned matrix, the output projection.

Before you tap anything in the figure: in “The keys to the cabinet are lost so I can’t open it”, which word should “are” look at to know it’s are and not is? And what does “it” refer to?

Same sentence, four patterns. The previous-word head draws a neat diagonal. The syntax head connects “are” to “keys”, skipping right past “cabinet”, the nearer noun and exactly the trap that makes people write “the keys to the cabinet is”. The reference head sends “it” to “cabinet”, because you open cabinets, not keys. And the topic head ignores grammar entirely, keeping a running tally of the content words, “keys” and “cabinet” most.

Where did the word order go?

Something odd has been hiding in the formula. Apart from the mask’s rule about who can see whom, nothing in it knows where a token is, and a weighted sum doesn’t care what order you add things in. If you shuffled “on”, “the” and “river” in Fig. 2, the bank row would come out exactly the same. Give a bare attention layer “dog bites man.” and “man bites dog.” and the final full stop gets an identical mix of the same three words. For language, that’s a problem.

So transformers add position information. The original transformer added a fixed pattern of waves to each embedding, different for each position. Many newer models, Llama among them, use RoPE (rotary position embeddings) instead, and the idea is lovely. Split a query or key vector into pairs of numbers and treat each pair as an arrow on a clock face. Before scoring, rotate every arrow by an angle proportional to its token’s position: a little for token 1, twice that for token 2, and so on. Rotate a query and a key this way and the angle between them shifts by an amount that depends only on how far apart the two tokens are, so their score now carries their distance. (Values aren’t rotated: position changes who gets attention, not what is handed over.) Different pairs turn at different speeds, like a clock’s hands. A Llama 3 head’s 64 pairs range from about 57° per token down to one turn every 2.5 million tokens or so, so some pairs are sensitive to the word next door and others to words pages away.

The rest of the layer

Attention moves information between tokens. That’s only half a layer. The other half works on each token on its own.

That half is the feed-forward network (FFNFeed-forward network (FFN)Per-token layer in a transformer that detects patterns and adds learned features.), built from the neurons of Part 3: each takes a weighted sum of the token’s vector and passes it through an activation function. Read each neuron as a pattern detector asking something like “is this about money?”. When the answer is yes, it fires, and adds the features that tend to go with that pattern back into the vector. Real neurons are rarely this tidy (many respond to a jumble of unrelated patterns, and most features are spread across many neurons), but the detect-and-add arithmetic is the same. Attention gathered the clue (“money” is nearby); the FFN acts on it (this is the kind of bank with accounts and deposits). A good share of what a model “knows” appears to live here.

Most of the weights live there too. The FFN is typically 3 to 4 times wider than the model’s vector: in Llama 3 8B, each token’s 4,096 numbers fan out to 14,336 neurons and come back to 4,096. Llama uses a gated variant (SwiGLU) with three matrices of 4,096 × 14,336, about 176 million weights per layer, against about 42 million for the layer’s attention. Across the whole model, the FFNs hold roughly 70% of its 8 billion parameters.

The residual stream. Neither half replaces the token’s vector. Each computes a change and adds it to a running vector called the residual stream: the lookup starts it, attention adds to it, the FFN adds to it, the next layer adds more. Each sub-layer edits a shared document instead of rewriting it, so the word’s own information passes through every layer by default (a layer would have to learn to subtract it), and training gets a clean path back to every layer (Part 6).

Normalization. Before each sub-layer reads the stream, it rescales its own copy of the vector to a standard size and multiplies each number by a learned gain; the stream itself is left alone. This layer normalization (Llama uses a cheaper variant, RMSNorm) means every sub-layer sees inputs on the scale it was trained for, however large the stream has grown over 32 layers, which keeps training stable. One last normalization runs just before the unembedding.

Stack the layers, 32 of them in Llama 3 8B, and that’s the transformer. At the top, the vector at the last position, which by now carries the context of the whole input, becomes the prediction. One more matrix, the unembedding (the output layer, or output table, that Part 2 and Part 3 counted), scores it against every token in the vocabulary: one dot product per token, giving one logitLogitRaw, unnormalized score for each vocabulary token before softmax. each, 128,256 of them for Llama 3 8B. Softmax turns those into probabilities, and the sampler from Part 1 picks one.

Here is the whole “bank” example in one figure. It gives the toy’s vector two extra slots, outdoors and accounts, which every embedding leaves at 0 and only the feed-forward neurons write to. (Real models don’t reserve labeled slots, but their FFNs do write new features into the same vector.) Its vocabulary has just six tokens: ”.”, “and”, “account”, “vault”, “watching” and “fishing”. Guess which one leads after each sentence. Then the harder one: switch attention off. What happens to the feed-forward neurons, which attention never touches?

Follow the river sentence down. Attention adds 0.79 to nature. The “riverside?” neuron computes 2 × 1.29 − 2 × 0.53 − 0.5 ≈ 1.0, so it fires: it writes 1.21 into “outdoors” and turns “ambiguous” down. At the top, “watching” and “fishing” lead. Switch to money and a different neuron fires, writing “accounts”, and “account” leads with 44%.

Now switch attention off. The neurons go dark, though the FFN itself is untouched: it only ever sees this one token’s vector, and without attention that vector says nothing about rivers or money. The predictions slump to the same hedge as “I walked to the bank”, where there’s no clue to find. Now the reverse: attention on, feed-forward off. “bank” still carries 0.79 of extra nature, but no neuron turns it into “outdoors”, and “and” (26%) edges out “fishing” (19%). Gathering the clue isn’t enough; something has to act on it.

Why this design won

The transformer was introduced in 2017 in a paper from Google titled Attention Is All You Need, as a model for translating between languages. The language models of the day were mostly recurrent neural networks (RNNs), which read text the way you do, one token at a time, carrying a running summary from step to step. That has two costs. To process token 500, an RNN must first finish token 499, so the 500 steps run one after another: a GPU can parallelize the arithmetic inside a step, but not the steps themselves. And everything token 500 knows about token 3 has had to survive 497 updates of that summary.

A transformer drops the recurrence. Every position’s queries, keys and values come out of the same big matrix multiplications, all at once, and attention connects token 500 to token 3 directly. Here is the mask’s second payoff: in training, every position of a sequence predicts its own next token simultaneously, because the mask guarantees that none can peek at its answer. One pass over a 4,000-token document yields about 4,000 training examples, computed in parallel. That’s what let models scale to trillions of training tokens on GPUs. Generation is still one token at a time, as in Part 1, since you can’t attend to a word that hasn’t been written yet. (GPT-style models keep only the “decoder” half of the 2017 design, the half with the mask.)

The bill for looking back

That directness has a price. Every token compares its query with the key of every token before it, and its own. For n tokens that’s 1 + 2 + … + n = n(n + 1)/2 comparisons per head per layer, roughly n²/2. Double the text and the comparisons quadruple. It’s the same triangle as the chat bill in Part 1, now inside a single pass.

Generation adds a twist. Each new token needs the keys and values of every earlier token, in every layer, and recomputing them would redo the whole triangle at every step. But they never change: thanks to the mask (its third payoff), an earlier token can’t see a later one, so its keys and values are final the moment they’re computed. So the server stores them. That’s the KV cacheKV cacheStored attention keys and values from earlier tokens. Speeds up generation within a request.. The prompt is processed in one parallel pass (prefill, from Part 3) that fills the cache; after that, each new token computes only its own query, key and value, adds its key and value to the cache, and compares its query against everything stored. The work per new token grows with n instead of n². The price is memory.

With the cache, each step adds one row of comparisons; without it, every step repaints the whole triangle, and writing 40 tokens costs 11,480 comparisons instead of 820. The memory arithmetic is short. A key and a value, in each of 32 layers, 1,024 numbers each, at 2 bytes per number, make 131,072 bytes per token. Llama 3.1 8B, the same architecture, accepts 128k tokens, and past about 122,000 of them a single conversation’s cache outweighs the model itself. At the full 128k it’s about 17 GB.

That 1,024 is the discount promised earlier. Llama 3 8B’s 32 query heads share 8 sets of keys and values, four query heads to each set. This is grouped-query attention (GQA). Switch the figure to one set per head and the cache is four times bigger: 69 GB at 128k, which no longer fits beside the weights on an 80 GB GPU.

Every weight in this post was one I chose by hand. A real model has billions of them, and nobody chooses any. Next: where they come from, starting from random numbers and a game of guessing the next token, played trillions of times.