Self-Attention & Multi-Head

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

Read it right to left. The product $QK^\top$ forms a score matrix whose entry $(i,j)$ is the dot product of query $i$ with key $j$, a raw measure of compatibility. Softmax, applied across each row, turns those scores into non-negative weights that sum to one, so each token ends up with a probability distribution over the sequence. Multiplying that weight matrix by $V$ produces, for every token, a weighted average of the value vectors. Because softmax and matrix multiplication are both differentiable, the whole lookup is trainable end to end by gradient descent.

Why divide by $\sqrt{d_k}$?

If the components of $q$ and $k$ are independent with zero mean and unit variance, their dot product $q\cdot k=\sum_{i=1}^{d_k} q_i k_i$ has variance $d_k$, so its typical magnitude grows like $\sqrt{d_k}$. Feed large scores into softmax and it saturates: one weight approaches $1$, the rest approach $0$, and the gradient through softmax nearly vanishes. Dividing by $\sqrt{d_k}$ rescales the scores back to unit variance, keeping softmax in a responsive, well-conditioned regime so learning stays stable.

1. 1ProjectMultiply each token embedding by learned matrices W_Q, W_K, W_V to produce a query, key, and value vector for that token.
2. 2ScoreDot every query with every key (the matrix product QKᵀ). Entry (i, j) measures how compatible token i's query is with token j's key.
3. 3ScaleDivide all scores by sqrt(d_k) so their variance stays near 1 and softmax does not saturate.
4. 4NormalizeApply softmax along each row, turning scores into non-negative weights that sum to 1.
5. 5AggregateMultiply the weight matrix by V so each token's output is a weighted average of value vectors.

Worked example: two tokens by hand

Take two tokens with $d_k=2$. Suppose the projections give the vectors below, and we compute the output for token 1, whose query is $q_1=[1,\,1]$, attending over both tokens.

Scores. $q_1\cdot k_1 = (1)(1)+(1)(1) = 2$, and $q_1\cdot k_2 = (1)(1)+(1)(-1)=0$.

Scale. Divide by $\sqrt{d_k}=\sqrt{2}\approx1.414$, giving scaled scores $\approx[1.414,\,0]$.

Softmax. $e^{1.414}\approx4.11$ and $e^{0}=1$, so the sum is $\approx5.11$ and the weights are $\approx[0.80,\,0.20]$.

Aggregate. Output $=0.80\,[1,\,0]+0.20\,[0,\,1]=[0.80,\,0.20]$.

Token 1 pulls about 80% of its new representation from token 1's value and 20% from token 2's, a genuine, differentiable blend rather than a hard pick.

Multi-head attention

A single attention operation forces every relationship, syntax, coreference, positional cues, through one set of $Q$/$K$/$V$ projections. Multi-head attention instead runs $h$ attention operations in parallel, each with its own learned projections into a lower-dimensional subspace (typically $d_k=d_{\text{model}}/h$). Each head can specialize: one may track subject–verb agreement, another may attend to the previous token, another to long-range topic words. The heads' outputs are concatenated and passed through a final linear projection $W^O$:

$\text{MultiHead}(Q,K,V)=\text{Concat}(\text{head}_1,\dots,\text{head}_h)\,W^O,\qquad \text{head}_i=\text{Attention}(QW_i^Q,\,KW_i^K,\,VW_i^V)$

Because each head works in a subspace with $d_k < d_{\text{model}}$, the total cost is comparable to one full-width attention, yet the layer gains the ability to attend to several patterns at once.

Further reading

• Attention Is All You Need, Vaswani et al. (2017), which introduced scaled dot-product and multi-head attention.
• The Illustrated Transformer, Jay Alammar's visual walkthrough of Q/K/V and attention.
• The Annotated Transformer, Harvard NLP's line-by-line PyTorch implementation.