Scaled Dot-Product & Multi-Head Attention

In Session 17 we saw what attention computes. Now let us write the exact formula used inside every transformer, understand the small but very important “÷ √dₖ” in it, and see why transformers use not one but many attention heads at the same time.

01The formula

This one line is the heart of the transformer:

Here Q, K and V are matrices: one row per token. If the sentence has n tokens and each key has dₖ numbers, then Q·Kᵀ is an n × n table of scores: every token compared with every token, all in one matrix multiplication.

02Why divide by √dₖ?

A dot product adds up dₖ multiplications. The more numbers you add, the bigger the result can become. If query and key values are roughly random with spread 1, the spread of their dot product grows like √dₖ:

Big scores are bad for softmax. When one score is much bigger than the others, softmax gives it almost all the weight (close to 1.0) and the rest get almost 0. The model stops mixing information, and the gradients become tiny, so learning slows down. Dividing by √dₖ brings the scores back to a spread of about 1.

03Masks: what a token is not allowed to see

The optional mask step sets some scores to minus infinity before softmax, so their weight becomes exactly 0. Two common masks:

• Padding mask: sentences in a batch have different lengths, so short ones are padded with [PAD]. Nobody should attend to padding.
• Causal (look-ahead) mask: in GPT-style models, a token must not see future tokens while learning to predict the next one.

04Multi-head attention

One attention head produces one set of weights, so it can focus on one kind of relationship at a time. But language has many relationships at once: grammar, meaning, which noun a pronoun refers to, nearby words. So transformers run several attention heads in parallel, each with its own W_Q, W_K and W_V.

Here is what different heads typically learn. Switch between heads:

Does it cost more?

Not really. Each head works on a smaller slice (64 instead of 512 numbers), so 8 heads together cost about the same as one big head. You get many views for the price of one.

05Code

Scaled attention in NumPy

Same code as Session 17, with one change: / np.sqrt(d_k).

import numpy as np
np.random.seed(42)

X = np.random.randn(4, 8)
Wq, Wk, Wv = [np.random.randn(8, 4) for _ in range(3)]
Q, K, V = X @ Wq, X @ Wk, X @ Wv
d_k = K.shape[1]

scores = Q @ K.T / np.sqrt(d_k)               # the only change
weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights /= weights.sum(axis=1, keepdims=True)
print(weights.round(2))

Compare with Session 17: the 0.97 row became 0.85 and the 0.75 row became 0.62. The weights are softer, so more words contribute.

Multi-head attention in PyTorch

import torch
import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
x = torch.randn(1, 4, 512)                    # 1 sentence, 4 tokens, 512 numbers each

out, weights = mha(x, x, x)                   # self-attention: Q, K and V all come from x
print(out.shape)
print(weights.shape)                          # averaged over the 8 heads by default

06Summary

• Attention = softmax(Q·Kᵀ / √dₖ) · V, done for all tokens in one go.
• Dividing by √dₖ keeps scores small, so softmax doesn’t saturate and gradients stay healthy.
• Masks hide padding and, in decoders, the future.
• Multi-head attention runs several smaller heads in parallel; each learns a different relationship, at about the same cost.

07Quick quiz

08References

• Vaswani et al., Attention Is All You Need (2017), section 3.2
• PyTorch docs: nn.MultiheadAttention