Part 1: Query-Key Similarity via Scaled Dot Products

The first matrix product $QK^\top$ computes all pairwise similarities between queries and keys in one operation. For $Q \in \mathbb{R}^{2 \times 2}$ and $K \in \mathbb{R}^{3 \times 2}$, the product $QK^\top \in \mathbb{R}^{2 \times 3}$ has 6 entries (2 queries $\times$ 3 keys), each representing the dot product $q_i^\top k_j$. The transpose $K^\top$ is essential: $Q$ has shape $(2, 2)$, $K$ has shape $(3, 2)$, so $Q @ K$ would fail (inner dimensions $2 \ne 3$ mismatch), but $Q @ K.T$ succeeds ($(2, 2) \times (2, 3) \to (2, 3)$). Scaling by $1/\sqrt{d_k}$: For random unit vectors $q, k \in \mathbb{R}^{d_k}$, the dot product $q^\top k$ has variance $d_k$ (sum of $d_k$ independent products, each with variance 1). For $d_k = 512$ (GPT-3), unscaled scores would have standard deviation $\sim 23$, pushing softmax into saturation (e.g., $\exp(20) / (\exp(20) + \exp(0)) \approx 0.9999$, one key dominates, gradients vanish). Scaling by $1/\sqrt{d_k}$ normalizes variance to $1$, keeping scores in a moderate range ($\sim [-2, 2]$) where softmax remains sensitive to all inputs. Why not normalize by $d_k$? Dividing by $d_k$ (linear scaling) would make scores $\sim 1/d_k$ (tiny for large $d_k$), causing underflow in softmax. The $\sqrt{d_k}$ balances: large enough to prevent saturation, small enough to avoid underflow.

Part 2: Softmax Normalization to Attention Weights

Row-wise softmax converts raw scores to probability distributions: $A_{ij} = \exp(\text{scores}_{ij}) / \sum_k \exp(\text{scores}_{ik})$. Properties: (1) Non-negative: $\exp(\cdot) > 0$ ensures $A_{ij} \ge 0$. (2) Row-stochastic: $\sum_j A_{ij} = 1$ (each query distributes 100% attention across keys). (3) Differentiable: Softmax gradients $\partial A_{ij} / \partial \text{scores}_{ik} = A_{ij} (\delta_{jk} - A_{ik})$ enable backpropagation. Why softmax vs. other normalizations? (a) Max-normalization ($A_{ij} = \text{scores}_{ij} / \max_k \text{scores}_{ik}$) violates row-stochasticity (sums $\ne 1$). (b) L1-normalization ($A_{ij} = \text{scores}_{ij} / \sum_k \text{scores}_{ik}$) can produce negative weights if scores are negative. (c) Softmax (exponential + normalization) guarantees non-negativity and stochasticity, and has a maximum entropy interpretation (Boltzmann distribution). Temperature scaling: Softmax with temperature $\tau$: $A_{ij} = \exp(\text{scores}_{ij} / \tau) / \sum_k \exp(\text{scores}_{ik} / \tau)$. As $\tau \to 0$, attention becomes “hard” (argmax, one-hot). As $\tau \to \infty$, attention becomes uniform (all keys equally weighted). Standard attention uses $\tau = 1$.

Part 3: Value Aggregation via Weighted Sum

The second matrix product $AV$ computes attention-weighted aggregations: $O = AV \in \mathbb{R}^{2 \times 2}$. Each output row is $O_i = \sum_{j=1}^3 A_{ij} V_j$, a convex combination of value vectors (since $\sum_j A_{ij} = 1$, $A_{ij} \ge 0$). Geometric interpretation: Each query $i$ produces an output $O_i$ lying in the convex hull of $\{V_1, V_2, V_3\}$ (the set of all weighted averages). Why separate keys and values? In principle, we could set $V = K$ (attention weights determined by and applied to the same vectors). However, decoupling allows keys to encode what information is available (searchable metadata) while values encode what to retrieve (actual content). In practice, $Q, K, V$ are learned projections of the input: $Q = XW_Q$, $K = XW_K$, $V = XW_V$, where $W_Q, W_K, W_V$ are trained end-to-end. Self-attention vs. cross-attention: In self-attention (e.g., BERT, GPT), $Q, K, V$ all come from the same sequence $X$. In cross-attention (e.g., encoder-decoder translation), $Q$ comes from decoder states, $K, V$ from encoder outputs—allowing the decoder to “attend” to source sentence positions when generating target words.

Why This Matters for ML

Transformer architecture dominance: Transformers replaced RNNs/LSTMs as the default sequence model because attention enables full parallelization—all positions processed simultaneously ($O(n^2 d)$ total, but parallelizable across $n$ positions), whereas RNNs require sequential processing ($O(nd)$ per step, $n$ sequential steps, total $O(n^2 d)$ but sequential). This $100-1000\times$ wall-clock speedup on GPUs enabled scaling to billions of parameters (GPT-3: 175B, PaLM: 540B). Quadratic bottleneck: For sequence length $n$ and dimension $d$, attention costs $O(n^2 d)$. For $n = 10^4$ (long documents), this becomes prohibitive ($10^8$ pairwise scores per head). Efficient attention variants address this: sparse attention (Longformer: $O(n \log n)$ via local+global masks), low-rank attention (Linformer: $O(nd)$ via key/value projection to rank $k \ll n$), kernel attention (Performer: $O(nd)$ via random feature approximation of softmax kernel). Multi-head attention: Standard Transformers use $h = 8$-16 parallel heads, each with $d_k = d_{\text{model}} / h$. Heads learn complementary attention patterns (e.g., head 1: syntactic dependencies, head 2: semantic similarity, head 3: positional locality). Positional encodings: Raw attention is permutation-invariant (shuffling input tokens doesn’t change attention weights). Transformers add sinusoidal or learned positional embeddings to inject sequence order information: $X_{\text{pos}} = X + \text{PE}(pos)$, where $\text{PE}(pos)$ encodes absolute or relative positions.

ML Examples and Patterns

Self-attention layer (single head): Q = X @ W_Q, K = X @ W_K, V = X @ W_V project input $X \in \mathbb{R}^{n \times d_{\text{model}}}$ to queries/keys/values via learned matrices $W_Q, W_K, W_V \in \mathbb{R}^{d_{\text{model}} \times d_k}$. Then compute scores = Q @ K.T / sqrt(d_k), A = softmax(scores, axis=-1), O = A @ V. Output $O \in \mathbb{R}^{n \times d_k}$ is projected back to model dimension via $W_O \in \mathbb{R}^{d_k \times d_{\text{model}}}$.

Multi-head attention: Run $h$ attention heads in parallel with $d_k = d_{\text{model}} / h$ per head, then concatenate: O_multi = Concat(head_1, ..., head_h) @ W_O, where head_i = Attention(X @ W_Q^i, X @ W_K^i, X @ W_V^i). PyTorch implementation: nn.MultiheadAttention(embed_dim=512, num_heads=8).

Masked attention (autoregressive): GPT-style models mask future positions to prevent lookahead: mask = torch.triu(torch.ones(n, n), diagonal=1).bool(), then scores.masked_fill_(mask, -1e9) before softmax. This zeroes attention to future tokens, enforcing left-to-right generation.

Cross-attention (encoder-decoder): In translation, decoder queries attend to encoder keys/values: Q_dec = X_dec @ W_Q, K_enc = X_enc @ W_K, V_enc = X_enc @ W_V, then O = Attention(Q_dec, K_enc, V_enc). Decoder state at position $i$ can attend to all encoder positions (no causal mask).

Positional encoding: PE[pos, 2i] = sin(pos / 10000^(2i/d)), PE[pos, 2i+1] = cos(pos / 10000^(2i/d)) for $i=0, \ldots, d/2-1$. Add to embeddings: X_input = Embedding(tokens) + PE. Alternative: learned positional embeddings nn.Embedding(max_len, d_model).

Attention dropout: A = F.dropout(softmax(scores), p=0.1, training=True) randomly zeros attention weights during training to prevent overfitting to specific dependencies.

Connection to Linear Algebra Theory

Attention as convex combination: Each output $O_i = \sum_j A_{ij} V_j$ is a point in the convex hull of $\{V_1, \ldots, V_{n_k}\}$ (all weighted averages with non-negative weights summing to 1). This is a projection onto the value subspace, but weights are data-dependent (learned from query-key similarity) rather than fixed (like orthogonal projection in least squares).

Dot product as similarity metric: For normalized vectors $\|q\|_2 = \|k\|_2 = 1$, the dot product $q^\top k = \cos(\theta)$ is the cosine of the angle between $q$ and $k$. High dot product ($\cos(\theta) \approx 1$) means $\theta \approx 0$ (aligned vectors, high similarity). Attention amplifies this via softmax: similar pairs get disproportionately large weights.

Softmax as exponential family: Softmax $A_{ij} = \exp(s_{ij}) / Z_i$ (where $Z_i = \sum_k \exp(s_{ik})$) is the maximum entropy distribution subject to expected score constraint $\mathbb{E}[s] = \sum_j A_{ij} s_{ij}$. This is a Boltzmann distribution in statistical mechanics, with “temperature” $\tau = 1$ (or $\tau = \sqrt{d_k}$ if unscaled).

Row-stochastic matrices as Markov transitions: Attention matrix $A$ (rows sum to 1) is a transition matrix for a Markov chain on keys. Applying $A$ repeatedly ($A^2, A^3, \ldots$) corresponds to multi-hop random walks on the key-value graph. Some architectures (e.g., Graph Attention Networks) explicitly interpret attention as graph traversal.

Rank-1 outer products: Each attention contribution $A_{ij} V_j$ is a rank-1 matrix (outer product of scalar weight $A_{ij}$ with vector $V_j$). Attention sums $n_k$ such rank-1 contributions: $O_i = \sum_j A_{ij} V_j$, analogous to SVD reconstruction $X \approx \sum_{k=1}^r \sigma_k u_k v_k^\top$.

Kernel methods interpretation: Attention can be written as $O_i = \sum_j \kappa(q_i, k_j) V_j / \sum_j \kappa(q_i, k_j)$, where $\kappa(q, k) = \exp(q^\top k / \sqrt{d_k})$ is an unnormalized Gaussian kernel. This enables kernel approximation techniques (random Fourier features, Nyström) for efficient attention (Performer, Linear Transformers).

Numerical and Implementation Notes

Shape discipline: $Q \in \mathbb{R}^{n_q \times d_k}$ (2 queries, 2-dim), $K \in \mathbb{R}^{n_k \times d_k}$ (3 keys, 2-dim), $V \in \mathbb{R}^{n_k \times d_v}$ (3 values, 2-dim). Scores: $(n_q, n_k) = (2, 3)$. Attention weights: $(n_q, n_k) = (2, 3)$. Output: $(n_q, d_v) = (2, 2)$. Check dimensions before every matrix multiply: $Q @ K.T$ requires $Q.shape[1] == K.shape[1]$ ($d_k$ match), $A @ V$ requires $A.shape[1] == V.shape[0]$ ($n_k$ match).

Gotcha 1: Softmax axis. softmax(scores, axis=1) normalizes across keys (axis=1, columns) for each query (axis=0, rows). Using axis=0 would normalize across queries, which is meaningless (each key would distribute attention to queries, backwards). Always normalize across the key dimension.

Gotcha 2: Scaling factor. Use $1/\sqrt{d_k}$, not $1/\sqrt{d_v}$ or $1/\sqrt{n_k}$. The goal is to counteract dot product variance growth with key dimension ($\text{Var}(q^\top k) = d_k$ for unit vectors).

Gotcha 3: Masked attention. For autoregressive models (GPT), mask future positions: mask = np.triu(np.ones((n_q, n_k)), k=1) (upper triangular), then scores = scores - 1e9 * mask (set future scores to $-\infty$, softmax $\to 0$). Don’t use scores[mask] = -inf (in-place modification breaks autograd).

Gotcha 4: Batch dimensions. Production attention has shape (batch, num_heads, seq_len, seq_len). Ensure softmax normalizes over the last dimension (keys): F.softmax(scores, dim=-1). Broadcasting rules: (B, H, N_q, N_k) @ (B, H, N_k, D_v) -> (B, H, N_q, D_v).

Gotcha 5: Numerical stability. Softmax can overflow ($\exp(1000) \to \infty$) or underflow ($\exp(-1000) \to 0$). Use log-sum-exp trick: $\log \sum_j \exp(x_j) = \max_j x_j + \log \sum_j \exp(x_j - \max_j x_j)$. PyTorch F.softmax handles this automatically.

Gotcha 6: FlashAttention. For long sequences ($n > 2048$), materializing full $QK^\top \in \mathbb{R}^{n \times n}$ exceeds GPU memory. FlashAttention (Dao et al., 2022) fuses attention computation in CUDA kernels, tiling to avoid storing full attention matrix—$3\times$ memory reduction, $2\times$ speedup.

Numerical and Shape Notes

Verification checks: (1) Row sums: A.sum(axis=1) should equal [1., 1.] (each query’s weights sum to 1). (2) Non-negativity: A >= 0 all elements (softmax guarantees this). (3) Output shape: O.shape == (n_q, d_v) (2 queries, 2-dim values). (4) Numerical equivalence: Forward pass should match backprop gradients (use torch.autograd.gradcheck for validation).

Cost analysis: For $n_q = n_k = n = 512$ (typical sentence), $d_k = d_v = 64$ (per-head dimension): - $QK^\top$: $n^2 d_k = 512^2 \times 64 \approx 17M$ FLOPs - Softmax: $n^2 = 512^2 \approx 262K$ FLOPs (negligible) - $AV$: $n^2 d_v = 512^2 \times 64 \approx 17M$ FLOPs - Total: $\sim 34M$ FLOPs per head. For 8 heads: $\sim 270M$ FLOPs per layer. GPT-3 (96 layers, 16k context): $\sim 4 \times 10^{12}$ FLOPs per forward pass.

Memory: Attention matrix $A$ requires $n^2$ storage. For $n = 16384$ (Longformer), $A$ is $16384^2 \approx 268M$ elements ($\sim 1$ GB for float32). This is the memory bottleneck of Transformers—efficient variants (sparse, low-rank) reduce this to $O(n)$ or $O(n \log n)$.

Tolerance: Softmax is numerically stable (log-sum-exp trick), but attention weights can be tiny ($< 10^{-6}$) for dissimilar query-key pairs. Check for underflow: A.min() should be $> 0$ (if $\approx 0$, may indicate numerical issues or extreme sparsity).

Pedagogical Significance

This example is the foundational demonstration of attention mechanisms in modern deep learning:

Key takeaways: (1) Attention = weighted aggregation: Dynamically select and combine information based on query-key similarity. (2) Query-key-value paradigm: Separates “what to search for” (query) from “what’s available” (key) and “what to retrieve” (value). (3) Softmax normalization: Converts raw scores to probabilities (row-stochastic, interpretable weights). (4) Scaling prevents saturation: $1/\sqrt{d_k}$ keeps scores in moderate range for stable gradients. (5) Permutation invariance: Raw attention has no notion of order; positional encodings inject sequence structure.

Common misconceptions addressed: (a) “Attention requires recurrence”—No, Transformers are fully feedforward (parallel), unlike RNNs. (b) “Attention is approximate”—No, it’s an exact weighted sum, fully differentiable. (c) “Softmax axis doesn’t matter”—Critical—must normalize across keys (each query distributes attention). (d) “Attention always attends uniformly”—No, exponential weighting (softmax) creates sharp distributions for high-similarity pairs. (e) “Transformers don’t need positional info”—False—raw attention is permutation-invariant, positional encodings are essential.

Connection to other examples: Example 70 (attention as data-dependent projection), Example 72 (query-key similarity via inner products), Example 79 (sparse attention masks), Example 74 (multi-head as low-rank approximation). This example provides the implementation foundation for understanding BERT, GPT, Vision Transformers, and diffusion models—all built on variants of this core mechanism.

Why pedagogically powerful: It isolates the core attention operation (scaled dot-product attention) in minimal code (~10 lines), showing how three matrices ($Q, K, V$) and two operations (matmul + softmax) produce dynamic context aggregation. The small dimensions (2 queries, 3 keys) make manual calculation feasible, demystifying the “black box” of Transformers. Students see that attention is fundamentally linear algebra (matrix products + softmax), not magic. This is the gateway to understanding modern NLP, computer vision, and generative AI—attention is the computational primitive underlying language models (GPT-4, Gemini), vision models (ViT, CLIP), and multimodal systems (DALL-E, Stable Diffusion).