Attention in Transformers, Step by Step | Deep Learning Chapter 6

3Blue1Brown. Attention in Transformers, Step by Step. Deep Learning Chapter 6. Hosted by Grant Sanderson. Duration: approximately 26 minutes. This is the clearest, most visual explanation of the attention mechanism you will find anywhere, breaking down the exact mathematics that let large language models understand context.

Why Attention Exists

Grant opens with the fundamental problem. When text enters a transformer, each word gets converted into a high-dimensional vector called an embedding, essentially a long list of numbers where different directions correspond to different aspects of meaning. One direction might encode gender, another might encode size, another might encode formality. But at the start, these embeddings only capture what a word means in isolation. They carry zero context.

Consider the word "mole." It means completely different things in "American shrew mole," "one mole of carbon dioxide," and "take a biopsy of the mole." After the initial embedding step, mole gets the same vector in all three cases because the lookup table has no reference to surrounding words. The job of the attention mechanism is to let surrounding embeddings pass information into this one, effectively moving the generic mole vector toward a specific direction in embedding space that captures which meaning is actually intended.

Grant gives another beautiful example with the word "tower." On its own, it points in some generic direction associated with large tall structures. If preceded by "Eiffel," the mechanism should update it toward a direction correlated with Paris, France, and steel. If preceded by "miniature Eiffel," it should shift further so it no longer correlates with large or tall things. The attention block is the machinery that makes all of this possible.

Queries, Keys, and the Attention Pattern

Grant introduces the three core matrices using an intuitive example. Imagine a sentence like "a fluffy blue creature roamed the verdant forest." We want the adjectives to update the meanings of their corresponding nouns.

First, each word asks a question. The noun "creature" generates a query vector by multiplying its embedding through a query matrix. You can think of this query as encoding the question "are there any adjectives sitting in front of me?" This query vector lives in a much smaller dimensional space than the embedding, typically 128 dimensions versus 12,288 for the embedding.

Simultaneously, every word generates a key vector through a separate key matrix. The key is like an answer to potential queries. The adjectives "fluffy" and "blue" would produce keys that closely align with the query produced by "creature." Irrelevant words like "the" would produce keys that point in unrelated directions.

To measure relevance, you compute the dot product between every possible key-query pair, producing a grid of scores. Large positive dot products mean strong alignment. This grid gets normalized column by column using softmax so that each column becomes a probability distribution summing to one. The result is the attention pattern, a map showing how much each word should influence every other word.

Grant shows a compact mathematical notation for all of this. Q and K represent the full arrays of query and key vectors. The grid of dot products is QK transposed, divided by the square root of the key dimension for numerical stability, then softmax is applied column by column.

Masking and Why Later Tokens Cannot Cheat

There is an important training detail. During training, the model simultaneously predicts the next token for every position in the sequence, not just the last one. This means you cannot let later words influence earlier words, because that would leak the answer.

The solution is masking. Before applying softmax, all entries where a later token would influence an earlier one get set to negative infinity. After softmax, those entries become zero, but the columns remain properly normalized. This is why the attention pattern in GPT models always has that characteristic triangular shape. Even though masking is primarily a training concern, it is always applied.

Grant also notes that the attention pattern is a square grid whose size equals the context length squared. This is why scaling context windows is such a bottleneck for large language models. A context of 100,000 tokens means 10 billion entries in this grid. Recent research has explored variations to make attention more scalable, but the fundamentals remain the same.

Value Vectors and Updating Embeddings

Now comes the actual information transfer. A third matrix called the value matrix multiplies each word's embedding to produce a value vector. You can think of the value vector as encoding "if this word turns out to be relevant to something else, what exactly should be added to that something's embedding?"

For each position, you take the weighted sum of all value vectors, where the weights come from that word's column in the attention pattern. The word "creature" would get large contributions from the value vectors of "fluffy" and "blue" because those had high attention weights, and near-zero contributions from everything else. This weighted sum produces a change vector that gets added to the original embedding, yielding a refined embedding that now encodes the contextually richer meaning of "fluffy blue creature."

In practice, the value map is factored into two smaller matrices rather than one enormous square matrix. Grant calls these the value down matrix, which projects the embedding into the smaller key-query space, and the value up matrix, which projects back up to the full embedding space. This factorization keeps the parameter count manageable while maintaining the same conceptual operation.

Multi-Headed Attention and Scale

Everything described so far is a single head of attention. A full attention block runs many heads in parallel, each with its own distinct set of key, query, and value matrices. GPT-3 uses 96 attention heads per block.

Each head learns a different type of contextual relationship. One head might learn adjective-noun associations. Another might learn that "they crashed the" implies something about the physical state of whatever follows. Another might learn that "wizard" anywhere near "Harry" should push the embedding toward Harry Potter rather than Prince Harry.

All 96 heads produce their own proposed changes for each embedding position. These are summed together and added to the original embedding. One technical detail: in practice, the value up matrices from all heads are combined into a single large output matrix associated with the entire multi-headed attention block.

The parameter count is staggering. Each attention head has about 6.3 million parameters across its four matrices. With 96 heads per block, that is 600 million parameters per attention block. GPT-3 has 96 layers, each with its own attention block, bringing the total attention parameters to just under 58 billion. That sounds enormous, but it is actually only about a third of GPT-3's total 175 billion parameters. The majority comes from the multi-layer perceptron blocks that sit between the attention layers.

The Bigger Picture

Grant zooms out to show how data flows through the full transformer architecture. After each attention block, data passes through a multi-layer perceptron, then into another attention block, then another perceptron, repeating this pattern across all 96 layers. Each pass allows embeddings to absorb more and more nuanced information from their increasingly sophisticated surroundings.

By the deep layers of the network, embeddings no longer represent just words with some adjective context. They encode sentiment, tone, whether the text is poetry, what scientific principles are relevant, and all manner of abstract high-level concepts. The final vector in the sequence, which produces the actual next-token prediction, must somehow contain all the information from the entire context window that is relevant to what comes next.

Grant closes by noting that one of the biggest reasons attention has been so successful is not any specific type of behavior it enables, but the fact that it is extremely parallelizable. Matrix multiplications can run on GPUs in a massively parallel fashion, and the lesson of the last decade is that scale alone provides huge qualitative improvements in model performance. Attention enables that scale.

Key Takeaways

The attention mechanism works in three steps. First, query and key matrices identify which words are relevant to which other words by computing dot product similarity scores. Second, these scores are normalized into an attention pattern using softmax. Third, value matrices determine what information to extract from relevant words, and weighted sums of value vectors update each embedding to encode contextual meaning. Multi-headed attention runs many independent versions of this process in parallel, each learning different types of relationships. The result is that static word embeddings are progressively transformed into rich contextual representations that capture far more than any individual word. GPT-3 devotes about 58 billion parameters to attention across 96 layers of 96 heads each, yet this is only a third of the full model.