The attention mechanism

Summary

Attention forms one permitted weighted value mixture for each query. Queries $Q\in {\mathbb{R}}^{n_q\times d}$ and keys $K\in {\mathbb{R}}^{n_k\times d}$ determine weights. Values $V\in {\mathbb{R}}^{n_k\times d_v}$ supply the mixed payload; key and value rows must remain paired.

Normalize each query’s allowed scores separately, then multiply the resulting $n_q\times n_k$ matrix by $V$. Without dropout or later projections, every nonempty row is a convex value combination. This primitive appears in sequence, vision, graph, and memory models; an architecture’s dominant cost depends on its actual dimensions.

The mechanism step by step

For a single query $q\in {\mathbb{R}}^d$ and set of $n$ key-value pairs:

1. Score each key against the query: $s_i=q^{\top }{k}_i/\sqrt{d}$.
2. Normalize with softmax: ${\alpha }_i=\exp (s_i)/{\sum }_j\exp (s_j)$. The ${\alpha }_i$ sum to 1. They form an attention distribution.
3. Aggregate: output $={\sum }_i{\alpha }_i{v}_i$.

Take one-dimensional query $q=1$, keys $(\log 6,\log 3,0)$, and values $(1,4,7)$. Exponentiated scores are $(6,3,1)$, so weights are $(0.6,0.3,0.1)$. The output is $0.6(1)+0.3(4)+0.1(7)=2.5$. Changing a value changes the mixture’s payload without changing these query-key scores.

If corresponding query-key coordinates are independent, zero-mean, and unit-variance, their summed dot product has variance $d$. Dividing by $\sqrt{d}$ gives unit variance under those assumptions. Learned correlations or different coordinate scales can change that analysis; scaling is not a universal nonsaturation guarantee.

Self-attention vs. cross-attention

• Self-attention: queries, keys, and values derive from the same sequence; a mask may restrict visible positions.
• Cross-attention: queries derive from one source and paired keys/values from another; their sequence lengths can differ.

Multi-head attention

Apply learned projections for several heads, mix values separately, concatenate the outputs, and apply an output projection. Head dimensions and counts are architecture choices; semantic specialization must be demonstrated rather than assumed.

In modern LLMs, heads are reduced via grouped-query attention where multiple Q heads share K/V heads.

Causal (autoregressive) masking

New task: Exclude future or padded positions from each mixture.

Picture and change: Set forbidden scores to $-\infty$ before normalization; keep absolute position and cache-offset conventions explicit.

Check: Excluding the third key in the example renormalizes weights to $(2/3,1/3,0)$ and changes output to $2$. Simply zeroing its old weight without renormalization gives $1.8$, a different operation.

Use when: Only a permitted subset may contribute. An entirely masked row has no categorical distribution; reject it or define a separate zero-output policy explicitly.

Connection to retrieval

New task: Return one retrieved item instead of a value mixture.

Picture and change: Replace aggregation with the maximum-score key’s associated item.

Check: The example selects the first key and value $1$, not the attention output $2.5$.

Use when: A hard candidate is required. Positive-temperature softmax preserves score order, but a weighted sum is not nearest-neighbor retrieval. Cosine and raw dot-product rankings agree only under appropriate normalization, not in general.

Cost

• Score and mixture arithmetic costs $O(n_q{n}_k(d+d_v))$, excluding projections.
• Explicit scores and weights need $O(n_q{n}_k)$ intermediate storage, in addition to inputs and $O(n_q{d}_v)$ output.

FlashAttention tiles exact dense attention without retaining the full score matrix. Its storage still includes dimension-dependent inputs, outputs, row statistics, and tile workspace.