Scaled Dot-Product Attention

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. From Retrieval to Attention
5. The QKV Formulation
6. Why Scale by Root-dk
7. Attention as Soft Dictionary Lookup
8. Masking: Causal and Padding
9. Computational Complexity
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive scaled dot-product attention from first principles.
2. Prove why the $1/\sqrt{d_k}$ scaling is necessary for stable softmax.
3. Interpret attention weights as a soft retrieval mechanism.
4. Derive the quadratic complexity $O(T^2d)$ and identify the bottleneck.
5. Implement causal masking mathematically.

Notation

• $\mathbf{Q} \in \mathbb{R}^{T_q \times d_k}$ — query matrix
• $\mathbf{K} \in \mathbb{R}^{T_k \times d_k}$ — key matrix
• $\mathbf{V} \in \mathbb{R}^{T_k \times d_v}$ — value matrix
• $\mathbf{A} \in \mathbb{R}^{T_q \times T_k}$ — attention weight matrix
• $d_k$ — key/query dimension
• $T$ — sequence length

Core Intuition

Attention answers: "For each position in the output, which input positions are most relevant?" It computes a weighted average of values, where the weights are determined by the similarity between queries and keys. Unlike fixed-weight layers (convolution, MLP), attention weights are data-dependent — they change based on the input content.

From Retrieval to Attention

Hard retrieval: Given a query $\mathbf{q}$, find the key $\mathbf{k}_j$ most similar and return its value $\mathbf{v}_j$:

Soft retrieval (attention): Replace the hard $\arg\max$ with a soft distribution:

This is differentiable and allows gradient-based learning.

The QKV Formulation

For all queries simultaneously:

Step by step:

1. Compute similarity scores: $\mathbf{S} = \mathbf{Q}\mathbf{K}^T \in \mathbb{R}^{T_q \times T_k}$. Entry $S_{ij} = \mathbf{q}_i^T\mathbf{k}_j$.
2. Scale: $\mathbf{S} \leftarrow \mathbf{S}/\sqrt{d_k}$.
3. Normalize: $\mathbf{A} = \text{softmax}(\mathbf{S})$ (row-wise). Each row sums to 1.
4. Aggregate: $\text{output} = \mathbf{A}\mathbf{V}$. Each output row is a weighted average of value rows.

Why Scale by Root-dk

Theorem. If $q_i, k_j \overset{\text{iid}}{\sim} \mathcal{N}(0, 1)$, then $\mathbf{q}^T\mathbf{k} = \sum_{i=1}^{d_k}q_ik_i$ has:

Proof. Each term $q_ik_i$ has mean 0 and variance $\text{Var}(q_i)\text{Var}(k_i) = 1$. Sum of $d_k$ independent terms has variance $d_k$. $\blacksquare$

Problem: For large $d_k$ (e.g., 128), dot products have standard deviation $\sqrt{128} \approx 11$. The softmax input has values in $[-30, +30]$, pushing softmax into saturation (near one-hot), causing:

• Near-zero gradients through softmax.
• Attention concentrating on a single key.

Solution: Divide by $\sqrt{d_k}$ to normalize variance to 1, keeping softmax in its sensitive regime.

Attention as Soft Dictionary Lookup

Interpretation: Think of attention as querying a dictionary:

• Keys = addresses/labels of stored information.
• Values = stored content.
• Query = what you're looking for.
• Output = content blended by relevance.

The attention weights $\alpha_{ij}$ tell us: "how much does output position $i$ read from input position $j$?"

Properties:

• Each row of $\mathbf{A}$ is a probability distribution (sums to 1, non-negative).
• Attention is permutation equivariant w.r.t. key-value pairs (order doesn't matter without positional encoding).
• Attention is not equivariant to query permutation in the causal case.

Masking: Causal and Padding

Causal mask (for autoregressive models):

After softmax, positions with $-\infty$ get weight 0: position $i$ can only attend to positions $\leq i$.

Padding mask: For variable-length sequences padded to the same length, mask padding positions with $-\infty$ to prevent attending to padding tokens.

Computational Complexity

Total: $O(T^2 d)$ compute, $O(T^2)$ memory.

The quadratic scaling in $T$ is the fundamental bottleneck of standard attention. For $T = 128K$ tokens: the attention matrix has $128K^2 \approx 16$ billion entries — infeasible to store in GPU memory.

Common Pitfalls

Pitfall 1. Forgetting the scaling factor. Without $1/\sqrt{d_k}$, attention weights collapse to near-one-hot for large $d_k$, and gradients vanish.

Pitfall 2. Applying causal mask after softmax instead of before. The mask must be added to logits (before softmax) to properly zero out future positions.

Pitfall 3. Confusing $T_q$ and $T_k$. In self-attention they're equal; in cross-attention (encoder-decoder), they differ.

Summary

• Attention computes data-dependent weighted averages: $\text{softmax}(\mathbf{QK}^T/\sqrt{d_k})\mathbf{V}$.
• Scaling by $1/\sqrt{d_k}$ prevents softmax saturation.
• Interpretable as soft dictionary lookup.
• Causal masking enforces autoregressive structure.
• Quadratic in sequence length: $O(T^2d)$ compute, $O(T^2)$ memory.

Exercises

Exercise 1. Compute the attention output for $\mathbf{Q} = [1, 0]$, $\mathbf{K} = [[1,0],[0,1]]$, $\mathbf{V} = [[1,2],[3,4]]$ with $d_k = 2$.

Exercise 2. Prove that $\text{Var}(\mathbf{q}^T\mathbf{k}) = d_k$ under the stated assumptions.

Exercise 3. Show that attention is equivariant to permutation of key-value pairs.

Exercise 4. Compute the exact memory (in bytes, FP16) needed to store the attention matrix for $T=8192, H=32$.

Exercise 5. Derive the gradient $\partial\mathcal{L}/\partial\mathbf{Q}$ through the attention operation.