Scaled Dot-Product✓ Mathematical
1 min read
◆ The PatternThe mechanism that lets every token look at every other token
Self-attention is the core operation of transformers. Each token produces a Query (what am I looking for?), a Key (what do I contain?), and a Value (what do I output?). Attention weights are the softmax of query-key dot products.
Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V
Scale by √d_k to prevent softmax saturation as dimension grows. Output is a weighted mix of value vectors.
Q = XW_Q K = XW_K V = XW_V where W ∈ ℝ^(d_model × d_k)
Linear projections from the input. No bias in most modern architectures.
Complexity is O(n² · d) — quadratic in sequence length. This is the fundamental bottleneck that drives context window research. Each attention weight tells you how much token i "pays attention to" token j.
The √d_k scaling is crucial. Without it, dot products grow proportionally with dimension, pushing softmax into regions with tiny gradients.
Interactive — attention weight heatmap
Python — self-attention from scratch#
import torch
import torch.nn.functional as F
def self_attention(x, W_q, W_k, W_v):
"""x: (batch, seq_len, d_model)"""
Q = x @ W_q # (batch, seq_len, d_k)
K = x @ W_k
V = x @ W_v
d_k = Q.size(-1)
scores = Q @ K.transpose(-2, -1) / d_k**0.5
weights = F.softmax(scores, dim=-1)
return weights @ V # (batch, seq_len, d_v)Pattern bridge: Every token attending to every other token — a complete correlation matrix computed at each layer. In markets, VWAP weights each price by volume — attention over the session.
Performance in practice
- O(n²) is the bottleneck: For a 128K context window, the attention matrix has 16 billion entries. Flash Attention reduces memory from O(n²) to O(n) by tiling and fusing CUDA kernels
- Flash Attention 2 gives 2-4x speedup over standard attention on A100 GPUs. It's now default in PyTorch 2.0+ via
torch.nn.functional.scaled_dot_product_attention - GPT-4 likely uses mixture-of-experts with grouped-query attention (GQA) to handle 128K context at reasonable cost
- For inference: KV-cache means attention cost is O(n) per new token, not O(n²). Cache size = 2 × n_layers × n_heads × d_head × seq_len × dtype_bytes
When to use this
✓ Use when: Building or understanding any transformer model. Debugging attention patterns to understand model behavior. Designing custom architectures. Understanding why context length is limited.
✗ Skip when: Using LLMs via API — attention is handled for you. Working with very long sequences (>100K tokens) where linear attention variants (Mamba, RWKV) may be more efficient.