I follow the equations but not the why. Why dot products for similarity, why softmax, and why divide by √d_k specifically? Looking for the mental model, not just the formula. (Beginner-friendly answers very welcome.)
Accepted answer
Mental model: it's a soft dictionary lookup. Each token emits a query; every token also exposes a key and a value. The query is compared to all keys (dot product = unnormalized similarity), softmax turns those scores into weights that sum to 1, and you return a weighted average of the values. So a token 'reads' a blend of the other tokens, weighted by relevance.
@jihoon Dot product is cheap (one matmul over the whole sequence) and the projections Wq/Wk are learned, so the model can shape the geometry however it likes — effectively it learns its own similarity metric. Cosine would throw away magnitude, which sometimes carries signal.
√d_k 부분은 분산 때문입니다. q와 k 성분이 대략 평균0·분산1이고 독립이면, 내적 q·k의 분산은 d_k가 돼요(성분 곱을 d_k개 더하니까). d_k가 크면 점수 크기가 √d_k 스케일로 커지고, softmax에 큰 값이 들어가면 거의 one-hot이 돼서 한 토큰만 보게 됩니다. √d_k로 나누면 분산을 다시 1로 맞춰주는 거예요.
@minseok 정확합니다. 그리고 그게 그냥 '날카로워지는' 문제가 아니라 학습 문제예요 — softmax가 포화되면 기울기가 거의 0이 돼서 학습이 안 됩니다. √d_k는 그 포화를 막는 정규화죠.
Mild contrarian note: the √d_k story is a clean post-hoc justification, but in practice it's a normalization heuristic that happened to train well. The variance argument holds under iid unit-variance assumptions that aren't exactly true after a few layers.
@tbecker I'd push back a little — the variance derivation is standard and not hand-wavy: Var(Σ q_i k_i) = d_k under the stated assumptions, so dividing by √d_k restores unit variance. It's a normalization, yes, but a principled one.
@weizh Fair — I'm not disputing the math under iid; my gripe is people present √d_k as a deep insight when it's 'keep the pre-softmax logits at a sane scale'. We agree on the mechanism, just on the marketing.
Has anyone actually ablated removing the scaling on a real model rather than a toy? Curious how badly it breaks in practice vs theory.
@yukis 작게 재현해봤는데, d_k가 작을 땐 차이가 미미했어요. 근데 head dim을 키우니까(예: 128) 스케일 빼면 초반 학습이 확 불안정해지고 loss가 튀더라고요. 이론이랑 일치합니다.
From cognitive science, the analogy that helped me: it's content-addressable memory. Instead of retrieving by an address/index, you retrieve by what you're looking for (the query) matching the content (keys). Softmax just makes the retrieval soft instead of all-or-nothing.
This is incredibly helpful, thank you all. Follow-up: where does multi-head fit? Is each head just a different similarity function?
@priyanair Basically yes. Each head projects q/k/v into a different subspace, so heads can specialize — one tracks syntax, another coreference, another positional patterns. You then concatenate. The famous 'induction heads' are concrete evidence that different heads learn distinct, interpretable jobs.
Small caveat to add to @amir_r: heads aren't guaranteed to specialize cleanly, and many are redundant (you can prune a lot of them). So 'different similarity function per head' is the intent; in practice it's messier.
Outsider question: is this related to kernel methods at all? It smells like a similarity-weighted average, which is what a kernel smoother does.
@hcole Good instinct — yes, attention can be read as a data-dependent kernel: softmax(qk) is like a (non-symmetric, learned) kernel weighting, and the output is a Nadaraya–Watson-style weighted average of values. The 'linear attention' line of work makes that connection explicit by swapping softmax for an explicit feature map.