Scores and the √dₖ scaling

$S = QK^\top$ is $(T \times d_k)(d_k \times T) = T \times T$. Row $i$ is query $i$, column $j$ is key $j$: how much token $i$ is interested in token $j$.

“cat” asks with $q = [0, 2]$; “cat” answers with $k = [1, 2]$. $0 \cdot 1 + 2 \cdot 2 = 4$: the strongest match in the matrix.

Write “sat”’s scores against the three keys.

Row 3 of $S$. “sat” is most interested in “cat”.

With $d_k = 2$, every score is divided by $\sqrt 2 \approx 1.414$. The order inside each row doesn’t change; only the gaps shrink.

From Part 1: with unit-variance entries, a dot product of length $d_k$ has standard deviation $\sqrt{d_k}$. A real head has $d_k = 128$, so raw scores would swing by about 11, and softmax would hand nearly all the weight to one key with tiny gradients for the rest. Dividing brings the spread back to about 1, for any head width.

scores = q @ k.transpose(-2, -1) / d_k ** 0.5    # (B, h, T, T)

• $S = QK^\top$, $T \times T$, $S_{ij} = q_i \cdot k_j$.
• Example: $S = [2\ 2\ 2;\ 0\ 4\ 2;\ 1\ 3\ 2]$.
• Scale by $1/\sqrt{d_k}$ to keep the variance near 1.