Types of Attention Mechanisms

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Multi-Head Attention (MHA)
5. Multi-Query Attention (MQA)
6. Grouped-Query Attention (GQA)
7. Sliding Window Attention
8. Sparse Attention Patterns
9. Linear Attention
10. Cross-Attention
11. FlashAttention (IO-Aware)
12. Comparison Table
13. Common Pitfalls
14. Summary
15. Exercises

Learning Objectives

1. Distinguish MHA, MQA, and GQA and derive their KV-cache sizes.
2. Derive sliding window attention and compute its effective receptive field.
3. Explain how linear attention avoids the quadratic bottleneck.
4. Understand FlashAttention's IO-aware tiling strategy.
5. Choose the appropriate attention variant for a given constraint (memory, latency, context length).

Notation

• $H$ — number of attention heads
• $G$ — number of KV groups (for GQA)
• $w$ — window size (sliding window)
• $T$ — sequence length
• $d_k$ — dimension per head

Core Intuition

Standard multi-head attention is powerful but has two bottlenecks: (1) $O(T^2)$ memory for the attention matrix, and (2) large KV-cache at inference. Different attention variants trade off expressiveness for efficiency, each targeting a specific bottleneck.

Multi-Head Attention (MHA)

Each head has independent Q, K, V projections:

KV-cache size per token: $2 \times H \times d_k \times L$ values (2 for K and V, across all layers).

Used in: Original Transformer, BERT, GPT-2, GPT-3.

Multi-Query Attention (MQA)

Key idea: All heads share a single set of K, V projections; only Q is per-head.

KV-cache reduction: $H\times$ smaller (only 1 K/V head instead of $H$).

Tradeoff: Slight quality loss (~1% on benchmarks) for massive inference speedup. KV-cache for $T=4096, L=32, d_k=128, H=32$: MHA needs 1 GB, MQA needs 32 MB.

Used in: PaLM, Falcon, StarCoder.

Grouped-Query Attention (GQA)

Key idea: Group $H$ heads into $G$ groups. Each group shares K, V.

where $g(i) = \lfloor iG/H \rfloor$ maps head $i$ to its group.

Interpolation:

• $G = H$: full MHA (every head has its own K, V)
• $G = 1$: MQA (all heads share one K, V)
• $1 < G < H$: GQA (balanced tradeoff)

KV-cache size: $2 \times G \times d_k \times L$ per token. For $G=8, H=32$: $4\times$ reduction vs MHA.

Used in: LLaMA-2 (70B), Mistral, Gemma.

Sliding Window Attention

Each token attends only to the $w$ nearest tokens:

Complexity: $O(Tw)$ instead of $O(T^2)$ — linear in $T$ for fixed $w$.

Effective receptive field: After $L$ layers of window size $w$, a token can "see" up to $L \times w$ positions away (information propagates through the residual stream).

Used in: Mistral (window $w=4096$ with full attention every few layers), Longformer.

Sparse Attention Patterns

Fixed patterns (predefined sparsity):

• Strided: Attend to every $s$-th position → $O(T\sqrt{T})$.
• Block-sparse: Divide into blocks, attend within block + selected global tokens.
• BigBird: Combination of random, window, and global tokens.

Learned patterns:

• Routing Transformer: Cluster tokens via k-means, attend within clusters.
• Reformer: LSH-based grouping of similar tokens.

Complexity: Generally $O(T\sqrt{T})$ or $O(T\log T)$.

Linear Attention

Replace softmax with a kernel decomposition:

Standard: $\text{Attn}(\mathbf{q}, \mathbf{K}, \mathbf{V}) = \frac{\sum_j\exp(\mathbf{q}^T\mathbf{k}_j)\mathbf{v}_j}{\sum_j\exp(\mathbf{q}^T\mathbf{k}_j)}$

Linear: $\text{Attn}(\mathbf{q}, \mathbf{K}, \mathbf{V}) = \frac{\phi(\mathbf{q})^T\sum_j\phi(\mathbf{k}_j)\mathbf{v}_j^T}{\phi(\mathbf{q})^T\sum_j\phi(\mathbf{k}_j)}$

where $\phi$ is a feature map (e.g., $\phi(x) = \text{elu}(x) + 1$).

Key insight: Compute $\sum_j\phi(\mathbf{k}_j)\mathbf{v}_j^T$ once ($O(Td_kd_v)$), then query in $O(d_kd_v)$ each. Total: $O(Td^2)$ — linear in $T$.

Tradeoff: Expressiveness is reduced; performance gap vs softmax attention, especially for long-range dependencies.

Used in: Linear Transformer, RWKV, RetNet.

Cross-Attention

Queries from one sequence, keys/values from another:

Use cases:

• Encoder-decoder models (T5, BART)
• Text-conditioned image generation (Stable Diffusion)
• Retrieval-augmented generation (retrieved docs as keys/values)

FlashAttention (IO-Aware)

FlashAttention computes exact attention but reorders computation to minimize GPU memory reads/writes:

Problem: Standard attention materializes the $T \times T$ attention matrix in HBM (slow GPU memory).

Solution: Tile the computation into blocks that fit in SRAM (fast on-chip memory):

1. Load blocks of Q, K, V from HBM to SRAM.
2. Compute local attention for each block (with online softmax normalization).
3. Write output back to HBM.

Result: $O(T^2d)$ FLOPs (same) but $O(T)$ HBM memory (vs $O(T^2)$). Enables 2–4x speedup and much longer sequences.

Used in: Nearly all modern LLM training and inference.

Comparison Table

• MHA: Full expressiveness, $O(T^2d)$ compute, $O(H \cdot d_k)$ KV per token per layer
• MQA: Shared K/V, $O(T^2d)$ compute, $O(d_k)$ KV per token per layer
• GQA: Grouped K/V, $O(T^2d)$ compute, $O(G \cdot d_k)$ KV per token per layer
• Sliding Window: $O(Twd)$ compute, $O(w \cdot d_k)$ KV per token
• Linear: $O(Td^2)$ compute, $O(d^2)$ state per layer
• Flash: $O(T^2d)$ compute, $O(T)$ memory (IO-optimal)

Common Pitfalls

Pitfall 1. Assuming linear attention is always better for long sequences. While asymptotically better, it often underperforms softmax attention on standard benchmarks due to reduced expressiveness.

Pitfall 2. Using MQA without quality validation. For smaller models, the quality drop from shared K/V can be significant. GQA with $G=4$–$8$ is often a better tradeoff.

Pitfall 3. Confusing algorithmic complexity with wall-clock time. FlashAttention has the same FLOP count but is much faster due to better memory access patterns.

Summary

• MHA: Maximum expressiveness, baseline cost.
• MQA/GQA: Reduce KV-cache for efficient inference.
• Sliding window: Linear complexity with bounded context.
• Linear attention: Replace softmax with kernel tricks for $O(T)$.
• FlashAttention: IO-optimal implementation of exact attention.
• Choose based on your bottleneck: training memory → Flash; inference KV → GQA; context length → sparse/linear.

Exercises

Exercise 1. Compute the KV-cache size (in GB) for MHA vs GQA ($G=8$) for a model with $L=32, H=32, d_k=128, T=32768$.

Exercise 2. Derive the effective receptive field after $L=32$ layers of sliding window attention with $w=4096$.

Exercise 3. Show that linear attention with $\phi(\mathbf{x}) = \mathbf{x}$ (identity) is equivalent to a linear layer.

Exercise 4. Prove that GQA with $G=H$ is identical to MHA and with $G=1$ is identical to MQA.

Exercise 5. Explain why FlashAttention uses $O(T)$ HBM memory despite computing the exact $T \times T$ attention.