Query, key, and value projections

Let XRnseq×dmodel\mathbf{X} \in \mathbb{R}^{n_\text{seq} \times d_\text{model}} represent the model’s activations at a particular layer – each xiRdmodel\mathbf{x}_i \in \mathbb{R}^{d_\text{model}} is the activation at token position i[nseq]i \in [n_{\text{seq}}].

For each attention head hh, we project each activation xi\mathbf{x}_i to corresponding query, key, and value vectors:

qi(h)=WQ(h)xi,ki(h)=WK(h)xi,vi(h)=WV(h)xi,\begin{aligned} \mathbf{q}_i^{(h)} &= \mathbf{W}_Q^{(h)} \mathbf{x}_i, \\ \mathbf{k}_i^{(h)} &= \mathbf{W}_K^{(h)} \mathbf{x}_i, \\ \mathbf{v}_i^{(h)} &= \mathbf{W}_V^{(h)} \mathbf{x}_i, \end{aligned}

where the linear maps WQ(h),WK(h),WV(h)Rdhead×dmodel\mathbf{W}_Q^{(h)}, \mathbf{W}_K^{(h)}, \mathbf{W}_V^{(h)} \in \mathbb{R}^{d_{\text{head}} \times d_{\text{model}}} are learned parameters.

The resulting vectors qi(h),ki(h),vi(h)Rdhead\mathbf{q}_i^{(h)}, \mathbf{k}_i^{(h)}, \mathbf{v}_i^{(h)} \in \mathbb{R}^{d_{\text{head}}} live in a much lower-dimensional space than the original activations (i.e. dheaddmodeld_{\text{head}} \ll d_{\text{model}}).

Intuitively, we can think of the projections as follows:

  • The query vector qi(h)\mathbf{q}_i^{(h)} represents what information xi\mathbf{x}_i looks for.
  • The key vector ki(h)\mathbf{k}_i^{(h)} represents what information xi\mathbf{x}_i contains.
  • The value vector vi(h)\mathbf{v}_i^{(h)} represents what information xi\mathbf{x}_i 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 xi\mathbf{x}_i at position ii, we check to see which past activations {xjji}\{ \mathbf{x}_j \mid j \leq i \} contain information that the current activation xi\mathbf{x}_i is looking for.

We can formulate this using the language of query and key vectors: we check to see which past key vectors {kj(h)ji}\{ \mathbf{k}_j^{(h)} \mid j \leq i \} are similar to the current query vector qi(h)\mathbf{q}_i^{(h)}.

We can compute the similarity between a query vector and key vector by simply taking their dot product:

scoreij(h)=qi(h)kj(h)dhead.\begin{aligned} \text{score}_{i \rightarrow j}^{(h)} = \frac{\mathbf{q}_i^{(h)} \cdot \mathbf{k}_j^{(h)}}{\sqrt{d_{\text{head}}}}. \end{aligned}

Here, the subscript iji \rightarrow j indicates that position ii looks at position jj – I use this convention throughout.1

Why scale by 1dhead\frac{1}{\sqrt{d_{\text{head}}}}?

We scale by 1dhead\frac{1}{\sqrt{d_{\text{head}}}} to ensure that the dot products don’t grow with dheadd_{\text{head}}. This scaling is important because larger dot products would cause the softmax function to saturate, resulting in vanishing gradients.

To see how scaling by 1dhead\frac{1}{\sqrt{d_{\text{head}}}} prevents the dot products from growing with dheadd_{\text{head}}, let’s assume qi(h)\mathbf{q}_i^{(h)} and ki(h)\mathbf{k}_i^{(h)} to be drawn from N(0,I)\mathcal{N}(\mathbf{0}, \mathbf{I}). Then qi(h)ki(h)\mathbf{q}_i^{(h)} \cdot \mathbf{k}_i^{(h)} has a mean of 0 and variance of dheadd_{\text{head}} – each summand of the dot product is distributed as N(0,1)\mathcal{N}(0, 1), and there are dheadd_{\text{head}} such terms (recall that the variance of the sum of independent random variables is the sum of their variances). Scaling the resulting quantity by 1dhead\frac{1}{\sqrt{d_{\text{head}}}} results in a variance of (1dhead)2dhead=1\left( \frac{1}{\sqrt{d_{\text{head}}}}\right)^2 \cdot d_{\text{head}} = 1.

For causal attention, we set scoreij(h)=\text{score}_{i \rightarrow j}^{(h)} = -\infty for all j>ij > i. 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:

attentionij(h)=exp(scoreij(h))k=1nseqexp(scoreik(h)).\text{attention}_{i \rightarrow j}^{(h)} = \frac{\exp(\text{score}_{i \rightarrow j}^{(h)})}{\sum_{k=1}^{n_{\text{seq}}} \exp(\text{score}_{i \rightarrow k}^{(h)})}.

Intuitively, attentionij(h)\text{attention}_{i \rightarrow j}^{(h)} describes how strongly information from token position jj should be transferred to token position ii. We operationalize this by weighting each value vector vj(h)\mathbf{v}_{j}^{(h)} by attentionij(h)\text{attention}_{i \rightarrow j}^{(h)}:

weighted_valuei(h)=j=1nseqattentionij(h)vj(h).\begin{aligned} \text{weighted\_value}_i^{(h)} = \sum_{j=1}^{n_{\text{seq}}} \text{attention}_{i \rightarrow j}^{(h)} \mathbf{v}_j^{(h)}. \end{aligned}

Finally, we map this vector weighted_valuei(h)Rdhead\text{weighted\_value}_i^{(h)} \in \mathbb{R}^{d_{\text{head}}} back to the original dimension dmodeld_{\text{model}}:

attention_outi(h)=WO(h)weighted_valuei(h)+bO(h),\begin{aligned} \text{attention\_out}_i^{(h)} = \mathbf{W}_O^{(h)} \text{weighted\_value}_i^{(h)} + \mathbf{b}_O^{(h)}, \end{aligned}

where WO(h)Rdmodel×dhead\mathbf{W}_O^{(h)} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{head}}} and bO(h)Rdmodel\mathbf{b}_O^{(h)} \in \mathbb{R}^{d_{\text{model}}} are learned parameters.

Multi-head attention

The above description focused on a single head. In practice, we feed the activation xi\mathbf{x}_i through multiple, say nheadsn_{\text{heads}}-many, attention heads in parallel.

For each head h[nheads]h \in [n_{\text{heads}}], we compute the attention output attention_outi(h)\text{attention\_out}_i^{(h)} as described above. We then sum the outputs across all heads:

multi_head_attention_outi=h[nheads]attention_outi(h).\text{multi\_head\_attention\_out}_i = \sum_{h \in [n_{\text{heads}}]} \text{attention\_out}_i^{(h)}.

It is usually the case that dmodel=dheadnheadsd_{\text{model}} = d_{\text{head}} \cdot n_{\text{heads}}. 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 dmodel=512d_{\text{model}} = 512 with nheads=8n_{\text{heads}} = 8 and dhead=64d_{\text{head}} = 64.

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 [t1,,tn][t_1, \ldots, t_n] as input. We want to generate the next token tn+1t_{n+1}. We can do this by running the transformer over the whole sequence [t1,,tn][t_1, \ldots, t_n], and then sampling tn+1t_{n+1}.

Next, we want to generate tn+2t_{n+2}. Naively, we could run the transformer over the entire sequence [t1,,tn,tn+1][t_1, \ldots, t_n, t_{n+1}], and then sampling tn+2t_{n+2}. But it turns out that this is really wasteful!

There are two key observations to notice:

  1. In a causal transformer, activations at positions 1,,n1, \ldots, n will be exactly the same whether we run the transformer over the sequence [t1,,tn][t_1, \ldots, t_n] or [t1,,tn,tn+1][t_1, \ldots, t_n, t_{n+1}]. Adding new tokens doesn’t change how previous tokens are processed.
  2. When running the transformer at position n+1n+1, 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 hh, we maintain:

key_cache(h):[k1(h),k2(h),,kt(h)]value_cache(h):[v1(h),v2(h),,vt(h)]\begin{aligned} \text{key\_cache}^{(h)} &: [\mathbf{k}_1^{(h)}, \mathbf{k}_2^{(h)}, \ldots, \mathbf{k}_t^{(h)}] \\ \text{value\_cache}^{(h)}&: [\mathbf{v}_1^{(h)}, \mathbf{v}_2^{(h)}, \ldots, \mathbf{v}_t^{(h)}] \end{aligned}

When running inference at position t+1t+1, 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 t+1t+1.

This allows us to run the forward pass on just one token position!

However, the cache does incur a memory cost of O(2nseqnlayersnheadsdhead)O(2 \cdot n_{\text{seq}} \cdot n_{\text{layers}} \cdot n_{\text{heads}} \cdot d_{\text{head}}). Doing vanilla forward passes without caching requires O(nseqdmodel+nseq2)O(n_{\text{seq}} \cdot d_{\text{model}} + n_{\text{seq}}^2) 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 nseqnlayersnheadsn_{\text{seq}} \cdot n_{\text{layers}} \cdot n_{\text{heads}} 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 nheadsn_{\text{heads}} different query projections WQ(h)\mathbf{W}_Q^{(h)}, so each head still computes its own query vector:
qi(h)=WQ(h)xi.\mathbf{q}_i^{(h)} = \mathbf{W}_Q^{(h)} \mathbf{x}_i.
  • We now share one key matrix WK\mathbf{W}_K and one value matrix WV\mathbf{W}_V for all heads. Hence, the keys and values become the same for each head:
ki(h)=WKxi,vi(h)=WVxi.\begin{aligned} \mathbf{k}_i^{(h)} &= \mathbf{W}_K \mathbf{x}_i, \\ \mathbf{v}_i^{(h)} &= \mathbf{W}_V \mathbf{x}_i. \end{aligned}

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 O(nseqnlayersnheadsdhead)O(n_{\text{seq}} \cdot n_{\text{layers}} \cdot n_{\text{heads}} \cdot d_{\text{head}}) to O(nseqnlayersdhead)O(n_{\text{seq}} \cdot n_{\text{layers}} \cdot d_{\text{head}}).

However, MQA is not as expressive as standard MHA, since each head must share the same keys and values.

Grouped-query attention

A nice schematic displaying the difference between multi-head, multi-query, and 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:

  1. We partition the nheadsn_{\text{heads}} heads into gg groups.
  2. Each group r[g]r \in [g] has a shared key projection WK(r)\mathbf{W}_K^{(r)} and a shared value projection WV(r)\mathbf{W}_V^{(r)}.
  3. All heads within group rr use the same key and value projections, but each head in that group keeps its own query projection:
qi(h)=WQ(h)xi,ki(h)=WK(r)xi(for h assigned to group r),vi(h)=WV(r)xi(for h assigned to group r).\begin{aligned} \mathbf{q}_i^{(h)} &= \mathbf{W}_Q^{(h)} \mathbf{x}_i, \\ \mathbf{k}_i^{(h)} &= \mathbf{W}_K^{(r)} \mathbf{x}_i \quad (\text{for } h \text{ assigned to group } r), \\ \mathbf{v}_i^{(h)} &= \mathbf{W}_V^{(r)} \mathbf{x}_i \quad (\text{for } h \text{ assigned to group } r). \end{aligned}

GQA is a middle-ground between full MHA and MQA.

Compared to MHA, GQA reduces the memory cost of the KV cache from O(nseqnlayersnheadsdhead)O(n_{\text{seq}} \cdot n_{\text{layers}} \cdot n_{\text{heads}} \cdot d_{\text{head}}) to O(nseqnlayersgdhead)O(n_{\text{seq}} \cdot n_{\text{layers}} \cdot g \cdot d_{\text{head}}).

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

Rather than caching each key and value individually, multi-head latent attention (MLA) caches a compressed latent, and then projects this compressed latent to keys and values.
[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 xi\mathbf{x}_i at position ii, we first project it to a compressed latent vector ciKV\mathbf{c}_i^{\text{KV}}:

ciKV=WDKVxi,\mathbf{c}_i^{\text{KV}} =\mathbf{W}_{\text{DKV}} \mathbf{x}_i,

where WDKVRdc×dmodel\mathbf{W}_{\text{DKV}} \in \mathbb{R}^{d_c \times d_{\text{model}}} is a learned down-projection matrix, projecting from dmodeld_{\text{model}} down to dcd_c. Note that for this compression to be effective, we choose dcnheadsdheadd_c \ll n_{\text{heads}} \cdot d_{\text{head}}.

This latent vector is then expanded into keys and values for each head:

ki(h)=WUK(h)ciKV,vi(h)=WUV(h)ciKV,\begin{aligned} \mathbf{k}_i^{(h)} &= \mathbf{W}_{\text{UK}}^{(h)} \mathbf{c}_i^{\text{KV}}, \\ \mathbf{v}_i^{(h)} &= \mathbf{W}_{\text{UV}}^{(h)} \mathbf{c}_i^{\text{KV}}, \end{aligned}

where WUK(h),WUV(h)Rdhead×dc\mathbf{W}_{\text{UK}}^{(h)}, \mathbf{W}_{\text{UV}}^{(h)} \in \mathbb{R}^{d_{\text{head}} \times d_c} are learned up-projection matrices,2 projecting from dcd_c to dheadd_{\text{head}}.

For queries, MLA similarly uses a compressed representation:

ciQ=WDQxi,qi(h)=WUQ(h)ciQ,\begin{aligned} \mathbf{c}_i^{\text{Q}} &= \mathbf{W}_{\text{DQ}} \mathbf{x}_i, \\ \mathbf{q}_i^{(h)} &= \mathbf{W}_{\text{UQ}}^{(h)} \mathbf{c}_i^{\text{Q}}, \end{aligned}

where WDQRdc×dmodel\mathbf{W}_{\text{DQ}} \in \mathbb{R}^{d_c' \times d_{\text{model}}} is a learned down-projection matrix, and WUQ(h)Rdhead×dc\mathbf{W}_{\text{UQ}}^{(h)} \in \mathbb{R}^{d_{\text{head}} \times d_c'} is a learned up-projection matrix.

During inference, we only need to cache the latent vectors ciKV\mathbf{c}_i^{\text{KV}}, not the full keys and values. While caching the full keys and values as in MHA requires O(nseqnlayersnheadsdhead)O(n_{\text{seq}} \cdot n_{\text{layers}} \cdot n_{\text{heads}} \cdot d_{\text{head}}) memory, caching the latent vectors requires only O(nseqnlayersdc)O(n_{\text{seq}} \cdot n_{\text{layers}} \cdot d_c) memory, where dcnheadsdheadd_c \ll n_{\text{heads}} \cdot d_{\text{head}}.

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:

scoreij(h)qi(h)kj(h)=(WUQ(h)ciQ)(WUK(h)cjKV)=(ciQ)(WUQ(h))WUK(h):=WUQK(h)Rdc×dccjKV.\begin{aligned} \text{score}_{i \rightarrow j}^{(h)} &\propto \mathbf{q}_i^{(h)} \cdot \mathbf{k}_j^{(h)} \\ &= (\mathbf{W}_{\text{UQ}}^{(h)} \mathbf{c}_i^{\text{Q}})^{\top} (\mathbf{W}_{\text{UK}}^{(h)} \mathbf{c}_j^{\text{KV}}) \\ &= \left(\mathbf{c}_i^{\text{Q}}\right)^{\top} \underbrace{\left(\mathbf{W}_{\text{UQ}}^{(h)}\right)^{\top} \mathbf{W}_{\text{UK}}^{(h)}}_{:= \mathbf{W}_{\text{UQK}}^{(h)} \in \mathbb{R}^{d_c' \times d_c}} \mathbf{c}_j^{\text{KV}}. \end{aligned}

Thus, we can “roll” WUK(h)\mathbf{W}_{\text{UK}}^{(h)} into WUQ(h)\mathbf{W}_{\text{UQ}}^{(h)}, and just compute affinity scores between the compressed query and key vectors.

We can similarly “roll” WUV(h)\mathbf{W}_{\text{UV}}^{(h)} into WO(h)\mathbf{W}_{\text{O}}^{(h)}.

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 nseqn_{\text{seq}}, attention considers every preceding key. This means that the computation at position nseqn_{\text{seq}} is O(nseq)O(n_{\text{seq}}), and so producing nseqn_{\text{seq}} tokens autoregressively requires O(nseq2)O(n_{\text{seq}}^2) 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 kk candidates, then the expensive full-attention step becomes O(k)O(k) 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-kk candidates with the highest scores. Note that this still requires scanning all previous tokens to find the candidates, and therefore is still O(nseq)O(n_\text{seq}). 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-kk candidates.

For the query at position ii and each preceding position jj, the indexer computes:

Iij=h=1HIwi,hIReLU ⁣(qi,hIkjI),I_{i \rightarrow j} = \sum_{h=1}^{H_I} w_{i,h}^I \cdot \text{ReLU}\!\left( \mathbf{q}_{i,h}^I \cdot \mathbf{k}_j^I \right),

where HIH_I is the number of indexer heads, qi,hI,kjIRdI\mathbf{q}_{i,h}^I, \mathbf{k}_j^I \in \mathbb{R}^{d_I} are low-dimensional query and key projections, and wi,hIRw_{i,h}^I \in \mathbb{R} is a learned scalar weight derived from xi\mathbf{x}_i that controls how much head hh’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 kk positions with the highest scores become the candidate set. Full attention then runs over only those kk 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 nheadsn_{\text{heads}} 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 CC. Since the query at position LL pairs with all LL preceding positions (including itself), the total cost is simply CLC \cdot L.

For vanilla attention, every pair goes through the full multi-head mechanism, so:

FLOPsvanilla(L)=CattnL.\text{FLOPs}_\text{vanilla}(L) = C_\text{attn} \cdot L.

For DSA, two different operations run at different costs: the indexer scans all LL tokens (at CidxC_\text{idx} per pair), then full attention runs on only kk of them (at CattnC_\text{attn} per pair):

FLOPsDSA(L)=CidxL+Cattnk.\text{FLOPs}_\text{DSA}(L) = C_\text{idx} \cdot L + C_\text{attn} \cdot k.

Let’s estimate CattnC_\text{attn} and CidxC_\text{idx} for DeepSeek-V3.2. For each (query, key) pair, full MLA attention does two things per head: a QK dot product over dc=512d_c = 512 dimensions (2dc2d_c FLOPs), and a value aggregation step that scales the dcd_c-dimensional cached latent by the attention weight and accumulates it (another 2dc2d_c FLOPs).3 That’s 4dc4d_c per head, giving:

Cattn4nheadsdc=4128512=262,144 FLOPs per pair.C_\text{attn} \approx 4 \cdot n_\text{heads} \cdot d_c = 4 \cdot 128 \cdot 512 = 262{,}144 \text{ FLOPs per pair.}

The indexer only scores (no value aggregation) using HI=64H_I = 64 heads with dI=128d_I = 128-dimensional dot products:

Cidx2HIdI=264128=16,384 FLOPs per pair.C_\text{idx} \approx 2 \cdot H_I \cdot d_I = 2 \cdot 64 \cdot 128 = 16{,}384 \text{ FLOPs per pair.}

The ratio Cidx/Cattn1/16C_\text{idx} / C_\text{attn} \approx 1/16: 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.

Interactive FLOP comparison: vanilla MLA (CattnLC_\text{attn} \cdot L) vs. DSA (CidxL+CattnkC_\text{idx} \cdot L + C_\text{attn} \cdot k). For L<kL < k, we assume DSA falls back to full attention over all token positions.

References

References cited in the text are listed first, in order of citation; additional references follow, ordered alphabetically.

  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.
  2. Fast transformer decoding: One write-head is all you need [link]
    Noam Shazeer. arXiv preprint. 2019.
  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.
  4. DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model [link]
    DeepSeek-AI. arXiv preprint. 2024.
  5. DeepSeek-V3.2: Pushing the frontier of open large language models [link]
    DeepSeek-AI. arXiv preprint. 2025.
  6. 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.
  7. An analogy for understanding transformers [link]
    Callum McDougall. 2023.
  8. The annotated transformer [link]
    Sasha Rush, Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman. 2022.

Footnotes

  1. Some prefer the reverse convention, jij \rightarrow i, since information flows from jj to ii. I find iji \rightarrow j more natural for attention patterns. 

  2. Note that the head-specific matrices WUK(h)\mathbf{W}_{\text{UK}}^{(h)} and WUV(h)\mathbf{W}_{\text{UV}}^{(h)} may not actually be “up-projection” matrices, as we previously specified that dcnheadsdheadd_c \ll n_{\text{heads}} \cdot d_{\text{head}}, not necessarily that dcdheadd_c \ll d_{\text{head}}. The original paper works with WUK\mathbf{W}_{\text{UK}} and WUV\mathbf{W}_{\text{UV}} as matrices projecting from dcd_c to nheadsdheadn_{\text{heads}} \cdot d_{\text{head}}, rather than notating a separate matrix for each head, and through this lens WUK\mathbf{W}_{\text{UK}} and WUV\mathbf{W}_{\text{UV}} are true “up-projection” matrices. 

  3. A dd-dimensional dot product requires dd multiplications and dd additions, totaling 2d2d FLOPs. The value aggregation has the same cost structure: dd scalar-by-vector multiplications plus dd additions to an accumulator. This estimate excludes the smaller RoPE score term.