Producto punto escalado y multi-head

When key vectors are long, dot-product scores grow huge and softmax becomes razor-sharp. That kills gradients, so we need to scale the scores down.

The fix is to divide each score by the square root of the key dimension. This keeps the numbers in a stable range before softmax.

Put it together: score, scale, softmax, then blend values. This four-step recipe is the famous scaled dot-product attention.

A single attention head learns just one way to relate tokens. But language has many relationships at once, so one head is limiting.

Multi-head attention runs several attention heads side by side, each with its own Q, K, V projections, each attending in a different subspace.

You do not make the model wider. You split the model dimension across heads, so each head works on a smaller slice in parallel.

To run heads in parallel, you reshape Q, K, V to add a heads axis, then attend within each head independently.

One head may track grammar, another may link a pronoun to its noun. Multiple heads capture diverse patterns the same input contains.

After each head produces its output, you concatenate them back into one vector of the original model dimension.

A last output projection mixes the concatenated head results, letting the model combine what every head discovered.

In practice PyTorch ships nn.MultiheadAttention, so you rarely hand-roll the reshapes. Knowing the math still helps you debug.

Let's confirm why we scale the scores.

You saw how scaling stabilizes softmax and how multi-head attention splits the work across parallel heads, then concatenates and projects. Great progress!