Attention Variants: Multi-Head / Cross / Sparse / Linear

L3-05 covered standard Self-Attention:

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

This piece: the attention family beyond that formula.

1. Multi-Head Attention

One sentence: run multiple Attentions in parallel, concat, project.

Why

A single Attention learns one “relationship pattern” — neighbor, syntax, co-occurrence. Multiple heads → each learns different patterns → richer.

Math

Single head:  head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)

Multi-head:   MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O

• $h$ = number of heads (typical 8, 12, 16, 32)
• Each head’s dim $d_k = d_{\text{model}} / h$
• Total parameters ≈ single head (each head is smaller)

Intuition

Like having 8 analysts read the same paragraph:

• Head 1: subject-verb relationships
• Head 2: pronoun resolution
• Head 3: sentiment
• …

Then stitch their notes together.

In practice: mechanistic interpretability work (L6-04) shows different heads do learn different things.

2. Self vs Cross Attention

Self-Attention

Q, K, V all come from the same sequence:

encoder input → Q, K, V → Attention(Q, K, V)

Used in: BERT, GPT, ViT, etc.

Cross-Attention

Q comes from one sequence, K and V come from another:

Q from decoder
K, V from encoder

→ Attention(Q_dec, K_enc, V_enc)

Lets decoder “query” the encoder. Used in:

• Original Transformer encoder-decoder (translation)
• Stable Diffusion’s text → image conditioning

3. Causal Attention (Masking)

GPT-class autoregressive models need: don’t see future tokens when predicting the next.

Standard attention matrix    causal mask
┌──────────────┐         ┌──────────────┐
│ * * * * * *  │         │ *            │
│ * * * * * *  │         │ * *          │
│ * * * * * *  │   →     │ * * *        │
│ * * * * * *  │         │ * * * *      │
│ * * * * * *  │         │ * * * * *    │
│ * * * * * *  │         │ * * * * * *  │
└──────────────┘         └──────────────┘
                          Lower triangle; upper triangle masked

Implementation: set upper-triangle attention scores to $-\infty$ before softmax → those positions get 0 weight.

mask = torch.triu(torch.ones(L, L), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))
attn = F.softmax(scores, dim=-1)

All decoder-only LLMs use causal attention.

4. Sparse Attention

Problem: O(n²) is the bottleneck

Standard attention computes $n \times n$ matrix.

• 1k tokens → 1M ops
• 10k → 100M
• 100k → 10G ops

Long-context costs come mostly from this.

Sliding Window

Each token only attends to nearby $w$ tokens. $O(n \cdot w)$.

• Pro: linear in n
• Con: loses long-range
• Used in: Longformer, Mistral

Strided / Dilated

Fixed-stride hops to attend to distant tokens. Combine local + global → captures both.

Global Tokens

Some tokens (like [CLS]) attend everywhere; rest sparse. Tradeoff.

Used in: BigBird, ETC.

5. Linear Attention

Standard:

$$\text{softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) V$$

Bottleneck: $QK^T$ — $O(n^2 d)$.

Linear attention approximates this with $O(n d^2)$ (much smaller when n >> d):

$$\text{attn} \approx \phi(Q)(\phi(K)^T V)$$

Where $\phi$ is a kernel function (e.g., ELU+1, random features).

Notable work:

• Performer (2020): random feature approximation of softmax
• Linear Transformer (2020): kernel-based
• Mamba (2024): state-space-based linear recurrence

But: linear attention is measurably weaker than standard for hard long-context tasks. Most production LLMs still use standard + FlashAttention.

6. FlashAttention (Same Math, Faster)

L7-01 / L7-03 covered this. Quick recap:

Same softmax(QK^T)V computation, but keep intermediate results in GPU SRAM not HBM. 2-5× speedup with no precision loss.

This is algorithmic identity, engineering optimization — and is used in all modern Transformer training (FlashAttention 2/3).

7. Multi-Query / Grouped-Query Attention (GQA)

Problem: KV cache explosion

In inference, each head has its own K, V cache. 32 heads → 32× memory for long context.

Multi-Query Attention (2019)

All heads share one K, V — only Q is per-head:

Standard MHA: 32 (Q, K, V) heads
MQA:          32 Q heads + 1 shared (K, V)

KV cache 32× smaller. But quality drops noticeably.

Grouped-Query Attention (2023, compromise)

Split 32 Q heads into 8 groups, each group shares K, V:

32 Q heads + 8 sets of (K, V)

KV cache 4× smaller, quality almost unchanged.

Used by LLaMA 2/3, Mistral, etc.

8. Rotary Position Embedding (RoPE)

L3-05 didn’t dwell on positional encoding. Quick coverage:

Self-Attention itself doesn’t know order — must inject position.

Method	How
Absolute (original Transformer)	Add fixed sin/cos vector per position
Learned (BERT, GPT-2)	Learn an embedding per position
Relative (T5)	Bias in attention based on distance
RoPE (LLaMA, Qwen, etc.)	Rotate Q, K in complex plane by angle = position

RoPE advantages:

• Naturally extrapolates to longer sequences
• No extra parameters
• Fast to compute

Standard for modern LLMs.

9. The Future

State Space Models (SSM)

Mamba — drop attention entirely, use linear recurrence. 5-10× faster, quality close to Transformer.

Hybrid Architectures

Some layers attention, others SSM. Examples: Jamba (Mamba + Transformer), Falcon Mamba.

Long-context Techniques

Ring Attention, Sequence Parallel, sparse + chunked combinations — let 1M-token context be feasible.

Selection Quick-Reference

Scenario	Use
Standard LLM training	MHA + GQA + RoPE + FlashAttention
Long-context needed	Sliding window or SSM
Multimodal (image + text)	Self-Attention + Cross-Attention combination
Maximum inference speed	MQA / GQA
Research new directions	Mamba-style SSM

Next recommended: L3-08 Transformer Architecture or L3-09 BERT vs GPT.