Attention and Self-Attention

Phase 2 · Document 02 · Transformer Foundations Prev: 01 — Token Embeddings · Next: 03 — Feed-Forward Layers

Table of Contents

1. Why This Matters
2. Core Concept
3. Mental Model
4. Hitchhiker's Guide
5. Warmup Readings
6. Deep Readings and External References
7. Key Terms
8. Important Facts
9. Observations from Real Systems
10. Common Misconceptions
11. Engineering Decision Framework
12. Hands-On Lab
13. Verification Questions
14. Takeaways
15. Artifact Checklist

1. Why This Matters

Attention is the mechanism — the reason transformers beat everything before them, and the source of their biggest costs. It explains why context length is quadratic in the naive case, why the KV cache exists, why FlashAttention and sliding-window attention were invented, and why long context is expensive. Almost every serving optimization you'll meet (Phase 7) is an attempt to make attention cheaper. You don't need to derive it, but understanding what it computes — and what it costs — is what separates "I call the API" from "I can reason about why this is slow and expensive."

2. Core Concept

Plain-English primer (the attention formula, built from zero)

The famous formula Attention(Q,K,V) = softmax(QKᵀ/√d_k)·V looks scary but is just the four moves from 2.00's math primer in sequence. Here it is from zero (full math intuition: 2.00).

• Q, K, V are three vectors per token, each made by multiplying the token's vector by a learned matrix:
  • Query (Q) = "what am I looking for?"
  • Key (K) = "what do I offer?"
  • Value (V) = "what information do I carry?" Analogy: searching a library — your query is what you want, each book's key is its catalog entry, its value is the actual contents you take.
• Q·Kᵀ = compare my Query to every Key with a dot product (the similarity score from 2.00). High score = "that token is relevant to me."
• /√d_k = scale down so the scores don't get huge when vectors are long (d_k = the head's dimension). Purely keeps the next step numerically stable.
• softmax(...) = turn those scores into weights that sum to 1 (the softmax from 2.00) — "how much attention to pay to each token."
• · V = take the weighted average of the Values. The result for each token is a blend of the information from the tokens it decided were relevant.

Tiny worked example (3 tokens; token 3 attending back):

import math
def dot(a,b): return sum(x*y for x,y in zip(a,b))
def softmax(s):
    m=max(s); e=[math.exp(x-m) for x in s]; t=sum(e); return [x/t for x in e]

# Query of token 3, Keys/Values of tokens 1,2,3 (toy 2-dim vectors)
q3 = [1.0, 0.0]
K  = [[1.0,0.0], [0.0,1.0], [0.9,0.1]]   # token1 aligns with q3, token2 doesn't
V  = [[10,0],    [0,10],    [9,1]]
d_k = 2
scores  = [dot(q3,k)/math.sqrt(d_k) for k in K]   # relevance of each token to q3
weights = softmax(scores)                          # → ~[0.46, 0.18, 0.36]: tokens 1&3 win
out = [sum(weights[i]*V[i][j] for i in range(3)) for j in range(2)]
# out ≈ [7.8, 1.4]  → mostly the info from tokens 1 and 3 (the relevant ones)

• Self-attention = Q, K, V all come from the same sequence (tokens attending to tokens).
• Causal mask = a token may only attend to itself and earlier tokens (we hide the future by setting its scores to −∞ before softmax) — this is what makes "predict the next token" a fair task.
• Multi-head = run several of these in parallel with different learned Q/K/V matrices, each catching a different kind of relationship (grammar, reference, position), then concatenate.

With that, the rest of this doc is what attention costs and how production shrinks that cost.

Plain English

For each token, attention asks: "Of all the tokens I'm allowed to look at, which ones matter for me right now, and what should I take from them?" It then builds a weighted blend of the others' information. That's how a pronoun finds its referent, a verb finds its subject, or a closing bracket finds its opener.

Technical depth — Q, K, V

Each token's vector is projected (via learned weight matrices) into three roles:

• Query (Q): "what am I looking for?"
• Key (K): "what do I offer?"
• Value (V): "what information do I carry?"

Attention for a token = compare its Q to every K, turn the scores into weights, and take the weighted sum of the Vs:

Attention(Q, K, V) = softmax( Q·Kᵀ / √d_k ) · V

• Q·Kᵀ → a score matrix (token × token): how relevant each token is to each other.
• /√d_k → scaling to keep softmax stable.
• softmax → scores become weights summing to 1.
• ·V → each token's output is the weighted blend of all Values.

Self-attention, causal masking, multi-head

• Self-attention: Q, K, V all come from the same sequence (tokens attending to tokens).
• Causal mask: in decoder-only LLMs, a token may only attend to itself and earlier tokens (the future is masked to −∞ before softmax). This is what makes next-token prediction valid.
• Multi-head attention (MHA): run several attention computations in parallel with different learned projections ("heads"), each capturing different relationships (syntax, coreference, position), then concatenate. More heads ≈ richer relations.

The cost — why everything downstream exists

• Compute: the Q·Kᵀ score matrix is seq_len × seq_len → attention is O(n²) in sequence length. Double the context, quadruple this cost. This is the long-context tax.
• Memory (KV cache): during generation, the K and V of all prior tokens are reused every step, so they're cached — growing linearly with context × layers × heads (06).
• Mitigations you'll meet: FlashAttention (compute attention without materializing the full n² matrix — faster, less memory), GQA/MQA (share K/V across heads to shrink the KV cache), sliding-window attention (each token attends only to a recent window — linear cost, used by Mistral/Gemma).

3. Mental Model

For each token i (a Query), look back at tokens 0..i (Keys):
        scores  = Qᵢ · Kⱼ      (how much j matters to i)
        weights = softmax(scores / √d_k)
        outᵢ    = Σⱼ weightsⱼ · Vⱼ      (blend of what those tokens carry)

CAUSAL MASK:  i can only see j ≤ i  (no peeking at the future)
MULTI-HEAD:   do this H times in parallel with different Q/K/V projections, concat

COST:
  score matrix is n×n  →  compute ∝ n²   (the long-context tax)
  K,V of past tokens reused every step →  CACHE them (KV cache)

CHEAPER VARIANTS:  FlashAttention (no n² materialization) ·
                   GQA/MQA (fewer K/V → smaller cache) · sliding window (local → linear)

4. Hitchhiker's Guide

What to understand first: attention = relevance-weighted blending of other tokens' Values, restricted to the past (causal), done in parallel heads. And that it costs O(n²).

What to ignore at first: the exact softmax/scaling derivation and backprop. Know the shapes and the cost.

What misleads beginners:

• "Attention is memory." It's a weighting over the current context, not stored knowledge (that's the FFN/weights).
• "More heads always better." Beyond a point, heads add memory/compute with diminishing returns; many models now share K/V (GQA) to save memory.
• "Long context is just a bigger number." It's an O(n²) compute tax and a linear KV-memory tax.

How experts reason: they connect attention to cost — "this is slow because the prompt is long (n² prefill)," "we're KV-memory-bound, switch to GQA or shrink context," "enable FlashAttention." They know which attention variant a model uses because it changes serving memory.

What matters in production: whether the model uses GQA/MQA (KV-cache size), sliding-window (long-context cost), and whether your serving engine uses FlashAttention (speed/memory).

How to verify: check the model config for num_key_value_heads (< num_attention_heads ⇒ GQA) and any sliding-window setting; benchmark prefill time vs prompt length to see the n² effect.

Questions to ask: Does this model use MHA, GQA, or MQA? Sliding window? Does the serving stack use FlashAttention? What's the head count and head dim (KV cost)?

What silently gets expensive: long prompts (quadratic prefill), full-MHA models at high concurrency (large KV cache), and disabling FlashAttention on long context.

5. Warmup Readings

6. Deep Readings and External References

7. Key Terms

8. Important Facts

• Attention output = softmax-weighted sum of Values, weights from Q·Kᵀ relevance.
• LLMs use causal self-attention: a token attends only to itself and earlier tokens.
• Multi-head runs several attentions in parallel with different projections, then concatenates.
• Naive attention is O(n²) in sequence length — the core long-context cost.
• GQA/MQA share K/V across query heads, drastically shrinking the KV cache (most modern big models use GQA).
• FlashAttention computes attention without materializing the n² matrix — faster and lower-memory; standard in production.
• Sliding-window attention (e.g. Mistral, Gemma) makes long-context cost roughly linear.
• Attention is weighting over current context, not stored knowledge (that's weights/FFN).

9. Observations from Real Systems

• Llama 2/3, Qwen, Mistral use GQA (num_key_value_heads < num_attention_heads) specifically to shrink the KV cache and serve more concurrency.
• Mistral / Gemma document sliding-window attention to bound long-context cost.
• vLLM, TGI, SGLang ship FlashAttention(-style) kernels by default; it's a baseline expectation, not an exotic option.
• PagedAttention (vLLM) manages the K/V that attention reuses — a direct downstream of this mechanism (06).
• Benchmark dashboards show prefill time growing super-linearly with prompt length — the O(n²) tax made visible.

10. Common Misconceptions

11. Engineering Decision Framework

Diagnosing attention-driven cost:
  Slow first token on long prompts? → O(n²) prefill.
     → shorten/cache prompt; ensure FlashAttention; consider sliding-window model.
  KV cache too big / low concurrency? → too many K/V heads × context.
     → prefer a GQA/MQA model; shorter context; PagedAttention (Phase 7).
  Need cheap long context?           → sliding-window or retrieval (RAG) instead.

Selecting a model for serving (attention lens):
  high concurrency + long context → GQA + FlashAttention + (optionally) sliding window.
  check config: num_attention_heads vs num_key_value_heads, sliding_window, rope settings.

12. Hands-On Lab

Goal

Implement scaled dot-product attention from scratch, visualize an attention weight matrix, and observe the causal mask and the O(n²) score matrix.

Prerequisites

• Python 3.10+, numpy (no GPU needed).

Setup

pip install numpy

Steps

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x); return e / e.sum(axis=axis, keepdims=True)

def attention(Q, K, V, causal=True):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)            # (n, n) — the O(n^2) matrix
    if causal:
        n = scores.shape[0]
        mask = np.triu(np.ones((n, n)), k=1).astype(bool)
        scores[mask] = -1e9                    # block the future
    weights = softmax(scores)
    return weights @ V, weights

np.random.seed(0)
n, d = 5, 8                                    # 5 tokens, head_dim 8
Q = K = V = np.random.randn(n, d)              # toy: same source (self-attention)
out, w = attention(Q, K, V)
print("score matrix shape (n×n):", w.shape)
print("causal weights (row=token, lower-triangular):\n", np.round(w, 2))

Expected output

• w is 5×5, lower-triangular (each token only attends to itself + earlier) and each row sums to 1.

Debugging tips

• Rows not summing to 1? softmax axis wrong.
• Upper triangle nonzero? Mask not applied before softmax.

Extension task

Add multi-head: split d into H heads, run attention per head, concatenate. Then implement GQA by sharing one K/V across two Q heads and note the memory saving.

Production extension

Plot prefill latency vs prompt length against a real endpoint (reuse Phase 1.05 lab) and fit the curve — observe the super-linear (n²) growth.

What to measure

Shape and triangularity of the weight matrix; row sums; (extension) KV-size reduction from GQA; prefill-vs-length curve.

Deliverables

• Working scaled-dot-product attention with causal mask.
• A printed/plotted attention weight matrix.
• A note: how GQA shrinks the KV cache, and evidence of O(n²) prefill.

13. Verification Questions

Basic

1. What do Q, K, and V represent?
2. What does the causal mask enforce and why?
3. What does "multi-head" add?

Applied 4. Why is naive attention O(n²) in sequence length, and what does that imply for long prompts? 5. How does GQA reduce KV-cache memory?

Debugging 6. Prefill time quadruples when you double the prompt. Expected or a bug? Explain. 7. A full-MHA model OOMs under concurrency where a similar GQA model doesn't. Why?

System design 8. Choose attention features (MHA/GQA, sliding window, FlashAttention) for a long-context, high-concurrency service. Justify.

Startup / product 9. Your long-context feature is unprofitable. Explain via attention cost where the money goes and two architectural/product mitigations.

14. Takeaways

1. Attention = relevance-weighted blend of other tokens' Values (softmax(QKᵀ/√d)V), causal in LLMs.
2. Multi-head captures multiple relation types in parallel.
3. Naive attention is O(n²) — the fundamental long-context tax.
4. KV cache exists because attention reuses past K/V every decode step.
5. GQA/MQA, FlashAttention, sliding-window are the cost mitigations you'll meet in production.
6. Attention weights current context; it is not stored knowledge.

15. Artifact Checklist

• Code: from-scratch scaled-dot-product attention (causal).
• Visualization: an attention weight matrix.
• Extension: multi-head + a GQA memory-saving note.
• Benchmark: prefill latency vs prompt length (n² evidence).
• Notes: the attention→cost map (n², KV, mitigations).
• Decision record: attention features for one serving scenario.

Next: 03 — Feed-Forward Layers