Anatomy of an LLM: the attention mechanism

🚀 Intro

In the previous stop of this series - one of the branches of the overview post “Anatomy of an LLM” - we turned text into numbers: a sequence of words broke into tokens, and each token received its own vector of meaning from a trained embedding matrix. But we ended on an unsettling observation: those vectors still sit side by side separately. Each one carries the meaning of a word as if “in a vacuum,” cut off from its neighbors. And yet the sense of a word depends on context - “bank” means one thing by a river and another in finance.

The mechanism that lets vectors “talk” to one another - look at each other and exchange information - is attention. It is the heart of the Transformer and, with little exaggeration, one of the most important ideas in the entire history of machine learning. The architecture behind ChatGPT, Claude, and Gemini even took its name from the title of the paper that made it famous: Attention Is All You Need.

In this post we trace attention from its birth - before the Transformer, in translation networks - through the famous scaled formula, all the way to the modern variants that today let models read hundreds of thousands of tokens at once. And at the end we ask a question that still has no full answer: what do these “attention heads” actually do? Because here too - in the spirit of the whole series - the mechanism itself turns out to be mundane, while what emerges from it is once again surprisingly hard to fully understand.

💡 How to read this. The main thread is descriptive and self-contained. I tuck the formulas into expandable “for the curious” blocks - if you like math, open them; if not, skip them with no loss to understanding. Every claim has a footnote to its source paper at the end.

📋 TL;DR

Attention solved a bottleneck. Old translation networks squeezed an entire source sentence into a single fixed-length vector. Bahdanau, Cho, and Bengio (2014) let the model instead “search” the sentence and take from it whatever was relevant at the moment - that was the first attention mechanism.
The core: a scaled dot product. Each token produces three vectors - a Query, a Key, and a Value. A token compares its query against the keys of all the others, and the result (after passing through a softmax) decides how much of whose value to take.
Why we divide by $\sqrt{d_k}$ (read “the square root of d-k”). Without this scaling the dot products grow with the dimension and push the softmax into a region of nearly zero gradients. Dividing by the square root of the dimension restores a healthy range.
The causal mask. In a model that generates text, a token must not “see” the future - this is blocked with a triangular mask: forbidden positions get minus infinity added just before the softmax, so afterward their weight drops to zero.
Many heads at once. Instead of a single attention function, the model runs several to several dozen parallel heads, each in a smaller subspace - and each can pick up a different kind of dependency.
Quadratic cost. Attention compares every token with every other, so the cost grows with the square of the sequence length. Hence the whole engineering of long context: KV-cache, MQA/GQA, FlashAttention.
What do the heads do? Carefully. Attention weights are often mistaken for an explanation of the model’s decision - and they are not always one (Jain and Wallace). But mechanistic interpretability has found concrete circuits - induction heads - that may underlie in-context learning.

👉 Play with it: on the interactive page you have a section on the attention mechanism - watch live how a token spreads its attention across the other words of a sentence.

🌅 Before the Transformer: where attention came from

Attention did not arrive all at once with the Transformer. It was born earlier, as a patch for a specific, painful problem in machine translation.

Earlier translation models had an encoder-decoder shape: the encoder read the whole source sentence and squeezed it into a single fixed-length vector (the so-called context vector), and the decoder produced the translation from that one vector. The weakness is plain to see - the entire content of the sentence, however long, had to fit into one predetermined string of numbers. That is the bottleneck: the longer the sentence, the more was lost.

Bahdanau, Cho, and Bengio (2014) said it outright - “the use of a fixed-length vector is a bottleneck” - and proposed something clever¹: let the decoder, as it generates each next word, automatically and softly “search” the source sentence for the fragments relevant right now. Instead of one summary of the whole, the model got, each time, a different weighted mixture of all the input words. This was the first attention mechanism (so-called additive or “Bahdanau attention”), in which the alignment weights were computed by a small network with a $\tanh$ function (read “hyperbolic tangent”).

A year later, Luong, Pham, and Manning (2015) simplified the idea, showing multiplicative attention based on the plain dot product of two vectors². That simplification - comparing vectors by their dot product - is the direct ancestor of what every LLM does today. It is worth remembering, though: in Bahdanau and Luong, attention was still only an add-on to a recurrent network (reading text word by word). Only the Transformer dared to throw out recurrence entirely and rest everything on attention alone.

📐 For the curious: Luong's three scoring functions

Luong and colleagues compared three ways of computing the “alignment” between a decoder hidden state $h_t$ and a source state $\bar h_s$ :

$\text{score}_{\text{dot}}(h_t,\bar h_s)=h_t^{\top}\bar h_s$

$\text{score}_{\text{general}}(h_t,\bar h_s)=h_t^{\top}W_a\,\bar h_s$

$\text{score}_{\text{concat}}(h_t,\bar h_s)=v_a^{\top}\tanh\!\big(W_a[h_t;\bar h_s]\big)$

The symbol $\top$ (transpose) turns a column vector into a row vector, so $h_t^{\top}\bar h_s$ is simply the dot product - a single number telling how much two vectors “point the same way.” The dot variant (pure “dot product”) is the simplest, and it is exactly the one that became the seed of the $QK^{\top}$ operation in the Transformer. The concat variant roughly corresponds to the earlier additive Bahdanau attention.

🎯 The heart of the mechanism: Query, Key, Value

Here we reach the core. In the Transformer, each token produces three separate vectors from its own vector, by passing it through three different (trained) matrices. The easiest way to understand them is an analogy to searching in a library:

Query - “what am I, this token, looking for right now?”
Key - “how do I, this token, advertise myself; what is my label?”
Value - “what will I actually hand over if someone picks me?”

The mechanism works like a soft search. A token takes its query and compares it against the keys of all the tokens in the sentence (itself included). Each comparison is a dot product - a single number saying “how well do you two match.” These numbers pass through a softmax, which turns them into weights that sum to one - something like a distribution of “what fraction of my attention I devote to each of the other tokens.” Finally the token pulls in a weighted sum of the values of all tokens: a lot from those it matched strongly, little from the rest.

The effect is that each token updates its own vector by blending in information from its most relevant neighbors. The word “bank” in a sentence about a river will pull in information from words like “water” and “shore”; the same “bank” next to “loan” - from words about finance. Attention is the mechanism by which meaning becomes contextual.

That mysterious fraction $\frac{1}{\sqrt{d_k}}$ in the original formula still remains. It is not an ornament - it solves a real numerical problem. When the query and key vectors are long (and in large models they have hundreds of dimensions), their dot products become very large in absolute value. And large numbers at the softmax input have a ruinous effect: the softmax saturates, meaning almost all the weight lands on a single token, and its derivatives (gradients) fall to zero. The network stops learning. Dividing by $\sqrt{d_k}$ - the square root of the dimension of a single head - brings these numbers back to a healthy range precisely.

📐 For the curious: the attention formula and where the square root of the dimension comes from

The full definition of scaled dot-product attention (Vaswani et al., 2017):

$\text{Attention}(Q,K,V)=\text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V$

Here $Q$ , $K$ , $V$ are matrices in which each row is, respectively, the query, key, and value of one token. The product $QK^{\top}$ yields a square matrix of all “query-key” pairs.

Why exactly $\sqrt{d_k}$ ? Assume the individual components of the vectors $q$ and $k$ are independent, with mean $0$ and variance $1$ . Their dot product is a sum of $d_k$ such products:

$q\cdot k=\sum_{i=1}^{d_k} q_i k_i,\qquad \mathbb{E}[q\cdot k]=0,\qquad \text{Var}(q\cdot k)=\sum_{i=1}^{d_k}\text{Var}(q_i k_i)=d_k$

The symbol $\mathbb{E}$ is the expected value (“the mean”), $\text{Var}$ is the variance (“the spread”), and $\Sigma$ (sigma) denotes summation. The variance of the dot product thus grows linearly with the dimension $d_k$ , and the standard deviation (the typical spread) is $\sqrt{d_k}$ . Dividing the whole dot product by $\sqrt{d_k}$ brings its variance back to $1$ - and the softmax works in a region of healthy gradients instead of saturating³.

👁️ What you are allowed to look at: self-, cross-, and causal attention

Attention comes in several flavors, depending on who looks at whom:

Self-attention - $Q$ , $K$ , and $V$ come from the same sequence. Tokens look at one another. This is the basic mode in language models.
Cross-attention - the queries come from one sequence, while the keys and values come from another. This is how, for example, a decoder peeks at an encoder in translation. It is the direct heir of Bahdanau and Luong attention.
Causal attention (“masked”) - and here comes something absolutely crucial for text-generating models.

A model that writes text word by word must predict token $i+1$ using only what has already been written - tokens $\le i$ . If during training it could “peek” at future tokens, it would cheat: it would learn an answer it already sees. That would utterly ruin the learning. So every token must be forbidden from looking forward.

This is done elegantly, with a triangular mask. Just before the softmax, minus infinity is added to all the “forbidden” positions (the ones from the future). A softmax of minus infinity gives exactly zero - so afterward those positions have zero weight, as if they did not exist. A token can look only backward and at itself. This seemingly small trick is what makes the model autoregressive - able to generate text in one direction, from left to right.

📐 For the curious: how the triangular mask works

A mask matrix $M$ is added to the score matrix just before the softmax:

$\text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}+M\right)V,\qquad M_{ij}=\begin{cases}0 & \text{if } j\le i\\[2pt] -\infty & \text{if } j>i\end{cases}$

The index $i$ is the position of the token that looks, and $j$ is the position of the token it wants to look at. For positions from the future ( $j>i$ ) the mask entry is $-\infty$ (minus infinity). Since the softmax computes $e^{x}$ , and $e^{-\infty}=0$ , those positions get a weight of exactly zero after normalization. The others ( $j\le i$ ) get $0$ , i.e. they are left untouched. Hence the name “triangular” - the allowed cells form the lower triangle of the matrix.

🧠 Many heads at once

A single attention has one “perspective” - one way in which it compares tokens. But language is multi-dimensional: syntax, coreference (what “it” refers to), topical closeness, and order all matter at once. It is hard for one attention function to catch all of that at the same time.

The solution is multi-head attention. Instead of one full-dimensional attention, the Transformer runs $h$ heads in parallel, each in a smaller subspace. Each head has its own matrices for producing queries, keys, and values - so each can specialize in something different. The outputs of all the heads are glued back together and passed through one final matrix. The Transformer’s authors put it this way: many heads “allow the model to jointly attend to information from different representation subspaces at different positions”³.

Importantly, this does not significantly raise the cost. Since each of the $h$ heads works in a dimension $h$ times smaller, the total bill is close to the cost of a single full-dimensional head. The base Transformer had $h=8$ heads, each of dimension 64 (with a full model dimension of 512). In today’s large models there are several dozen heads.

📐 For the curious: the multi-head attention formula

$\text{MultiHead}(Q,K,V)=\text{Concat}(\text{head}_1,\dots,\text{head}_h)\,W^{O}$

$\text{head}_i=\text{Attention}\!\big(QW_i^{Q},\,KW_i^{K},\,VW_i^{V}\big)$

Each head $i$ has its own projection matrices $W_i^{Q}, W_i^{K}, W_i^{V}$ , which project the input into a subspace of dimension $d_k=d_{\text{model}}/h$ (read “d-model over h”). $\text{Concat}$ glues the outputs of all the heads back into one full-dimensional vector, and a shared matrix $W^{O}$ mixes them into the final output. In the base model: $d_{\text{model}}=512$ , $h=8$ , so $d_k=d_v=64$ ³.

⏳ Quadratic cost and the long-context problem

All of this has its price - and quite literally. Since every token compares its query against the key of every other token, then for a sequence of length $n$ one must compute an $n \times n$ matrix of comparisons. The cost therefore grows with the square of the sequence length. Doubling the context length is a fourfold increase in the bill; a tenfold lengthening - a hundredfold.

This is the main limitation of the Transformer and the reason why “long context” was a luxury for years. Remember from the previous post why we don’t cut text into single characters? Exactly because of this - longer sequences crash into this quadratic cost. The whole line of research and engineering of recent years is a fight against this quadraticity.

📐 For the curious: where the quadratic cost comes from and what a KV-cache is

The score matrix $QK^{\top}$ has size $n \times n$ (each token with each), and computing and storing it costs $O(n^2 \cdot d)$ in time and memory, where $d$ is the dimension³. This is a quadratic dependence on the length $n$ - the dominant cost at long context.

KV-cache. When generating text token by token, notice that the keys and values ( $K$ , $V$ ) of the already generated positions do not change. There is no point recomputing them at every step - they are buffered (this is exactly the “key-value cache”). The price: that buffer grows linearly with the sequence length $\times$ the number of layers $\times$ the number of heads, and repeatedly loading it from memory becomes the new bottleneck - this time of memory bandwidth, not computation⁴.

⚡ Modern variants: how long context was tamed

Several ideas fight today against the quadratic cost and the memory burden. Some are exact (they give an identical result, just computed more cleverly), some are approximate (they change the result to make it cheaper).

Multi-Query Attention (MQA), Shazeer 2019. Since it is loading those huge buffers of keys and values that hurts most, Shazeer proposed that all the query heads share a single pair of keys and values⁴. The KV-cache shrinks dramatically, decoding speeds up - with “only minor quality degradation.”

Grouped-Query Attention (GQA), Ainslie et al. 2023. A middle ground between full multi-head and MQA: the query heads are split into $G$ groups, and each group shares one key-value pair⁵. GQA “achieves quality close to” full multi-head attention “at speed comparable to MQA,” and on top of that lets an existing model be retrained into this form for as little as 5% of the original training cost. Hence its popularity: Llama 2 uses GQA in the largest variants (34B and 70B), and Llama 3 - in all sizes, with the context stretched to 128 thousand tokens⁶⁷.

FlashAttention (Dao et al. 2022). Here the idea is different - we do not simplify the math, we just manage memory more wisely. FlashAttention computes exactly the same attention (a result identical to the bit), but it organizes the computation in GPU memory (a “tiling” technique) so that it never materializes the full $n \times n$ matrix⁸. The effect: real speedups and memory usage linear in the sequence length instead of quadratic. FlashAttention-2 (Dao 2023) added another roughly twofold speedup thanks to a better division of work⁹. It is one of the quiet heroes of the long-context era - without it, today’s windows of hundreds of thousands of tokens would be far more expensive.

Sparse and linear attention. Other approaches change the very math of attention to escape the square. Longformer combines a local attention window with sparse global attention and scales linearly¹⁰; Performer approximates full softmax attention in linear time and memory¹¹. Unlike FlashAttention, these methods give an approximate result, not an exact one - which is sometimes good enough, but not always.

🔍 What attention heads actually do

Now that we know how attention computes, the natural question is: what do the individual heads actually do? It is tempting to reach for the popular image - “this head watches syntax, that one tracks coreference” - but here a lot of caution is needed. This is an area where it is easy to over-interpret.

First, a warning. Jain and Wallace (2019), in a paper with the telling title Attention is not Explanation, showed that attention weights “largely do not” provide reliable explanations of the model’s decisions¹². They are often uncorrelated with other importance measures, and worse - one can construct entirely different attention distributions that yield the same predictions. In other words: the fact that a head “looks” strongly at some word does not yet prove that this word was the cause of the answer. (The matter is contested - there is a rebuttal, Attention is not not Explanation by Wiegreffe and Pinter, 2019¹³ - but the very existence of the dispute teaches humility.)

Second, something astonishingly concrete. The strand of mechanistic interpretability - attempts to decompose a network into understandable “circuits” - identified a special kind of head: the induction head. Olsson and colleagues at Anthropic (2022) described a head that performs simple pattern completion¹⁴: having earlier seen the sequence [A][B], when it again encounters [A], it predicts [B]. It works through “prefix matching” and “copying” - the head searches for where the current token previously appeared and suggests what followed it.

What makes this important: the authors put forward a hypothesis (and stress that it is a hypothesis) that induction heads may constitute the mechanism of “the majority of all in-context learning” in large models - that remarkable ability by which a model learns a new task from a few examples given in the prompt itself, with no fine-tuning at all. In support they show, among other things, a sudden phase change early in training: at some moment the capacity for in-context learning appears abruptly, exactly when these circuits form in the network. This is careful empirical analysis with the uncertainty explicitly flagged - not a full explanation of all attention heads, but a real foothold of understanding.

🌗 Learning that nobody programmed

Let us pause here for a moment, because we have just touched something that ties this whole series together. The entire attention mechanism we have taken apart is - like everything in this anatomy - mundane: a few matrix multiplications, one softmax, a triangular mask. An engineer understands every one of these steps. And yet, from their repetition across a deep stack of layers, there emerges an ability nobody designed explicitly - in-context learning. There is no line in the code saying “learn from the examples in the prompt.” There is only the pressure from the previous posts - “predict the next token as well as you can” - and from it, somewhere in the middle of training, circuits crystallize on their own that can do something that looks like reasoning by analogy.

And this is no diffuse emergence. As we showed above, in-context learning appears almost abruptly - at a countable moment in training, when the induction heads form. All the more tempting, then, to read off the attention maps what the model understands. That is why I weighed my words here: an attention map shows where the model looks, but it is not yet an explanation of why it decides one way rather than another. That distinction - “where it looks” versus “why it decides so” - is the axis of caution I kept throughout the post.

What remains is the harder question: where does it come from - that a blind pressure on getting the next token right gives rise to a circuit that learns on the fly? It reaches beyond attention itself, and perhaps beyond this architecture. The vectors that in the previous post sat apart began here to “talk” to one another - and out of that conversation emerged a behavior that was not in the blueprint. Where that trail leads, I follow separately, in the essay on the human-LLM resonance.

🎯 What’s next

We have the heart of the Transformer. We now know how a token creates a query, key, and value, how it compares itself with others through a scaled dot product, why we divide by $\sqrt{d_k}$ , how the triangular mask makes sure the model does not peek at the future, and why we run many heads at once. We also know where the quadratic cost comes from and how engineering - KV-cache, GQA, FlashAttention - pushes the limits of long context.

But attention by itself is not yet the whole model. It is one building block - a very important one, but a block. A token that has gathered information from its neighbors still has to process it; the layers must be stacked into a deep tower so that training even converges; and a decision has to be made about how to draw a concrete word out of the final vectors. In the next stop of the series we assemble the full Transformer block: the processing layer (FFN), the residual connections and normalization that make it possible to build stacks of dozens of layers, and how, at the very end, the model turns numbers back into text.

If you would rather touch now instead of reading on, the interactive “Anatomy of an LLM” page lets you see how a token spreads its attention across the other words of a sentence. And if you are drawn to what truly emerges from the mechanism described here - and why simple rules give rise to abilities nobody programmed - I wrote about it separately.

D. Bahdanau, K. Cho, Y. Bengio (2014). Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015. arXiv:1409.0473. https://arxiv.org/abs/1409.0473 ↩
M.-T. Luong, H. Pham, C. D. Manning (2015). Effective Approaches to Attention-based Neural Machine Translation. EMNLP. arXiv:1508.04025. https://arxiv.org/abs/1508.04025 ↩
A. Vaswani et al. (2017). Attention Is All You Need (scaled dot-product attention, $\sqrt{d_k}$ , multi-head, $O(n^2 d)$ complexity). NeurIPS. arXiv:1706.03762. https://arxiv.org/abs/1706.03762 ↩ ↩² ↩³ ↩⁴
N. Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need (MQA, KV-cache, memory-bandwidth bottleneck). arXiv:1911.02150. https://arxiv.org/abs/1911.02150 ↩ ↩²
J. Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP. arXiv:2305.13245. https://arxiv.org/abs/2305.13245 ↩
H. Touvron et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models (GQA in the 34B/70B variants, context 4096). arXiv:2307.09288. https://arxiv.org/abs/2307.09288 ↩
A. Grattafiori et al. (2024). The Llama 3 Herd of Models (GQA in all sizes, context up to 128K). arXiv:2407.21783. https://arxiv.org/abs/2407.21783 ↩
T. Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS. arXiv:2205.14135. https://arxiv.org/abs/2205.14135 ↩
T. Dao (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691. https://arxiv.org/abs/2307.08691 ↩
I. Beltagy, M. E. Peters, A. Cohan (2020). Longformer: The Long-Document Transformer (sparse attention, linear scaling). arXiv:2004.05150. https://arxiv.org/abs/2004.05150 ↩
K. Choromanski et al. (2021). Rethinking Attention with Performers (FAVOR+ approximation, linear complexity). ICLR. arXiv:2009.14794. https://arxiv.org/abs/2009.14794 ↩
S. Jain, B. C. Wallace (2019). Attention is not Explanation. NAACL. arXiv:1902.10186. https://arxiv.org/abs/1902.10186 ↩
S. Wiegreffe, Y. Pinter (2019). Attention is not not Explanation (a rebuttal to Jain and Wallace). EMNLP. arXiv:1908.04626. https://arxiv.org/abs/1908.04626 ↩
C. Olsson et al. (2022). In-Context Learning and Induction Heads. Anthropic / Transformer Circuits Thread. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html ↩