When Queries and Keys Speak Different Languages

Open with the concrete scenario: a query vector of dimension 512, a key vector of dimension 64. Frame the driving question — can we still compute a meaningful attention score, and what changes about its behavior?

• Set up the asymmetric setup: Q ∈ R^512, K ∈ R^64
• Pose the question: do scores still behave normally?
• Hint that this is not just theoretical — every transformer scales by sqrt(d_k) for this reason

Ask the learner to commit to an initial prediction before any evidence is shown — what they think happens to the attention distribution when dimension increases.

• Force a single explicit commitment before evidence
• Surface the most common naive assumption explicitly

A simulation where the learner drags a slider for key dimension d_k and observes a live histogram of dot-product scores Q·K sampled from unit-variance vectors, plus the resulting softmax output. As d_k grows, the distribution widens and softmax saturates.

• Observe score variance scale linearly with d_k
• See softmax collapse toward one-hot as variance grows
• Compare multiple d_k values side by side
• Notice gradient-signal loss in saturated regions

Pin down the math behind what the simulator showed: if each component is zero-mean with variance 1, the dot-product has variance d_k. Show three concrete numbers — d_k = 16, 64, 512 — and the corresponding score standard deviations.

• Var(Q·K) = d_k when components are unit-variance and zero-mean
• Standard deviation grows as sqrt(d_k)
• Larger variance → sharper softmax → less informative gradients

Explain the mechanism: softmax exponentials amplify differences. A standard deviation of 22 means typical differences between scores are tens — so one key gets probability ~0.999 and all others get dust. Backprop through saturated softmax gives near-zero gradients.

• Softmax(x_i) = exp(x_i) / Σ exp(x_j)
• Large σ makes one exp dominate the sum
• Saturated softmax → near-zero gradient on all but the winning key
• Result: attention becomes 'pick one and ignore everything else'

Show the resolution: dividing the score by sqrt(d_k) restores unit variance regardless of dimension. This is exactly why 'Scaled Dot-Product Attention' in the original transformer paper divides by sqrt(d_k).

• Var(Q·K / sqrt(d_k)) = 1, independent of d_k
• Softmax receives a well-behaved distribution at every scale
• Gradients flow, training is stable
• This is the 'scaled' in scaled dot-product attention

Let the learner toggle the sqrt(d_k) factor on and off across d_k = 16, 64, 256 and observe the softmax output for the same set of keys. The boundary case: with very small d_k, scaling barely matters; with large d_k, omitting scaling breaks training entirely.

• Quantify when scaling matters and when it doesn't
• See the boundary: below d_k ≈ 16 the difference is cosmetic
• Above d_k ≈ 64 the difference is catastrophic
• Modern transformers all live on the 'matters' side

Transfer to the practical case: cross-modal attention where text queries attend over image keys (or vice versa), or retrieval where query and document embeddings have different sizes. The same variance argument applies — and is why projection layers normalize dimensions before attention.

• Cross-modal and cross-domain attention are the real-world version of this
• Projection heads map both sides to a shared d_k
• Scaling rule still applies after projection
• Failure mode in production: mismatched embeddings before projection

Resolve the driving question directly: dimension mismatch changes attention behavior by inflating score variance, saturating softmax, and killing gradients. The sqrt(d_k) scaling is the minimal correction that makes attention dimension-invariant.

• Restate the driving question and answer it plainly
• Tie together variance, softmax saturation, gradient flow, and the scaling fix
• Point forward: this is one reason modern architectures are careful about d_k