How is dot product used in self-attention?

What you need to know

The dot product in two lines

q = [2, 1],  k1 = [3, 0],  k2 = [-1, 2]q·k1 = 2x3 + 1x0 = 6        q·k2 = 2x(-1) + 1x2 = 0

k1 points roughly the same way as q, so the score is high. k2 is at a right angle to q, so the score is 0. Geometrically, q·k = |q| × |k| × cos(angle).

Dot product versus cosine similarity

Cosine similarity divides by both lengths, so only direction matters. The dot product keeps length. In attention this is useful: a model can make a key longer to make it "louder" for every query. The learned Wq and Wk matrices control both direction and length, and some models add normalisation to queries and keys (QK-norm) to keep lengths in check for training stability.

Why it scales well

For n tokens and dimension d, all pairwise scores are one matrix multiply:

Q (n x d)  times  K^T (d x n)  =  scores (n x n)

GPUs are built for exactly this operation (a GEMM). For n = 8,000 and one head, that is 64 million scores in one call. Across 32 heads and 32 layers, it is about 65 billion scores for one pass over the prompt — which is why long prompts cost real time and money.

Why the scaling factor is sqrt(d_k)

If query and key entries are random with variance 1, their dot product over d_k dimensions has variance d_k. For d_k = 128 the standard deviation is about 11.3. Dividing by sqrt(128) ≈ 11.3 brings it back to about 1, so softmax stays in a range where it can learn.

The older alternative

Additive (Bahdanau) attention scores a pair with a small neural network: v · tanh(W1·q + W2·k). Quality is similar, but it cannot be done as one big matrix multiply, so it is slower on modern hardware. Dot-product attention won for speed.

A real-life example

An e-commerce search assistant reads a long query: "Need a laptop for video editing, 16GB RAM minimum, under Rs 80,000, not too heavy because I travel." When the model processes "heavy", its query vector (after training) points in a direction that matches keys of tokens about the product — "laptop" — rather than "RAM" or "Rs". The dot product with "laptop" is high, so "heavy" is understood as a property of the laptop.

The cost side shows up in the same product: shoppers who paste a whole spec sheet (6,000 tokens) make the prompt pass about 36 times more expensive in attention than a 1,000-token query, since 6 squared is 36. The team caps pasted input and extracts the key specs first.

Follow-up questions to expect

• "Why not use cosine similarity in attention?" — Length carries useful signal and cosine adds extra normalisation cost; scaling by sqrt(d_k), and sometimes QK-norm, controls the magnitude instead.
• "Is the attention matrix symmetric?" — No. Queries and keys use different matrices, so how much A attends to B differs from how much B attends to A.
• "How does FlashAttention relate?" — It computes the same dot products and softmax in small tiles in fast on-chip memory, avoiding storing the n × n matrix.