Inside Transformers, Part 3: Scaled dot-product attention

Part 2 added position information. Now we can derive the calculation that turns token representations into context-aware representations.

Self-attention

Let a sequence contain T token representations, indexed from 0 to T - 1. For a query at position i, self-attention assigns a weight to each permitted key position j and then takes a weighted sum of the corresponding value vectors.

1. Compute a score for each permitted pair.
2. Normalize the scores with softmax so the weights for one query sum to one.
3. Use those weights to mix the value vectors.

With scores e_ij, weights a_ij, and value vectors v_j, the last two steps are

a_ij = softmax_j(e_ij)
o_i   = sum over permitted j of a_ij * v_j

The output o_i is a vector. It contains information gathered from the positions that query i was allowed to read. In an encoder, the permitted set can contain every source position. In a decoder, the causal mask limits it to j <= i.

From token vectors to queries, keys, and values

A raw dot product is a useful teaching example. If x_i and x_j are input vectors, we could use

e_ij = x_i^T x_j

and then apply softmax. This compares the same representation in both roles. A Transformer learns three roles instead. For each input vector x_i, separate learned matrices produce

q_i = x_i W_Q
k_i = x_i W_K
v_i = x_i W_V

The query says what the current position is looking for. The key says what a position can be matched on. The value carries the content returned by a match. The vectors and dot products in the next equations are therefore in learned projected spaces. They are not the original token embeddings.

For one query position i, the scaled attention equations are

e_ij = q_i^T k_j
s_ij = e_ij / sqrt(d_k)
a_ij = softmax_j(s_ij)
o_i   = sum over permitted j of a_ij * v_j

Here, d_k is the width of a key vector. The scale is applied before softmax. A decoder also applies its causal mask before softmax, so future positions do not receive any weight.

Suppose we want the output for token 2. Its query q_2 is compared with the keys k_0 through k_T-1, or only the permitted prefix in a causal decoder. Softmax turns those scores into weights a_20 through a_2,T-1. The output is the weighted sum of v_0 through v_T-1 over that same permitted set.

The score matrix makes the calculation easier to see. Row i contains the scores for query i, and column j corresponds to key j.

Stack the projected vectors into matrices Q, K, and V. The same operation for every query position can then be written as

O = softmax(Q K^T / sqrt(d_k)) V

The softmax is applied row by row. With a causal mask, the implementation excludes entries where j > i before that row-wise softmax.

Why divide by sqrt(d_k)?

The scale controls the spread of the scores entering softmax. Under the usual intuition, assume each component of q_i and k_j is independent, has mean zero, and has variance one. For

S = q_i^T k_j = sum from r=1 to d_k of q_ir * k_jr

each product has expected value zero and variance one. The independent terms therefore give

E[S]       = 0
Var(S)     = d_k
std(S)     = sqrt(d_k)

Dividing by sqrt(d_k) gives a score with variance one under these assumptions. That keeps the logits entering softmax from becoming increasingly spread out just because the key vectors are wider. The argument explains the scale, but it is an initialization heuristic rather than a claim that every trained query and key follows these exact distributions.

The individual output vectors from one head are concatenated with outputs from the other heads and passed through a final learned projection. Part 4 explains that multi-head step, cross-attention, and the causal mask on the decoder.