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" means Q, K, and V are projected from the same sequence. Each token can attend to positions allowed by the mask; causal self-attention excludes future positions.

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 counters dimension-dependent score growth under the assumptions above; it does not guarantee unit variance or prevent saturation for arbitrary learned Q and K.

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

Each head can learn a different weighting; interpretable specializations are possible, not guaranteed. With $d_k = d_v = d_{\text{model}} / h$, multi-head costs roughly the same as a single full-width head. With equivalent Q/K/V and output projections, both have $4d_{\text{model}}^2$ projection weights, excluding biases. Splitting the width across heads does not add parameters.

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? (Dense attention takes $O(n^2 d)$ arithmetic, plus $O(nd^2)$ projections. A naive implementation stores quadratic score matrices; FlashAttention avoids storing them in device memory but retains quadratic arithmetic. Sparse and linear variants change the attention computation.)
4. Systems-level: What does the KV cache do at inference, and why does it make autoregressive decoding cheaper? (It reuses past keys/values. Dense projections for the new token cost $O(d^2)$ per layer; its attention over a prefix of length $t$ still costs $O(td)$.)
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.

Sources and further reading

• Vaswani et al., Attention Is All You Need, §3.2: attention equations, head dimensions, and projection shapes.
• Dao et al., FlashAttention: exact attention with reduced memory traffic.