Attention Mechanism
Attention Mechanism
The attention mechanism allows neural networks to dynamically focus on relevant parts of the input. It is the core innovation behind the Transformer and all modern LLMs.
Scaled Dot-Product Attention
The fundamental attention operation:
Attention(Q, K, V) = softmax(QK^T / √d_k) V
Where:
- Q (Query) — What we're looking for
- K (Key) — What we can attend to
- V (Value) — The information content
- √d_k — Scaling factor to prevent softmax saturation
Intuition
Each token produces a query vector. The dot product with all key vectors determines attention weights. These weights are normalized via softmax, then used to compute a weighted sum of value vectors.
Self-Attention
In self-attention, Q, K, and V all come from the same sequence. Each token attends to every other token (including itself), allowing the model to build contextual representations.
Causal (Masked) Attention
In decoder-only models, each token can only attend to itself and preceding tokens. A causal mask sets future token attention weights to -∞ (zero after softmax).
Multi-Head Attention
Instead of one attention computation, the model runs multiple attention heads in parallel, each learning different relationship types (syntax, semantics, positional, etc.).
Modern Optimizations
- Grouped Query Attention (GQA) — Multiple query heads share fewer key/value heads
- Multi-Query Attention (MQA) — All query heads share one key/value head
- Flash Attention — IO-aware, memory-efficient exact attention
- Sliding Window Attention — Local attention within a fixed window
- Sparse Attention — Attend to a subset of positions
KV Cache
During autoregressive generation, key and value tensors from previous steps are cached to avoid recomputation, dramatically speeding up inference.
Related
- Transformer Architecture — Where attention lives
- multi-head-attention — Parallel attention heads
- positional-encoding — How position info is added
- Context Window — Limits on attention span