Attention Mechanism Theory

Assume $q_1, \ldots, q_{d_k}, k_1, \ldots, k_{d_k}$ are mutually independent random variables, each with mean $0$ and variance $1$. Then the dot product $q^\top k$ has:

$\mathbb{E}[q^\top k] = 0, \qquad \text{Var}(q^\top k) = d_k$

The scaled dot product $q^\top k / \sqrt{d_k}$ has variance $1$, regardless of $d_k$.

The dot product is a sum of $d_k$ independent terms $q_i k_i$, each with variance $1$. By the independence assumption, the variance of the sum is the sum of the variances: $d_k$. Without scaling, as $d_k$ grows, the dot products grow in magnitude, pushing the softmax inputs into regions where the gradient is near zero (softmax saturation). Dividing by $\sqrt{d_k}$ normalizes the variance to $1$, keeping the softmax in a regime with useful gradients.

Without scaling, a model with $d_k = 512$ would have dot products with standard deviation $\sqrt{512} \approx 22.6$. Softmax inputs of magnitude 20+ produce outputs extremely close to 0 or 1, with gradients on the order of $10^{-9}$. Training would effectively freeze. The $\sqrt{d_k}$ scaling is not a minor numerical convenience. It is essential for trainability.

The assumption that $q$ and $k$ entries are independent with unit variance holds approximately at initialization (with proper weight initialization) but may not hold after training. In practice, the model learns to calibrate its own attention logits, so the scaling factor becomes less critical later in training. However, removing it entirely still causes training instability.

Under the Vaswani et al. (2017) convention $d_k = d_v = d / h$ and ignoring bias terms, multi-head attention has the same weight parameter count as single-head attention with full head dimension $d$:

$\text{Params(MHA)} = h \cdot 3 \cdot d \cdot (d/h) + d^2 = 3d^2 + d^2 = 4d^2$

$\text{Params(single-head)} = 3 \cdot d \cdot d + d^2 = 4d^2$

However, MHA can represent richer attention patterns: each head can specialize in a different type of relationship (syntactic, semantic, positional), and the output projection $W_O$ learns how to combine these patterns.

The equality breaks under other widely used choices. Including biases adds $4d$ parameters. Multi-query attention (MQA) shares $W_K, W_V$ across all heads, giving $d^2 + 2d \cdot (d/h) + d^2$. Grouped-query attention (GQA) interpolates between these extremes. See the MQA/GQA page for the exact counts.

Multiple heads give the model multiple independent "attention channels." One head might track subject-verb agreement, another might track coreference, a third might focus on adjacent tokens. Single-head attention is forced to compress all these patterns into a single set of attention weights, which is a lossy compression. Multi-head attention avoids this by giving each pattern its own subspace.

Empirically, reducing to a single head significantly degrades performance. The multi-head structure is one of the most important design decisions in the transformer. Mechanistic interpretability research has shown that individual heads do specialize: there are "induction heads" that copy patterns, "previous token heads" that attend to the immediately preceding token, and "name mover heads" that track entities. The "one head, one function" picture is a simplification: many heads are polysemantic (a single head implements several behaviors that are entangled across inputs), and head ablation studies (Michel et al. 2019, Voita et al. 2019) find that a sizable fraction of heads can be pruned at inference time with limited loss, indicating that specialization coexists with substantial redundancy.

Scaled dot-product attention can be written as a Nadaraya-Watson kernel regression estimator:

$\text{output}_i = \sum_{j=1}^{n} \frac{\kappa(q_i, k_j)}{\sum_{\ell=1}^{n} \kappa(q_i, k_\ell)} \, v_j$

where the kernel function is $\kappa(q, k) = \exp(q^\top k / \sqrt{d_k})$.

This is a softmax (exponential dot-product) kernel: $\kappa(q, k) = \exp(\langle q, k \rangle / \sqrt{d_k})$. It is not a translation-invariant RBF kernel. Using $q^\top k = (\|q\|^2 + \|k\|^2 - \|q - k\|^2)/2$:

$\exp\left(\frac{q^\top k}{\sqrt{d_k}}\right) = \exp\left(\frac{\|q\|^2 + \|k\|^2}{2\sqrt{d_k}}\right) \cdot \exp\left(-\frac{\|q - k\|^2}{2\sqrt{d_k}}\right)$

The left factor breaks translation invariance: unlike the Gaussian RBF kernel, the softmax kernel depends on $\|q\|$ and $\|k\|$ separately, not just on $\|q - k\|$. Only when $\|q\|$ and $\|k\|$ are (approximately) constant across all query--key pairs does the softmax kernel reduce to an RBF kernel. Layer normalization controls the pre-projection norm but $q = W_Q x$ and $k = W_K x$ are not norm-constrained, so this equivalence is heuristic rather than exact. Tsai et al. (2019) treat the softmax kernel as an asymmetric dot-product kernel, which is the view we use here.

The Nadaraya-Watson estimator is a classical nonparametric regression method: to estimate a function value at a query point, take a weighted average of observed values, where the weights are determined by a kernel measuring similarity between the query point and each data point. Attention is doing exactly this: it estimates the output at each position as a kernel-weighted average of value vectors.

This connection has two major implications. First, it explains why attention works as a form of nonparametric in-context learning: the model can adapt its behavior at inference time by "regressing" on the input context, without updating weights. Second, it opens the door to efficient attention approximations via random feature maps for kernels (the "Performers" approach), which approximate the softmax kernel with $O(n)$ complexity.

The kernel interpretation is cleanest for a single attention head without learned projections. In practice, the learned $W_Q, W_K, W_V$ matrices transform the inputs before the kernel is applied, and multi-head attention applies multiple kernels simultaneously. The kernel analogy is useful for intuition but does not fully capture the representational power of learned multi-head attention.