Attention mechanism
At a glance
Definition
The attention mechanism is a neural network component that computes a weighted sum of values, where the weights come from comparing a 'query' against 'keys.' In plain terms: given a set of things to look at, attention learns how much each of them matters right now, then produces a weighted combination. The weights are learned from data — they aren't fixed rules.
Attention was originally invented to solve a specific problem in neural machine translation: how does a decoder producing French text decide which English words to focus on for the word it's currently generating? Bahdanau, Cho, and Bengio's 2014 paper introduced a learnable attention mechanism that let the decoder attend to different encoder positions at each output step. Translation quality jumped, and researchers quickly realized attention was useful far beyond translation.
The 2017 Attention Is All You Need paper generalized attention into self-attention — where a sequence attends to itself. Every token in a sentence looks at every other token to build a rich contextual representation. This turned out to be an extraordinarily general primitive: modern LLMs stack dozens of self-attention layers, and the pattern powers vision transformers, audio models, and multimodal systems.
There are several variants: self-attention (a sequence attends to itself), cross-attention (one sequence attends to another — used in encoder-decoder and multimodal models), multi-head attention (multiple attention operations in parallel), and efficient variants like grouped-query attention and FlashAttention. The underlying computation is the same in each: a learnable weighted sum.
Real-world example
What attention "attends to" — a concrete example
Consider the sentence 'The animal didn't cross the street because it was too tired.' When processing the word it, a well-trained attention head learns to focus heavily on animal (the referent) and less on other words. In the sentence 'The animal didn't cross the street because it was too wide,' the same attention head learns to focus on street instead — because that's what it refers to now. This dynamic, context-dependent focusing is what makes attention such a powerful primitive.
# Scaled dot-product attention (the core computation)
import torch
import torch.nn.functional as F
def attention(query, key, value):
"""
query, key, value: (batch, seq_len, dim)
Returns: attended output and attention weights
"""
d_k = query.size(-1)
# Compute attention scores
scores = torch.matmul(query, key.transpose(-2, -1))
scores = scores / (d_k ** 0.5) # scaling for stability
# Softmax gives attention weights that sum to 1
weights = F.softmax(scores, dim=-1)
# Weighted sum of values
output = torch.matmul(weights, value)
return output, weights
# Example: attend over a sequence of 10 tokens with dim=64
q = torch.randn(1, 10, 64)
k = torch.randn(1, 10, 64)
v = torch.randn(1, 10, 64)
out, attn_weights = attention(q, k, v)
# attn_weights[0, i, j] = how much position i attends to position j√d_k prevents softmax saturation as dimension grows. Every modern LLM uses this or a minor variant.When you'll encounter this
- You're using any LLM — every response is generated by repeatedly running attention over your conversation.
- You're debugging model behavior — attention visualization tools can show what the model is 'looking at' when it makes mistakes.
- You're designing prompts for retrieval — attention's tendency to focus on beginnings and ends explains the 'lost in the middle' problem.
- You're fine-tuning models — LoRA typically targets attention weight matrices because they encode task-relevant behavior.
- You're optimizing inference cost — the quadratic scaling of attention drives every context-window pricing model.
How it works
Compute queries, keys, values
Each input vector is multiplied by three learned matrices (W_q, W_k, W_v) to produce a query, key, and value vector. In self-attention, all three come from the same input sequence.
Compare queries with keys
Every query is compared with every key via dot product. High similarity → high score. This gives an attention score matrix of shape (sequence × sequence).
Scale and softmax
Scores are divided by √d_k (dimension) for numerical stability, then softmax is applied per row. Now each row is a probability distribution over positions to attend to.
Weighted sum of values
Each output position takes a weighted sum of value vectors, weighted by the attention distribution. Positions with higher attention contribute more.
Multi-head parallelism
The above happens H times in parallel with different learned Q/K/V matrices — different 'heads' can specialize in different patterns (syntax, semantics, coreference, etc.). Outputs are concatenated and projected back.
Causal masking (in LLMs)
For decoder-only LLMs, future positions are masked out before softmax so a token can only attend to itself and earlier positions. This is what makes autoregressive generation possible.
Efficient variants
Modern implementations use optimizations like FlashAttention (memory-efficient), GQA (fewer key/value heads), and sliding window for long context.
Common misconceptions
The name is metaphorical. Neural attention is a specific mathematical operation — a weighted sum with learned weights. It doesn't consciously focus on anything the way humans do. But the pattern of what gets weighted highly often mirrors what a human reader might consider relevant, which is why the metaphor stuck.
Attention predates Transformers — it was introduced by Bahdanau et al. in 2014 for adding attention to RNNs. The Transformer's contribution was showing that attention alone, without recurrence, was enough. Attention still shows up in non-Transformer models today, though pure Transformers dominate.
Empirically, more heads help up to a point, then plateau or hurt. Modern LLMs typically use 8-64 heads. Newer architectures like grouped-query attention reduce the number of independent heads to save memory during inference without much quality loss.
Related terms
Where you'll see this in practice
Learn more
- Read the deep-dive: How Transformers Work
- Compare Multi-head vs Grouped-Query Attention
- Browse the Deep Learning Architectures category
- Understanding Context Windows (and how attention shapes them)
Sources & further reading
- Neural Machine Translation by Jointly Learning to Align and Translate — Bahdanau, Cho, Bengio, 2014 (attention introduced). Read →
- Attention Is All You Need — Vaswani et al., 2017 (Transformer + self-attention). Read →
- FlashAttention — Tri Dao et al., 2022 (memory-efficient attention). Read →
- Grouped-Query Attention — Ainslie et al., 2023 (efficient inference-time attention). Read →
Frequently asked about Attention Mechanism
Self-attention is a specific type of attention where the query, key, and value all come from the same sequence — a sequence attending to itself. Regular attention could be a decoder attending to an encoder (cross-attention). Modern decoder-only LLMs use self-attention exclusively.
Because every position must be compared with every other position. For N tokens, that's N×N pairwise comparisons. Doubling context length quadruples attention cost. This is the fundamental scaling constraint on context windows.
Yes. Attention weights are just numbers between 0 and 1, so you can visualize them as heatmaps showing which positions attend to which. Tools like BertViz and interpretability research at Anthropic and elsewhere routinely visualize attention. This is one reason attention-based models are considered slightly more interpretable than pure MLPs.
Roughly, but with variations. Most use multi-head attention with learned Q/K/V projections. Newer models use grouped-query attention (fewer K/V heads to save memory), FlashAttention for efficient computation, and often sliding window attention for very long context. The core idea is identical.
State Space Models like Mamba use different mathematics to process sequences with linear rather than quadratic complexity. RWKV is an RNN-transformer hybrid. Linear attention variants approximate attention with linear complexity but with quality trade-offs. See our modern architectures category.