Scaled Dot-Product Attention

Compare scaled and unscaled attention on the same dot products, then derive the square-root scaling factor and its limits.

Scaled Dot-Product Attention

Scaled dot-product attention scores a query against the available keys, normalizes the scores into weights, and mixes their values. The query/key/value mechanism is covered in self-attention. Here, hold that mechanism fixed and examine one choice: dividing the dot products by √dₖ before softmax.

Interactive visualization

Start with Q0 and width 4. Its four raw scores are [4, 0, 2, −4]. Dividing by √4 = 2 gives [2, 0, 1, −2]. Compare the two weight distributions and the resulting value mixtures. Then change the query or head width; the same fixed vectors remain inspectable below the comparison.

The core formula

\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

Q contains the query vectors, with shape n_q × d_k.
K contains the key vectors, with shape n_k × d_k.
V contains the corresponding value vectors, with shape n_k × d_v.
QKᵀ contains one score per query–key pair, with shape n_q × n_k.
The output has shape n_q × d_v. Value width d_v need not equal query/key width d_k.

The dot product as similarity

A dot product adds the coordinate-wise products:

q \cdot k = \sum_{i=1}^{d_k} q_i k_i

A larger score receives a larger softmax weight within the same query row. The weights are proportional to exponentials of scaled scores, not to the dot products themselves.

The dot product depends on both vector direction and magnitude. For unit-length vectors it equals the cosine of their angle; in general the norms also contribute. A large score is a learned compatibility signal; it does not by itself prove that two tokens mean the same thing. Computing QKᵀ applies this comparison to every query–key pair at once.

Why scale by √dₖ

The original Transformer paper motivates the factor with a specific model of the coordinates: all query and key components are independent random variables with mean 0 and variance 1. Each product q_i k_i then has mean 0 and variance 1, and products from different coordinates are independent. Under those assumptions:

\operatorname{Var}\left(\sum_{i=1}^{d_k}q_i k_i\right)=d_k, \qquad \operatorname{Var}\left(\frac{q\cdot k}{\sqrt{d_k}}\right)=1

The standard deviation, rather than a bound on possible scores, is √dₖ. For width 512 this model gives a raw-score standard deviation of about 22.6. Correlated coordinates, nonzero means, or different variances can change this result. The authored vectors in the figure do not satisfy this random-coordinate model.

Large differences between scores can make softmax concentrate almost all weight on one key. Its derivatives can then become very small. Scaling moderates that tendency; it does not guarantee healthy gradients, stable training, or an even weight distribution for every input. Adding the same large constant to every score changes no weights at all.

For the same finite scores, dividing by √dₖ with d_k > 1 preserves their ordering and makes the distribution at least as spread out. The figure measures this with Shannon entropy, using natural logarithms:

H(w)=-\sum_{j=1}^{n_k}w_j\ln w_j, \qquad 0\leq H(w)\leq\ln n_k

Entropy is reported in nats. A uniform distribution over four keys reaches ln(4) ≈ 1.386 nats. Greater entropy here describes the weight distribution, not accuracy. Increasing head width in the figure also changes the dot products, so it is a separate comparison from scaling one fixed score row.

The computation, step by step

S = QK^T

Score every query against every key. For self-attention with n tokens, the score matrix is n × n; cross-attention can have different numbers of queries and keys.

S_{scaled} = \frac{S}{\sqrt{d_k}}

Divide by the square root of query/key width. An optional mask, such as the causal mask in masked attention, sets forbidden logits to −∞ before softmax. Each row needs at least one allowed key for this normalization to be defined.

A = softmax(S_{scaled})

Softmax over the key axis turns each query's allowed scores into nonnegative weights that sum to 1. Implementations subtract each row's maximum before exponentiation to avoid overflow; this does not change the distribution.

Output = AV

Each output is the weighted sum of the corresponding value vectors. The figure keeps those values fixed at width 2, so you can see the output change when only the query, head width, or scaling changes.

Dense self-attention performs O(n²(dₖ + dᵥ)) arithmetic, commonly written O(n²d) when the two widths are comparable. Materializing its scores or weights uses O(n²) elements per head. FlashAttention computes exact attention without storing the full score and weight matrices in global memory, improving memory use and data movement; it still evaluates all query–key pairs. Sparse attention restricts those pairs, while linear attention changes the formulation under its own feature and approximation assumptions.

Multi-query attention shares key and value projections across query heads to reduce KV storage. It retains this scaled-softmax scoring rule.