Attention Mechanism

Attention starts with a sequence of hidden vectors, one per token position. Three learned linear projections turn each hidden vector into a query, a key, and a value. The query asks what this position is looking for, the key advertises what another position contains, and the value is the information that will be copied if that position is selected.

These are not separate tokens or hand-written features. They are learned coordinate systems. During training, the model discovers projections that make useful relationships easy to score, such as a pronoun looking for its antecedent, an operator looking for its operands, or a closing bracket looking for the opening bracket.

For one token to decide whether another token matters, the model takes the dot product between the first token query and the second token key. If the vectors point in similar directions, the dot product is large. If they are unrelated or opposed, the score is smaller. This gives each query a relevance score against every key.

Across a sequence of length n, doing this for every query-key pair creates an n by n score matrix. Row i answers the question: for token i, how much should it attend to each token j before masking and normalization? This pairwise comparison is the source of both attention power and its quadratic cost.

Raw dot products grow larger as the key and query dimension d_k grows. If those scores become too large, softmax saturates: one position gets almost all the probability, gradients become tiny for the others, and training becomes less stable. Scaled dot-product attention divides each score by sqrt(d_k) before softmax.

Softmax is applied across each row of scores. The result is a set of nonnegative weights that sum to one for each query position. The output for that position is then a weighted sum of value vectors, so high-weight tokens contribute more information and low-weight tokens contribute little.

In self-attention, queries, keys, and values all come from the same sequence. Encoder-style models usually allow every token to attend to every other token. Decoder-only models use causal self-attention, where a mask blocks each position from attending to future positions so next-token prediction cannot leak the answer.

In cross-attention, the query usually comes from a decoder state, while keys and values come from another source, such as an encoder output. This is common in encoder-decoder translation models and can also appear in multimodal systems where text queries attend over image or audio representations.

A single attention head has one query-key-value space. Multi-head attention splits the hidden representation into several smaller heads, runs scaled dot-product attention independently in each head, concatenates the head outputs, and projects the result back to the model hidden size.

The benefit is specialization. One head can learn local syntax, another can track coreference, another can match brackets or quotes, and another can retrieve task-specific tokens. Heads are not guaranteed to be human-interpretable, but the architecture gives the model several parallel channels for different relationships.

Plain attention compares vectors and does not inherently know whether a token came first or last. If the model received the same token vectors in a different order without positional information, the attention computation would not reliably distinguish the original sentence from a shuffled one.

Older transformers added learned or sinusoidal position vectors to token embeddings. Many modern LLMs use rotary position embeddings, or RoPE, which rotate query and key dimensions according to position before the dot product. RoPE makes relative distances visible inside attention scores and tends to work better for long-context decoder-only models than simple absolute position tables.

If query and key components have roughly unit variance, their dot product has variance proportional to d_k. Larger head dimensions therefore produce larger raw scores. Large scores push softmax into sharp, saturated distributions where one entry dominates and gradients for the rest become weak.

Dividing by sqrt(d_k) keeps score magnitudes in a healthier range across different head sizes. This does not change which keys are more aligned with a query, but it changes the temperature of the softmax so training is more stable and attention does not collapse too early.

For a prompt of length n, self-attention creates an n by n score matrix in every layer and head. That makes prompt prefill expensive for long contexts because every token can compare with every other token. Memory for activations and intermediate attention data also grows quickly during training or large-batch inference.

The KV cache reduces generation cost after prefill by saving old keys and values instead of recomputing them for every new token. Grouped-query attention reduces serving memory by letting multiple query heads share fewer key-value heads. These optimizations reduce decoding and memory pressure, but they do not make full prompt prefill free. The context-windows-and-kv-cache lesson goes deeper on this tradeoff.

Learned absolute position embeddings store a vector per position, which is simple but tied to the trained context range. Sinusoidal embeddings use fixed waves, which avoid a learned table and expose relative patterns indirectly, but they are less common in frontier decoder-only LLMs today.

RoPE applies a position-dependent rotation to query and key dimensions before the dot product. Because the same rotation structure is used across positions, the resulting attention score naturally depends on relative distance. This is one reason RoPE is popular in modern decoder-only models and why long-context extensions often focus on RoPE scaling.

Attention weights are not a complete explanation of model behavior, but they are a useful window into some learned circuits. Researchers have found heads that track syntax, copy previous tokens, link pronouns to names, attend from a closing delimiter to its opener, or focus on induction patterns that repeat text structure.

A careful interview answer should avoid saying attention alone equals explanation. The attention output is later transformed by feed-forward layers and residual streams. Still, attention heads provide concrete evidence that transformers learn reusable relationship detectors rather than only memorizing surface statistics.

The n by n attention score matrix makes long prompts expensive, especially during prefill. Production systems control this with context limits, prompt compression, retrieval, batching policy, paged KV caches, and model variants designed for efficient attention.

Masks must be applied before softmax and must use values that behave correctly in the chosen precision. Implementations also subtract the row maximum before exponentiation so softmax stays stable in FP16, BF16, and quantized serving paths.

The KV cache grows with layers, sequence length, batch size, and the number of key-value heads. Multi-query and grouped-query attention reduce cache size by sharing keys and values across query heads, improving throughput with a small architecture tradeoff that must match the trained model configuration.

This small Python example computes one attention layer without model libraries. It shows the exact order: dot products, scaling, softmax, then weighted sum of values.

The causal mask changes the score matrix before normalization. Future positions become impossible choices, so their softmax weights are effectively zero.