The Attention Mechanism: How LLMs Understand Context

A clear explanation of self-attention — the mathematical operation at the heart of every transformer that allows language models to understand relationships between words.

Key takeaways
  • Attention allows every word to directly access every other word, solving problems like pronoun resolution that sequential processing struggled with.
  • Attention uses a library analogy where queries search against keys to find relevant information stored in values, computed as softmax(QK^T / √d_k)V.
  • Decoder models use masked attention where tokens can only attend to previous tokens, enabling autoregressive generation without cheating.
  • Early attention heads learn syntactic patterns while later heads capture semantic relationships, with some heads perfectly tracking pronoun-antecedent relationships.
  • FlashAttention, Multi-Query Attention, and KV caching address the O(n²) computational cost of standard attention for long sequences.

The Problem Attention Solves

Consider the sentence: 'The trophy didn't fit in the suitcase because it was too big.' What does 'it' refer to — the trophy or the suitcase? Humans resolve this instantly using world knowledge (trophies are rigid, suitcases flex) and syntactic understanding. Pre-transformerTransformerThe neural network architecture that underpins virtually all modern LLMs, introduced in 2017, built around self-attention mechanisms that process entire sequences in parallel.Learn more → AI systems struggled with exactly this kind of long-range dependency.

Attention solves this by giving every word direct access to every other word. Instead of processing text sequentially and hoping relevant context survives through many steps, attention computes relationships between all pairs of words simultaneously. The model can directly compare 'it' against both 'trophy' and 'suitcase' to determine which makes more sense in context.

Query, Key, and Value — Intuition

Attention is often described using a library analogy. Your Query is a search request ('What do I need to understand my role in this sentence?'). Each word in the sequence has a Key (a searchable description: 'I am the subject of the main clause'). Attention compares your Query against all Keys to find the most relevant words. The Values are the actual information those words contain.

Mathematically, Attention(Q, K, V) = softmaxSoftmaxA mathematical function that converts a vector of raw logits into a probability distribution, with temperature controlling the sharpness of the output.Learn more →(QK^T / √d_k)V. The division by √d_k (the square root of the key dimension) prevents dot products from becoming too large and pushing softmax outputs toward zero gradient. Softmax ensures the weights sum to 1, creating a proper probability distribution over sequence positions.

Masked Attention in Decoder Models

In decoderDecoderA transformer component that generates text tokens sequentially using causal attention, where each token can only attend to previous tokens in the sequence.Learn more →-only models (GPTGPTGenerative Pre-trained Transformer — the model architecture and family name behind OpenAI's most famous models, from GPT-2 to GPT-5.Learn more →, Llama, Claude), attention is causal or masked: each tokenTokenThe basic unit of text that an LLM processes — roughly corresponding to a word or word-piece. Models read input and produce output in tokens, which is also how API usage is measured and billed.Learn more → can only attend to tokens before it in the sequence. This is necessary for autoregressiveAutoregressiveA text generation approach where models produce sequences one token at a time, with each new token conditioned on all previously generated tokens.Learn more → generation — during training, the model learns to predict each next token without 'looking at' future tokens, which would be cheating.

During inferenceInferenceThe process of running a trained AI model to generate outputs — what happens when you send a prompt and receive a response.Learn more → (generation), the causal mask means the model generates one token at a time, attending only to the promptPromptThe input text sent to a language model — the question, instruction, or context that triggers a response.Learn more → and previously generated tokens. The KV cacheKV CacheA memory optimization that stores previously computed attention keys and values so the model doesn't recompute them when generating each new token.Learn more → stores the key and value tensors from previous tokens so they don't need to be recomputed for each new token — a critical efficiency optimization for long generations.

What Attention Learns to Do

InterpretabilityInterpretabilityResearch into understanding how AI models work internally, including mechanistic interpretability that maps specific behaviors to neural components.Learn more → researchers have studied what individual attention heads learn. Heads in early layers often learn syntactic patterns: subject-verb agreement, determiner-noun agreement, parenthetical structure. Heads in later layers tend to capture more semantic relationships: coreference resolution, semantic similarity, factual associations.

Some heads are strikingly interpretable. One famous example from GPT-2 analysis: a single attention head almost perfectly tracks which pronoun refers to which antecedent in complex sentences. This emergent specialization — not explicitly trained, just discovered through gradient descentGradient DescentAn optimization algorithm that iteratively adjusts model parameters by moving in the direction opposite to the gradient of the loss function.Learn more → on next-token prediction — hints at the surprising structure that scale unlocks.

Making Attention Efficient

Standard attention is O(n²) in sequence length: doubling the context quadruples memory and compute. This is the key bottleneck for long-context models. FlashAttentionFlashAttentionAn IO-aware attention algorithm that tiles computation in fast SRAM memory to reduce data movement, enabling efficient training and inference on long sequences.Learn more → rewrites the attention computation to be memory-efficient by computing attention in tiles that fit in GPU SRAM, dramatically reducing memory bandwidth requirements without changing the mathematical result.

Multi-Query Attention (MQA) and Grouped-Query AttentionGrouped-Query AttentionAn attention optimization technique where multiple attention heads share key and value projections in groups, reducing memory usage and speeding up inference while maintaining model quality.Learn more → (GQA) reduce the KV cache size by sharing key and value projections across multiple query heads. This is why models like Llama 3 can serve long-context requests efficiently at scale. Sparse attention, sliding window attention (Mistral), and other approximations trade some theoretical expressiveness for further efficiency gains.

Read next