Scaled Dot-Product Attention

Scaled dot-product attention is a function that maps a set of queries $Q$, keys $K$, and values $V$ to a weighted combination of the values. Each query is compared against every key using a dot product, producing a raw similarity score. These scores are scaled, converted into a probability distribution via softmax, and used as weights to combine the value vectors. Formally, for query matrix $Q \in \mathbb{R}^{n \times d_k}$, key matrix $K \in \mathbb{R}^{n \times d_k}$, and value matrix $V \in \mathbb{R}^{n \times d_v}$,

The term $QK^\top$ produces an $n \times n$ matrix of unnormalized similarity scores — one row per query, one column per key. The division by $\sqrt{d_k}$ is the scaling step that gives this mechanism its name, distinguishing it from plain "dot-product attention," which omits the factor and was observed empirically to train worse for large $d_k$. The word "scaled" is not decorative. It encodes a specific engineering fix motivated by a variance argument: as the dimensionality $d_k$ of query and key vectors increases, the magnitude of their dot product tends to increase as well, even when the vectors themselves are drawn from a fixed, dimension-independent distribution. Left unscaled, this growth pushes the softmax function into a regime with extremely peaked outputs and near-zero gradients almost everywhere — the same saturation phenomenon that plagues sigmoid units fed large inputs. The scale factor $\sqrt{d_k}$ is chosen precisely to cancel this growth, restoring the softmax's input to an $O(1)$ variance regime regardless of how large $d_k$ becomes.

The mechanism proceeds in four steps: project inputs into queries, keys, and values (already assumed available from the Self-Attention Mechanism setup); compute pairwise similarity scores via matrix multiplication; scale and normalize those scores into attention weights with softmax; and use the weights to produce a weighted sum of the value vectors.

Step 1 — Score computation. For every query vector $q_i$ and key vector $k_j$, compute the raw score $s_{ij} = q_i \cdot k_j$. Doing this for all pairs simultaneously is exactly the matrix product $QK^\top$, an $n \times n$ matrix.

Step 2 — Scaling. Divide every entry by $\sqrt{d_k}$:

Step 3 — Normalization. Apply softmax row-wise so each query's scores become a probability distribution over keys:

Step 4 — Weighted aggregation. Compute the output for query $i$ as the weighted sum of value vectors: $o_i = \sum_j \alpha_{ij} v_j$. In matrix form this is $\alpha V$. An optional mask step can be inserted between scaling and softmax to zero out (via $-\infty$) forbidden connections — this is the hook that Masked Self-Attention uses to prevent a decoder from attending to future tokens. ``` function ScaledDotProductAttention(Q, K, V, mask = None):

# Q: n_q x d_k K: n_k x d_k V: n_k x d_v d_k := number_of_columns(K) S := MatMul(Q, Transpose(K)) # n_q x n_k raw similarity scores S := S / sqrt(d_k) # scale step — the crux of this article if mask is not None: for i in range(n_q): for j in range(n_k): if mask[i][j] == BLOCKED: S[i][j] := -infinity # forces softmax weight -> 0 A := SoftmaxRowwise(S) # n_q x n_k attention weights, rows sum to 1 O := MatMul(A, V) # n_q x d_v output return O, A # A often returned for interpretability/analysis

function SoftmaxRowwise(S): for each row r in S: m := max(r) # numerical stability: subtract row max exps := [exp(x - m) for x in r] total := sum(exps) r := [e / total for e in exps] return S ``` Two implementation details matter for correctness. First, the row-wise max subtraction inside softmax is a numerically stability trick (it does not change the mathematical result, since softmax is invariant to additive shifts per row) but it is essential in floating point, since $\exp$ overflows quickly. Second, masking is applied by additive $-\infty$ (or a very large negative number in practice, such as $-10^9$) before the exponential, so that masked positions vanish from the normalization rather than being clipped afterward, which would leave probability mass unaccounted for.

Time complexity. Computing $QK^\top$ multiplies an $n \times d_k$ matrix by a $d_k \times n$ matrix, costing $O(n^2 d_k)$ scalar multiply-adds. The scaling and masking steps touch every entry of an $n \times n$ matrix, costing $O(n^2)$. Row-wise softmax also costs $O(n^2)$ (each of $n$ rows requires $O(n)$ work for max, exponentials, sum, and division). The final product $AV$ multiplies an $n \times n$ matrix by an $n \times d_v$ matrix, costing $O(n^2 d_v)$. Summing these, the total time is

treating $d_k, d_v$ as a common order-of-magnitude dimension $d$. This is the well-known quadratic-in-sequence-length cost of self-attention, and it is the reason long-context transformers require sparse or linear-attention approximations for very large $n$.

Space complexity. The intermediate score matrix $S$ (and its softmax-normalized version $A$) is $n \times n$, so the mechanism requires $O(n^2)$ memory in addition to $O(nd)$ for the inputs and outputs. For a sequence of length $n = 2048$ this is over four million floats per attention head — the dominant memory cost in naive transformer implementations, and the motivation behind memory-efficient variants such as FlashAttention (which recomputes rather than materializes $A$ in full).

Correctness of the scaling constant. The choice of $\sqrt{d_k}$ is not arbitrary; it follows from a variance calculation. Assume each component of $q$ and $k$ is drawn independently with mean $0$ and variance $1$ (a reasonable approximation immediately after typical weight initialization and layer normalization). Then the dot product is

Each term $q_i k_i$ has mean $\mathbb{E}[q_i]\mathbb{E}[k_i] = 0$ (by independence of $q_i, k_i$) and variance

using independence of $q_i$ and $k_i$ and $\mathbb{E}[q_i^2] = \mathrm{Var}(q_i) = 1$. Since the $d_k$ terms are independent, variances add:

So the standard deviation of the raw dot product grows as $\sqrt{d_k}$. Dividing by $\sqrt{d_k}$ restores unit variance:

regardless of how large $d_k$ is. This is exactly the guarantee needed: the input distribution to softmax stays at a fixed scale as the model's hidden dimension grows, instead of drifting toward extreme values that make softmax behave like a hard argmax.

Why unscaled scores hurt training. Softmax's Jacobian involves terms of the form $\alpha_i(\delta_{ij} - \alpha_j)$. When the input scores have large magnitude, softmax saturates: one score dominates, its weight approaches $1$, all others approach $0$, and $\alpha_i(1-\alpha_i) \to 0$ for the dominant index and $\alpha_i \alpha_j \to 0$ for the rest. Consequently gradients backpropagated through softmax shrink toward zero across almost the entire output — precisely the vanishing-gradient problem. Because $\mathrm{Var}(q\cdot k)$ scales linearly with $d_k$, larger head dimensions make this saturation more severe unless corrected; the $\sqrt{d_k}$ divisor is the minimal correction that keeps the variance constant across all choices of $d_k$, which is why it appears as a square root rather than a linear or logarithmic factor.

Tradeoff. The scaling constant is derived under an idealized independence and unit-variance assumption that need not hold exactly once queries and keys are learned rather than random; in practice it is an approximation that works remarkably well empirically and is validated by every major transformer architecture, but it is not a proof that gradients remain healthy under all trained weight distributions — layer normalization elsewhere in the network is what keeps the assumption approximately true throughout training.

Consider a toy setting with sequence length $n = 3$ and key/query dimension $d_k = 4$. Suppose after the query and key projections we have:

Step 1: raw scores $QK^\top$. Row 1 of $Q$ is $(1,0,1,0)$; dotting with each row of $K$:

• vs. $K_1 = (1,1,0,0)$: $1\cdot1+0\cdot1+1\cdot0+0\cdot0 = 1$
• vs. $K_2 = (0,0,1,1)$: $1\cdot0+0\cdot0+1\cdot1+0\cdot1 = 1$
• vs. $K_3 = (1,0,1,0)$: $1\cdot1+0\cdot0+1\cdot1+0\cdot0 = 2$

Repeating for rows 2 and 3 of $Q$ gives the full score matrix:

Step 2: scale. With $d_k = 4$, $\sqrt{d_k} = 2$:

Step 3: softmax per row. For row 1, $\exp(0.5) \approx 1.6487$, $\exp(1.0) \approx 2.7183$. Sum $= 1.6487+1.6487+2.7183 = 6.0157$. Weights: $0.274, 0.274, 0.452$.

Now compare against what happens without scaling (dividing by $1$ instead of $2$), using the raw scores directly in softmax for row 1: $\exp(1)=2.718$, $\exp(1)=2.718$, $\exp(2)=7.389$, sum $=12.825$, weights $0.212, 0.212, 0.576$. The unscaled distribution is already visibly more peaked on this tiny example ($0.576$ vs. $0.452$ on the top choice). In a real transformer $d_k$ is typically $64$ or larger, so raw scores have standard deviation around $\sqrt{64}=8$ rather than $2$; feeding scores with standard deviation $8$ into softmax produces near one-hot distributions almost everywhere, exactly the saturated regime that kills gradient flow. The scale-by-$2$ step above is a miniature version of the same correction that, at realistic dimensions, is the difference between a trainable and an untrainable network.

Step 4: weighted sum. Suppose

Output for row 1: $0.274\,(1,0) + 0.274\,(0,1) + 0.452\,(1,1) = (0.726, 0.726)$. This vector is the attention-weighted blend of value vectors for the first query — a soft interpolation dominated by, but not exclusively determined by, the third key/value pair.

The scaled dot-product formula was introduced as part of the original transformer architecture, designed to replace recurrent and convolutional sequence models entirely with attention-based computation. The scaling term was explicitly justified in that work by the variance argument reproduced above: for small $d_k$, additive and dot-product attention perform comparably, but for large $d_k$, dot-product attention underperforms unless the scores are scaled down, because unscaled dot products push softmax into regions with extremely small gradients.

Scaled dot-product attention is the innermost computational kernel of every transformer-based system in production today. Large language models such as GPT-family models, PaLM, and LLaMA compute this exact formula, many times per layer, across dozens of layers, for every token generated. Encoder-only models like BERT use it to build bidirectional contextual representations for classification and retrieval tasks. Vision transformers apply the identical formula to patches of an image instead of tokens of text, treating image regions as a sequence to be attended over. Speech models, protein-structure models such as AlphaFold's attention-based modules, and recommendation systems built on transformer encoders all reuse this same scaled dot-product core. Efficient-attention research — sparse attention, linear attention, FlashAttention-style memory-efficient kernels — is fundamentally about approximating or accelerating this exact computation without changing its mathematical semantics, which underscores how central the scaling correction is to making the base operation trainable in the first place. Understanding this formula precisely is the necessary foundation for understanding Multi-Head Attention, which runs several scaled dot-product attentions in parallel with different learned projections, and Masked Self-Attention, which reuses the identical formula with an added masking step to enforce causality in autoregressive generation.

Ask the tutor: Why does the variance of a sum of independent random variables add, and how does that fact justify the $\sqrt{d_k}$ scale factor specifically?