Attention Variants

The original multi-head attention (MHA) works well but is expensive — especially at inference time where KV cache memory scales with head count. This chapter traces the evolution of attention variants, each designed to solve a specific infrastructure bottleneck.

---

1. The KV Cache Problem

During autoregressive generation, the model stores Key and Value tensors for all previous tokens (the KV cache) to avoid recomputation. Per token per layer:

$$\text{KV cache per token} = 2 \times N_{kv} \times d_k \times \text{bytes\_per\_element}$$

For LLaMA-2-7B with MHA ($N_{kv}=32, d_k=128$, FP16):

• Per token per layer: $2 \times 32 \times 128 \times 2 = 16$ KB
• Per token (32 layers): $32 \times 16$ KB = 512 KB per token
• For 4K context: $4096 \times 512$ KB = 2 GB of KV cache alone

At 128K context, that’s 64 GB — most of an H100’s memory consumed by KV cache for a single request. Every attention variant below exists to reduce this cost.

---

2. Multi-Query Attention (MQA)

Idea: All Q heads share a single K and V head.

$$N_h \text{ Q heads}, \quad 1 \text{ K head}, \quad 1 \text{ V head}$$

KV cache reduction: $N_h \times$ smaller (e.g., 32× for 32 heads).

Tradeoff: Quality degrades because all Q heads share a single K/V representation. Different heads can no longer learn independent attention patterns (e.g., one for syntax, another for semantics). Used by PaLM, Falcon, but largely superseded by GQA.

---

3. Grouped-Query Attention (GQA)

Idea: A middle ground — group Q heads and have each group share one K/V head.

$$N_h \text{ Q heads}, \quad N_{kv} \text{ K/V heads} \quad (N_{kv} < N_h)$$

Each group of $N_h / N_{kv}$ Q heads shares one K/V head.

Config	Q heads	KV heads	Group size	KV cache reduction
MHA	64	64	1	1× (baseline)
GQA-8	64	8	8	8×
MQA	64	1	64	64×

Used by: LLaMA-2-70B (8 KV heads), Mistral-7B (8 KV heads), LLaMA-3 (8 KV heads).

Quality vs MHA: Negligible difference in practice — the KV cache savings far outweigh any quality loss.

---

4. Multi-Latent Attention (MLA)

Idea (DeepSeek-V2/V3): Instead of storing full K and V, compress them into a low-dimensional latent vector and reconstruct on the fly.

$$\mathbf{c}_t = \mathbf{h}_t \mathbf{W}^{DKV} \qquad \text{(compress: } d \to d_c \text{)}$$ $$\mathbf{K}_t = \mathbf{c}_t \mathbf{W}^{UK}, \qquad \mathbf{V}_t = \mathbf{c}_t \mathbf{W}^{UV} \qquad \text{(reconstruct during attention)}$$

Only the compressed $\mathbf{c}_t$ (dimension $d_c \ll d$) is stored in KV cache.

KV cache reduction: With $d_c = 512$ vs $N_{kv} \times d_k = 8 \times 128 = 1024$ (GQA-8), MLA achieves ~2× further reduction on top of GQA.

Key insight for infrastructure: The reconstruction adds compute (extra matrix multiplies during attention), but eliminates memory — trading compute for memory. At long contexts where KV cache is the bottleneck, this is a winning trade.

---

5. Linear Attention

Standard attention is $O(S^2)$ in both compute and memory (due to the $S \times S$ attention matrix). Linear attention avoids materializing this matrix:

$$\text{Standard}: \quad \text{softmax}(QK^\top)V \qquad O(S^2d)$$ $$\text{Linear}: \quad \phi(Q)(\phi(K)^\top V) \qquad O(Sd^2)$$

By changing the order of operations (compute $K^\top V$ first, which is $d \times d$, then multiply by $Q$), complexity drops from $O(S^2d)$ to $O(Sd^2)$.

Tradeoff: Linear attention approximates softmax attention — quality is generally lower, especially for tasks requiring precise token-to-token attention patterns. No production LLM uses pure linear attention, but hybrid architectures are emerging.

---

6. State-Space Models (Mamba)

Idea: Replace attention entirely with a recurrence that processes tokens sequentially with a fixed-size state:

$$h_t = \bar{A} h_{t-1} + \bar{B} x_t, \qquad y_t = C h_t$$

Where $\bar{A}, \bar{B}$ are discretized state transition matrices and $h_t$ is a fixed-size hidden state.

Advantages:

• $O(S)$ compute and memory (linear in sequence length)
• No KV cache — fixed state size regardless of context length
• Efficient parallel training via the “parallel scan” algorithm

Disadvantages:

• Cannot do exact “lookup” attention (e.g., “find the name mentioned in paragraph 3”)
• Quality gap vs Transformers on retrieval-heavy tasks

Hybrid architectures (Jamba, Zamba) interleave Mamba layers with attention layers. Typical ratio: 1 attention layer per 4-8 Mamba layers. Attention handles retrieval-heavy operations (“find the name in paragraph 3”) while Mamba handles bulk sequential processing. Jamba (AI21, 2024) uses a 1:7 ratio with 52B parameters, achieving competitive quality with much lower KV cache than a pure Transformer.

Infrastructure Implications

Architecture	Memory (inference)	Compute (inference)	Best for
MHA	$O(S)$ KV cache	$O(S^2)$ per step	Quality-critical, short context
GQA	$O(S/g)$ KV cache	$O(S^2)$ per step	Production LLMs (LLaMA, Mistral)
MLA	$O(S \cdot d_c/d)$ KV cache	$O(S^2)$ + reconstruct	Very long context (DeepSeek)
Linear Attn	$O(d^2)$ state	$O(d^2)$ per step	Streaming, infinite context
SSM (Mamba)	$O(d)$ state	$O(d)$ per step	Very long context, streaming

---

Self-Check

• Calculate KV cache size for MHA, GQA, and MLA for a given model and context length
• Explain why GQA has minimal quality loss compared to MHA
• Explain MLA’s “compress then reconstruct” approach and the compute-memory tradeoff
• Explain why linear attention changes the order of operations
• Explain why SSMs have O(1) memory per token but struggle with retrieval tasks
• Explain why hybrid architectures (Mamba + Attention) exist

---

References

• GQA: Training Generalized Multi-Query Transformer Models (Ainslie et al., 2023)
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Gu & Dao, 2023)
• Jamba: A Hybrid Transformer-Mamba Language Model (AI21, 2024)