Attention Mechanisms

Attention lets each position combine information from other positions using learned compatibility scores. A query is compared with keys, softmax turns the scores into weights, and those weights combine the value vectors.

Think of it as a soft key-value lookup. The useful questions are what gets scored, which positions are allowed to contribute, and how the resulting values are mixed.

Think of a soft key-value store. Each token emits three vectors:

• a query (q): "what am I looking for?"
• a key (k): "what do I offer?"
• a value (v): "what I'll hand over if you attend to me."

For a given query, you score it against every key (dot product = similarity), turn those scores into a probability distribution with softmax, and return the weighted sum of the values. A hard dictionary returns exactly one value for an exact key match; attention returns a blend, weighted by match strength, which is what makes it differentiable and trainable.

"Self-attention" just means Q, K, and V are all projected from the same sequence: each token attends to the whole sequence, including itself.

Scaled dot-product attention, for queries $Q \in \mathbb{R}^{n \times d_k}$, keys $K \in \mathbb{R}^{m \times d_k}$, and values $V \in \mathbb{R}^{m \times d_v}$:

The softmax is taken row-wise, so each query's attention weights sum to 1. Why divide by $\sqrt{d_k}$? If the components of $q$ and $k$ are independent with mean 0 and variance 1, their dot product $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ has variance $d_k$. For large $d_k$, the raw scores grow large in magnitude, pushing softmax into regions where its gradient is almost zero (it saturates toward a one-hot vector). Scaling by $\sqrt{d_k}$ normalizes the variance back to ~1, keeping gradients healthy.

Key takeaway: Attention is a soft, differentiable dictionary lookup; the $\sqrt{d_k}$ divisor exists purely to keep the variance of the scores near 1 so softmax doesn't saturate and kill gradients.

Multi-head attention runs $h$ attention operations in parallel on lower-dimensional projections, then concatenates:

Each head can specialize (one tracks syntax, another coreference, etc.). With $d_k = d_v = d_{\text{model}} / h$, multi-head costs roughly the same as single-head of full width.

1. Conceptual: Why does attention scale scores by $\sqrt{d_k}$, and what fails if you don't? (Variance of the dot product grows with $d_k$; unscaled scores saturate softmax and kill gradients.)
2. Implementation: How do you implement a causal mask, and where in the computation does it go? (Set future positions to $-\infty$ in the score matrix before softmax, so their weights become 0.)
3. Applied: What's the time and memory complexity of self-attention in sequence length $n$, and why is that a problem for long contexts? ($O(n^2 d)$ time and $O(n^2)$ memory from the $n \times n$ score matrix: the motivation for FlashAttention, sparse, and linear-attention variants.)
4. Systems-level: What does the KV cache do at inference, and why does it make autoregressive decoding far cheaper? (Caches past keys/values so each new token attends to stored K/V instead of recomputing them: turns per-token cost from quadratic-recompute into linear.)
5. Failure modes: Why use multiple heads instead of one wide head? (Separate subspaces let heads attend to different relations simultaneously; a single softmax can only express one weighting per query.)

Close this page. Derive scaled dot-product attention from scratch: write the formula, state the shapes of $Q$, $K$, $V$, explain the $\sqrt{d_k}$ term, and describe in one sentence how multi-head differs. Then check what you missed against Stage 3.