The Attention Mechanism, Decoded

A simple drawing to decode the attention mechanism as it actually works in modern LLMs (explanation very simplified, I admit):

Input Embeddings Each token becomes a vector representing its meaning only. Position comes later.

3 linear projections (W_Q, W_K, W_V) The same input is projected into 3 different representations:

→ Q (Query), which asks "What am I looking for?" → K (Key), which describes "What do I contain?" → V (Value), which carries "Here's my actual information"

+ RoPE on Q and K Position isn't baked into the input anymore. With Rotary Position Embeddings (used in LLaMA, Mistral, Qwen…), position is injected directly into Q and K through rotation. This is what lets models generalize to longer sequences than they were trained on.

Q · Kᵀ (dot product) Every query is compared against every key. And a key is just another token's description of itself. So this step answers: "how relevant is token B to token A?" for every pair in the sequence.

÷ √dk (scale) Without this, dot products get too large and softmax saturates. This keeps gradients stable.

Causal Mask Before softmax, a mask zeros out all future positions. Each token can only attend to itself and past tokens, never what comes next. This is what makes autoregressive generation possible, and what distinguishes decoder-only LLMs from encoders like BERT.

Softmax Turns the masked scores into probabilities.

Attention × V Multiply those probabilities by the values. Each token becomes a weighted mix of all previous tokens' information.

Grouped Query Attention (GQA) Modern LLMs don't use classic multi-head attention anymore. Instead of giving each head its own K and V, multiple Q heads share the same KV heads. LLaMA 3 for instance uses 32 Q heads but only 8 KV groups, cutting KV cache by 4x with minimal quality loss.

Concat + W_O All heads are concatenated and projected back into a single representation.

Output Ready for the next layer.