The attention mechanism is a computational primitive that lets a model weigh the relevance of each input token to every other token when building representations. Instead of processing sequences in fixed order like RNNs or CNNs, attention computes pairwise relationships in parallel, enabling the model to route information dynamically based on content. This single idea — learned, content-based routing — is what makes modern LLMs possible.
How attention works
At its core, attention takes three matrices derived from the same input: queries (Q), keys (K), and values (V). For each position, the model asks: “How relevant is every other position to me?” It answers by taking the dot product of its query vector against all key vectors, scaling by the square root of the key dimension, applying softmax to get a probability distribution, then using those weights to combine the value vectors.
import torch
import torch.nn.functional as F
def scaled_dot_product_attention(Q, K, V, mask=None):
"""
Q: (batch, heads, seq_len, d_k)
K: (batch, heads, seq_len, d_k)
V: (batch, heads, seq_len, d_v)
"""
d_k = Q.size(-1)
scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attn_weights, V)
return output, attn_weightsThe scaling factor 1/sqrt(d_k) prevents gradients from vanishing when the dot products grow large — a practical detail that matters in training. The mask parameter handles causal masking (preventing future tokens from influencing past ones) and padding masks.
Multi-head attention
Single attention heads can only capture one type of relationship. Multi-head attention runs the above computation in parallel with different learned projections, then concatenates the results and projects back to the model dimension.
class MultiHeadAttention(torch.nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
assert d_model % n_heads == 0
self.d_k = d_model // n_heads
self.n_heads = n_heads
self.w_q = torch.nn.Linear(d_model, d_model)
self.w_k = torch.nn.Linear(d_model, d_model)
self.w_v = torch.nn.Linear(d_model, d_model)
self.w_o = torch.nn.Linear(d_model, d_model)
def forward(self, x, mask=None):
batch, seq_len, _ = x.shape
Q = self.w_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
out, _ = scaled_dot_product_attention(Q, K, V, mask)
out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
return self.w_o(out)Each head learns different patterns — some track syntactic dependencies, others resolve coreference, others handle positional reasoning. The concatenation and output projection let the model mix these signals.
Why attention matters for LLMs
Before attention, sequence modeling meant recurrence or convolution. RNNs process tokens sequentially, creating a bottleneck: information from token 1 must pass through tokens 2 through N to reach token N+1. This makes long-range dependencies hard to learn and prevents parallelization.
Attention removes both constraints. Every token attends to every other token in a single layer. The path length between any two positions is constant — one attention layer. This enables:
Parallel training: The entire sequence processes at once. GPUs stay saturated.
Direct gradient flow: Gradients from the loss at position N flow directly to position 1 without passing through intermediate timesteps. This mitigates vanishing gradients.
Content-based routing: The model decides dynamically what to attend to. A verb can attend to its subject regardless of distance. A pronoun can attend to its antecedent. This is fundamentally different from fixed receptive fields in CNNs or the implicit routing in RNNs.
Compositionality: Stacking attention layers lets the model build increasingly abstract representations. Layer 1 might attend to adjacent tokens for morphology. Layer 6 might attend across clauses for discourse structure. Layer 12 might attend across documents for reasoning.
A concrete example: coreference resolution
Consider: “The lawyer asked the doctor if she could help with the case.”
The pronoun “she” could refer to either “lawyer” or “doctor.” A human resolves this by world knowledge — lawyers ask doctors for help with cases, not the reverse. An attention-based model learns this statistically.
At the layer where “she” is processed, its query vector will have high dot product with the key vector of “lawyer” if the model has learned that pattern. The attention weight for “lawyer” might be 0.7, “doctor” 0.2, other tokens 0.1. The output representation for “she” becomes a weighted blend: 0.7 * lawyer_representation + 0.2 * doctor_representation + …
This is not a lookup table. The same “she” in “The doctor asked the lawyer if she could help” produces different weights because the context — the query and key projections of surrounding tokens — changes.
You can visualize this with attention maps:
def visualize_attention(model, tokenizer, text, layer=11, head=0):
"""Extract attention weights for a specific layer and head."""
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
# outputs.attentions is tuple of (batch, heads, seq, seq) per layer
attn = outputs.attentions[layer][0, head].cpu().numpy()
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
return tokens, attnRunning this on the lawyer/doctor sentence shows the model attending from “she” to “lawyer” in the first version and to “doctor” in the second — same pronoun, different resolution, driven entirely by context.
Common misconceptions
“Attention is just a weighted average”
Technically true but misleading. The weights are learned functions of the input, not fixed. The query, key, and value projections are all parameterized and trained end-to-end. This means the model learns what to attend to and how to represent what it attends to. Calling it a weighted average obscures the learned routing behavior.
“Attention weights are interpretable”
Attention weights show where gradient flow concentrates, but they don’t necessarily correspond to human-interpretable reasoning. A head might attend to punctuation for syntactic structure, or to a seemingly unrelated token for a learned heuristic. Treating attention weights as explanations is a category error — they’re internal model states, not causal explanations.
“More heads always helps”
Diminishing returns set in quickly. The original Transformer used 8 heads for 512-dimensional models (64 per head). Modern models often use 32–96 heads for 4096–8192 dimensions. But each head adds parameters and compute. Some heads become redundant — they learn similar patterns. Pruning studies show you can remove 30–50% of heads post-training with minimal degradation.
“Attention is O(n²) so it doesn’t scale”
The quadratic complexity is real for standard attention. But this has spawned an entire subfield of efficient attention: sparse attention (Longformer, BigBird), linear attention (Performer, Linformer), flash attention (IO-aware kernel fusion), and sliding window attention. The O(n²) bound applies to the naive implementation, not the concept. Production systems at n4n.ai and elsewhere routinely handle 128k+ context windows using these techniques.
“Self-attention and cross-attention are fundamentally different”
They share the same mathematical form. Self-attention: Q, K, V all come from the same sequence. Cross-attention: Q comes from the decoder, K and V from the encoder. The mechanism is identical; only the source of the matrices differs. This unified interface is why decoder-only models (GPT) and encoder-decoder models (T5) can share the same attention implementation.
Attention variants you’ll encounter
Causal attention: Mask prevents attending to future tokens. Essential for autoregressive generation.
Bidirectional attention: No mask. Used in encoders (BERT) and prefix-LMs.
Sliding window attention: Each token attends only to a fixed window around it. Linear complexity. Used in Longformer, Mistral.
Grouped-query attention (GQA): Multiple query heads share key/value heads. Reduces KV cache memory during generation. Used in Llama 2 70B, Gemma.
Multi-query attention (MQA): Extreme case of GQA — all query heads share one key and one value head. Maximum KV cache savings, slight quality tradeoff.
Flash attention: Not a mathematical variant — a kernel fusion that computes attention without materializing the full (batch, heads, seq, seq) matrix in HBM. Massive speedup and memory reduction for long sequences.
Where to go deeper
The original paper “Attention Is All You Need” (Vaswani et al., 2017) remains the best starting point. For implementation details, read the annotated Transformer from Harvard NLP. For efficient variants, the Longformer and FlashAttention papers are practical. For interpretability debates, see “Attention is not Explanation” (Jain & Wallace, 2019) and the rebuttal “Attention is not not Explanation” (Wiegreffe & Pinter, 2019).
Understanding what is the attention mechanism means understanding learned routing. Everything else — multi-head, causal masking, efficient variants — is engineering on top of that core idea. If you grasp that queries route to keys to combine values, you can read any Transformer variant and know what it’s doing.