Attention
Query, key, and value projections
Let represent the model’s activations at a particular layer – each is the activation at token position .
For each attention head , we project each activation to corresponding query, key, and value vectors:
where the linear maps are learned parameters.
The resulting vectors live in a much lower-dimensional space than the original activations (i.e. ).
Intuitively, we can think of the projections as follows:
- The query vector represents what information looks for.
- The key vector represents what information contains.
- The value vector represents what information propagates.
Attention mechanism
The main functionality of the attention mechanism is to transfer information between token positions.
In order to determine which information should be transferred to activation at position , we check to see which past activations contain information that the current activation is looking for.
We can formulate this using the language of query and key vectors: we check to see which past key vectors are similar to the current query vector .
We can compute the similarity between a query vector and key vector by simply taking their dot product:
Here, the subscript indicates that position looks at position – I use this convention throughout.1
Why scale by ?
We scale by to ensure that the dot products don’t grow with . This scaling is important because larger dot products would cause the softmax function to saturate, resulting in vanishing gradients.
To see how scaling by prevents the dot products from growing with , let’s assume and to be drawn from . Then has a mean of 0 and variance of – each summand of the dot product is distributed as , and there are such terms (recall that the variance of the sum of independent random variables is the sum of their variances). Scaling the resulting quantity by results in a variance of .
For causal attention, we set for all . This prevents future token positions from transferring information to past token positions.
We then apply a softmax function to the scores to obtain the attention weights:
Intuitively, describes how strongly information from token position should be transferred to token position . We operationalize this by weighting each value vector by :
Finally, we map this vector back to the original dimension :
where and are learned parameters.
Multi-head attention
The above description focused on a single head. In practice, we feed the activation through multiple, say -many, attention heads in parallel.
For each head , we compute the attention output as described above. We then sum the outputs across all heads:
It is usually the case that . For example, the original transformer[1]Attention is all you need [link]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Advances in Neural Information Processing Systems. 2017. used with and .
KV caching
Once trained, a transformer is generally used to generate sequences of tokens autoregressively – one token at a time.
Consider for a moment how this actually works.
Let’s say we have a prompt as input. We want to generate the next token . We can do this by running the transformer over the whole sequence , and then sampling .
Next, we want to generate . Naively, we could run the transformer over the entire sequence , and then sampling . But it turns out that this is really wasteful!
There are two key observations to notice:
- In a causal transformer, activations at positions will be exactly the same whether we run the transformer over the sequence or . Adding new tokens doesn’t change how previous tokens are processed.
- When running the transformer at position , the only data that is needed from previous token positions are the keys and values.
This leads to an elegant optimization called KV caching. After generating each token, we store the keys and values for all positions processed so far. For each attention head , we maintain:
When running inference at position , we can simply use the cached keys and values to compute the attention output, and also update the cache with new keys and values for position .
This allows us to run the forward pass on just one token position!
However, the cache does incur a memory cost of . Doing vanilla forward passes without caching requires memory – we can store and compute the activations one layer at a time, but need to compute the attention scores for all token pairs.
Multi-query attention
Recall that in standard multi-head attention (MHA), each head has its own query, key, and value projections. While this design is very flexible, it can become memory-intensive as the context length grows, since we need to store different sets of keys and values.
Multi-query attention (MQA)[2]Fast transformer decoding: One write-head is all you need [link]
Noam Shazeer. arXiv preprint. 2019. changes this by sharing a single set of keys and values across all heads, but still having different queries per head.
Concretely:
- We maintain different query projections , so each head still computes its own query vector:
- We now share one key matrix and one value matrix for all heads. Hence, the keys and values become the same for each head:
This means each attention head “sees” the same keys and values, but they “look” at them differently via distinct query vectors.
This approach is very memory-efficient, since the memory cost of the KV cache is reduced from to .
However, MQA is not as expressive as standard MHA, since each head must share the same keys and values.
Grouped-query attention
[Source: Figure 2 of Ainslie et al. [3]GQA: Training generalized multi-query transformer models from multi-head checkpoints [link]
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023.]
Grouped-query attention (GQA)[3]GQA: Training generalized multi-query transformer models from multi-head checkpoints [link]
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023. is a middle-ground approach between full MHA and MQA.
The core idea is to form a smaller number of groups, each group sharing one set of keys and values, but still allowing multiple heads within that group to have distinct queries.
Concretely:
- We partition the heads into groups.
- Each group has a shared key projection and a shared value projection .
- All heads within group use the same key and value projections, but each head in that group keeps its own query projection:
GQA is a middle-ground between full MHA and MQA.
Compared to MHA, GQA reduces the memory cost of the KV cache from to .
Compared to MQA, GQA is more expressive, because there are multiple K/V sets – one per group – rather than a single shared K/V set across all heads.
Multi-head latent attention
[Source: Figure 3 of DeepSeek-AI [4]DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model [link]
DeepSeek-AI. arXiv preprint. 2024.]
Multi-head latent attention (MLA)[4]DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model [link]
DeepSeek-AI. arXiv preprint. 2024. is another technique to reduce the memory cost of the KV cache while maintaining model performance.
Given an activation at position , we first project it to a compressed latent vector :
where is a learned down-projection matrix, projecting from down to . Note that for this compression to be effective, we choose .
This latent vector is then expanded into keys and values for each head:
where are learned up-projection matrices,2 projecting from to .
For queries, MLA similarly uses a compressed representation:
where is a learned down-projection matrix, and is a learned up-projection matrix.
During inference, we only need to cache the latent vectors , not the full keys and values. While caching the full keys and values as in MHA requires memory, caching the latent vectors requires only memory, where .
Another cool property of MLA is that the keys and values don’t need to be computed explicitly. Recall that the attention scores are computed as:
Thus, we can “roll” into , and just compute affinity scores between the compressed query and key vectors.
We can similarly “roll” into .
Sparse attention
The techniques above – MQA, GQA, and MLA – all address the memory cost of the KV cache. But there’s another precious resource that we need to consider: compute. For each query at position , attention considers every preceding key. This means that the computation at position is , and so producing tokens autoregressively requires compute.
The key observation is that most post-softmax attention weights end up near zero. I.e., each query attends strongly to only a small subset of preceding tokens – the rest contribute negligibly to the output.
The key idea behind sparse attention is to first identify a small set of candidate key tokens – the ones actually worth attending to – and then to run full attention only over that set. If we always select a fixed candidates, then the expensive full-attention step becomes per query regardless of context length.
But how do we identify the candidate key tokens? Well, we can scan over all previous tokens and compute some sort of relevance score for each, and then select the top- candidates with the highest scores. Note that this still requires scanning all previous tokens to find the candidates, and therefore is still . But that scan can be much cheaper than full attention. The intuition: answering “is this token worth attending to?” is a simpler question than “how much should I attend to it, and what information does it contain?”, and should therefore be cheaper to compute.
DeepSeek Sparse Attention (DSA)[5]DeepSeek-V3.2: Pushing the frontier of open large language models [link]
DeepSeek-AI. arXiv preprint. 2025. implements this idea, using a “lightning indexer” that scans all previous tokens and assigns a relevance score to each, followed by full attention over only the top- candidates.
For the query at position and each preceding position , the indexer computes:
where is the number of indexer heads, are low-dimensional query and key projections, and is a learned scalar weight derived from that controls how much head ’s score contributes for this query. The indexer can get away with using fewer heads, lower-dimensional vectors, ReLU (instead of softmax), and running at lower precision. With these simplifications, the indexer can run much more cheaply than normal attention.
The positions with the highest scores become the candidate set. Full attention then runs over only those entries.
DeepSeek-V3.2 uses a single shared candidate set across all attention heads, rather than letting each head select independently. This is primarily a hardware convenience: per-head selection would create different irregular memory access patterns. MLA makes this additionally natural, since the latents are already shared across heads.
Cost analysis
How much compute does attention require per query token? For each preceding position, attention performs a fixed amount of work: computing a query-key dot product (to get the attention score) and weighting the corresponding value (to accumulate the output). Call this per-pair cost . Since the query at position pairs with all preceding positions (including itself), the total cost is simply .
For vanilla attention, every pair goes through the full multi-head mechanism, so:
For DSA, two different operations run at different costs: the indexer scans all tokens (at per pair), then full attention runs on only of them (at per pair):
Let’s estimate and for DeepSeek-V3.2. For each (query, key) pair, full MLA attention does two things per head: a QK dot product over dimensions ( FLOPs), and a value aggregation step that scales the -dimensional cached latent by the attention weight and accumulates it (another FLOPs).3 That’s per head, giving:
The indexer only scores (no value aggregation) using heads with -dimensional dot products:
The ratio : the indexer does roughly 6% of the work per pair. Note that this estimation only counts raw FLOPs. The indexer can also run at lower precision in optimized implementations, which can further improve wall-clock speed beyond this raw-FLOP comparison.
References
References cited in the text are listed first, in order of citation; additional references follow, ordered alphabetically.
- Attention is all you need [link]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Advances in Neural Information Processing Systems. 2017. - Fast transformer decoding: One write-head is all you need [link]
Noam Shazeer. arXiv preprint. 2019. - GQA: Training generalized multi-query transformer models from multi-head checkpoints [link]
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023. - DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model [link]
DeepSeek-AI. arXiv preprint. 2024. - DeepSeek-V3.2: Pushing the frontier of open large language models [link]
DeepSeek-AI. arXiv preprint. 2025. - A mathematical framework for transformer circuits [link]
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. Transformer Circuits Thread. 2021. - An analogy for understanding transformers [link]
Callum McDougall. 2023. - The annotated transformer [link]
Sasha Rush, Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman. 2022.
Footnotes
-
Some prefer the reverse convention, , since information flows from to . I find more natural for attention patterns. ↩
-
Note that the head-specific matrices and may not actually be “up-projection” matrices, as we previously specified that , not necessarily that . The original paper works with and as matrices projecting from to , rather than notating a separate matrix for each head, and through this lens and are true “up-projection” matrices. ↩
-
A -dimensional dot product requires multiplications and additions, totaling FLOPs. The value aggregation has the same cost structure: scalar-by-vector multiplications plus additions to an accumulator. This estimate excludes the smaller RoPE score term. ↩