Square Root of Attention Head Dimension

For example, if a Model uses an attention head dimension (\(d_{k}\)) of 64, then \(\sqrt{d_k} = \sqrt{64} = \mathbf{8}\).
This number 8 remains exactly the same for every single token processed, regardless of what the embedding numbers actually are. 
Why \(\sqrt{d_{k}}\) Prevents Exploding Gradients
When ChatGPT matches a “Query” vector with a “Key” vector, it multiplies them together using a dot product.
    1. The Variance Problem: If your vectors have 64 dimensions, you are multiplying 64 pairs of numbers and adding them together. If each individual number has a standard variance of 1, adding 64 of them together means the final dot product will have a variance of 64.
    2. The Softmax Trap: A variance of 64 means the dot product results can easily become very large positive or negative numbers (like +30 or -30). When you pass these large numbers into the Softmax function to get probabilities, it pushes the outputs to the extreme ends (exactly 0 or exactly 1). [1, 2, 3]
    3. Vanishing Gradients: When Softmax outputs are pushed to the absolute extremes (0 or 1), the mathematical gradient (slope) becomes virtually zero. This causes gradients to vanish, meaning the AI completely stops learning during training. [1, 2, 3]

How \(\sqrt{d_{k}}\) Fixes It
By dividing the dot product by \(\sqrt{d_{k}}\) (dividing by 8 in our example), the variance is pulled back down from 64 perfectly to 1. [1]
This keeps the numbers small, keeps the Softmax function in its “sweet spot” where slopes are steep, and ensures that gradients flow smoothly without exploding or vanishing.
 
=-=-=-=-=-=-=-=-=-=-=-
 

 $d_k$ isn’t a concept that only exists because of multi-head splitting. It’s the general term for “dimensionality of the query/key vectors,” full stop.

Multi-head attention is a special case that happens to divide $d_{\text{model}}$ by $h$ to get $d_k$; 

Single-head case

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

This formula existed before multi-head attention was introduced, and $d_k$ was already in it. In the single-head case, there’s no splitting — the query and key vectors simply are $d_k$-dimensional, where $d_k$ is whatever dimension the model designer chose for queries/keys. There’s no requirement that $d_k$ relate to a head count, because $h = 1$.

Put differently:

$d_k = d_{\text{model}}/h$

is the formula for the multi-head case specifically. Set $h=1$ in that formula and you get $d_k = d_{\text{model}}$ — i.e., in single-head attention, the query/key dimension simply equals the full model dimension (or whatever dimension was chosen for $Q$/$K$, which need not even equal $d_{\text{model}}$ in every architecture).

Symbol Meaning Size (single-head case)
$d_k$ Dimensionality of query/key vectors — defined independently of head count scalar, chosen directly (often $=d_{\text{model}}$)
$h$ Number of heads $1$

The general definition, stripped of the multi-head framing

$d_k$ answers one question only: “how many numbers long is a query/key vector?” That question is meaningful whether there’s 1 head or 8. What multi-head attention adds on top is a reason $d_k$ ends up smaller than $d_{\text{model}}$ — because the full dimension gets divided among heads. But the variable $d_k$ itself predates and doesn’t depend on that division; it’s just relabeled/resized when heads enter the picture.

So the accurate ordering of concepts is:

  1. $d_k$ = dimension of query/key vectors (exists in any attention mechanism, single- or multi-head).
  2. Multi-head attention adds the specific choice $d_k = d_{\text{model}}/h$, so that splitting into $h$ heads doesn’t blow up total parameter count.

Reference

Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. “Attention Is All You Need.” Advances in Neural Information Processing Systems 30.

Loading