Part 1: Query-Key Matching via Scaled Dot Products
The attention mechanism begins with computing pairwise similarities between queries and keys via the matrix product $Q K^\top \in \mathbb{R}^{n_q \times n_k}$. Each entry measures the dot product $q_i^\top k_j$—high values indicate strong alignment between query $i$ and key $j$. This is content-based addressing: the model dynamically decides which keys are relevant based on input content, unlike fixed indexing. The scaling by $1/\sqrt{d_k}$ prevents the softmax function from saturating when embedding dimensions grow large. Without scaling, dot products can become arbitrarily large (as dimension increases, random vectors have dot products with variance $\sim d_k$), pushing softmax into regions where gradients vanish. The $\sqrt{d_k}$ normalization keeps dot products at unit variance, ensuring stable gradients during training. This is why you’ll see scores = (Q @ K.T) / np.sqrt(d_k) in every transformer implementation.

Part 2: Softmax Normalization Creates Convex Weights
Row-wise softmax converts raw scores into probability distributions: $A_{ij} = \exp(S_{ij}) / \sum_{j'} \exp(S_{ij'})$. Each row $A_{i\cdot}$ has non-negative entries summing to 1, making them valid convex weights. This is crucial: the final output is a convex combination of value vectors, ensuring outputs lie in the convex hull of $V$’s rows. Softmax is differentiable everywhere (exponential smoothness enables clean gradients), temperature-controllable (scaling scores by $1/T$ adjusts sharpness), and invariant to adding constants to scores (subtracting row-max prevents overflow). The probabilistic interpretation is powerful: “query 1 attends 60% to key 0, 40% to key 2”—directly inspectable, unlike hidden activations in MLPs. However, recent work questions whether high attention weights truly indicate “importance” (attention might satisfy constraints without being causal), spawning research into attention flow, integrated gradients, and counterfactual probing.

Part 3: Value Aggregation via Weighted Sum
The final output $O = A V$ performs a weighted sum: each output row $O_{i\cdot} = \sum_j A_{ij} V_{j\cdot}$. High-attention keys contribute more to the output. This is differentiable memory access—the model “reads” from a key-value store, with read weights determined by query-key similarity. Unlike hard attention (sampling one key) or linear layers (fixed weights), soft attention blends all values smoothly, enabling gradient flow to all parameters. In self-attention ($Q=K=V$ projections of the same input), this creates context-aware representations: each token’s output depends on all other tokens it attends to. For cross-attention (e.g., encoder-decoder), queries come from the decoder, keys/values from the encoder, enabling the decoder to “look at” encoder outputs. The value matrix $V$ doesn’t need the same dimension as $Q,K$—often $d_v < d_k$ for efficiency, though standard transformers use $d_v = d_k = d_{\text{model}}/H$ (model dimension divided by number of heads).

Part 4: Matrix Products as Composition of Linear Maps
Attention composes two linear maps: (1) $Q,K,V$ are typically obtained by projecting input embeddings $X$ through learned matrices $W_Q, W_K, W_V$, and (2) the attention output is projected again via $W_O$. The forward pass is:
\[\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) W_O\]
where each head computes $\text{Attention}(XW_Q^h, XW_K^h, XW_V^h)$. This is a cascade of matrix products: $X \to XW_Q \to (XW_Q)(XW_K)^\top \to \sigma(\cdot) \to A(XW_V) \to \text{Concat} \to (\cdots)W_O$. Understanding this chain is essential for backpropagation: gradients flow backward through each product via the chain rule, decomposing into Jacobian-vector products. For example, $\nabla_{W_Q} \mathcal{L} = X^\top (\nabla_Q \mathcal{L})$ where $\nabla_Q = \nabla_A \cdot \partial A/\partial Q$ involves the softmax Jacobian. Efficient implementations fuse these operations to minimize memory transfers (Flash Attention achieves 2-4× speedups by reordering computations to fit in SRAM).

Why This Matters for ML
Attention is the computational engine of modern AI: transformers power GPT (autoregressive language models), BERT (bidirectional encoders), Vision Transformers (image classification), DALL-E/Stable Diffusion (text-to-image), AlphaFold2 (protein structure), Perceiver (general-purpose architectures), and recent audio/video models. The quadratic complexity $O(n^2 d)$ for self-attention on sequences of length $n$ is the main bottleneck—this has spawned entire research directions: sparse attention (Longformer, BigBird with predefined patterns), linear attention (Performer, RWKV approximating softmax with kernels), and Flash Attention (memory-optimized exact attention). Multi-head attention (running $H$ independent attention operations in parallel, then concatenating) allows models to attend to different representation subspaces simultaneously—empirically crucial for performance. Attention weights are often visualized to interpret model behavior (e.g., does a pronoun attend to its antecedent?), though this interpretability is debated. Beyond NLP, attention applies anywhere dynamic, content-based routing is needed: graph neural networks (nodes attend to neighbors), reinforcement learning (attend to relevant state components), and slot attention (object-centric representations).

Numerical and Implementation Notes
Use scipy.special.softmax(scores, axis=1) for numerically stable softmax (subtracts row-max before exponentiating to prevent overflow). For large-scale attention, consider: (1) mixed-precision (FP16 for forward pass, FP32 for gradients), (2) gradient checkpointing (recompute activations during backward pass to save memory), (3) fused kernels (xFormers, Flash Attention optimize memory access patterns), and (4) causal masking for autoregressive models (set $S_{ij} = -\infty$ for $j > i$ to prevent attending to future tokens). In multi-head attention, reshape matrices to include head dimension: Q.shape = (batch, num_heads, seq_len, d_k), then compute attention independently per head. The output projection $W_O$ mixes information across heads. Dropout is typically applied to attention weights (zeroing out some $A_{ij}$ entries) to prevent overfitting, though this slightly violates the “sum-to-one” property—implementations renormalize after dropout. For variable-length sequences, use padding masks (set attention weights to zero for padding tokens) to prevent attending to meaningless positions.

Shape Notes and Scaling Considerations
Standard dimensions: $Q,K \in \mathbb{R}^{n \times d_k}$, $V \in \mathbb{R}^{n \times d_v}$, scores $\in \mathbb{R}^{n \times n}$, output $\in \mathbb{R}^{n \times d_v}$. For multi-head attention with $H$ heads and model dimension $d_{\text{model}}$, typically $d_k = d_v = d_{\text{model}}/H$ (e.g., $d_{\text{model}}=512$, $H=8$, $d_k=64$). This keeps total parameter count constant while allowing diverse attention patterns. Batch dimensions prepend: (batch_size, seq_len, d_model). Memory usage is dominated by storing the attention matrix $A \in \mathbb{R}^{n \times n}$ per batch per head—for a 2048-token sequence with 32 heads and batch size 8, this requires $8 \times 32 \times 2048 \times 2048 \times 4 \text{ bytes} \approx 4.3 \text{GB}$ just for attention weights (in FP32), not including activations/gradients. Flash Attention avoids materializing the full $A$ matrix in HBM (high-bandwidth memory), instead computing attention in tiles that fit in SRAM (on-chip cache), reducing memory from $O(n^2)$ to $O(n)$ while maintaining exact computation. For extremely long sequences ($n > 10^4$), sparse or linear attention variants replace dense $O(n^2)$ with $O(n \log n)$ or $O(n)$ complexity by restricting which positions can attend to each other (local windows, fixed patterns, learned sparsity, or kernel approximations). - Interpretation: relate algebraic steps to geometry (subspaces, projections) and to ML behavior (generalization, stability).