Self-Attention, Visualized in 3D

Why divide by the square root of d?

Dot products of random d-dimensional vectors have variance proportional to d. With d_k = 64, raw scores can reach magnitudes where softmax saturates: one weight goes to ~1.0, gradients vanish, learning stalls.

• Dividing by sqrt(d_k) normalizes score variance back to ~1 regardless of head size.
• It is the difference between a soft, trainable distribution and a brittle argmax.
• Same idea as temperature in sampling: the scale factor controls how peaked the distribution is.

Static vs dynamic embeddings

Word2vec gave "bank" one frozen vector, an average of riverbanks and vaults. Self-attention makes embeddings contextual:

• Each layer rewrites a token's vector as a weighted mix of its neighbors' Values.
• In "river bank", the token pulls Value payloads from river; in "bank deposit", from deposit. Same input vector, different outputs.
• Switch between sentences above and click bank/bat to watch the arcs re-route. That re-routing is the meaning shift.
• Stack 12–96 layers and vectors encode syntax, coreference, and world knowledge.

Multi-head attention

One attention pattern per layer is not enough. Transformers run h parallel heads (e.g. 12 heads of d_k=64 inside d=768):

• Each head gets its own W_Q, W_K, W_V matrices, so each learns a different relation: one tracks subject-verb, one tracks adjacent words, one tracks coreference.
• Head outputs are concatenated and passed through an output projection W_O.
• The visualization above shows a single plausible head; real models superimpose dozens.

A worked micro-example (d = 2)

Tiny numbers, real mechanics. Suppose the token bank has query q = [1, 2], and two neighbors expose keys k_river = [1, 2], k_the = [-1, 0]:

1. Scores: q.k_river = 1+4 = 5, q.k_the = -1.
2. Scale by sqrt(2) = 1.41: 3.54 and -0.71.
3. Softmax: e^3.54 = 34.5, e^-0.71 = 0.49; weights = 0.99 and 0.01.
4. Output = 0.99 * v_river + 0.01 * v_the - "bank" becomes almost entirely river-flavored.

Every arc in the 3D scene is this arithmetic at d = 64–128 instead of 2.

Causal masks: why chatbots cannot peek ahead

• Encoders (BERT-style) let every token attend both directions - great for understanding a finished sentence.
• Decoders (GPT-style) apply a causal mask: score(i, j) is set to negative infinity for j > i, so softmax gives future tokens exactly zero weight.
• This is what makes next-token prediction honest: the model must commit to each word using only the past.
• In this visualizer all directions are shown (encoder-style); imagine deleting every arc that points rightward to see the decoder view.

Cost and context windows

• Every token attends to every other: n tokens means an n × n score matrix, so compute and memory grow quadratically.
• Doubling context from 4k to 8k tokens roughly quadruples attention cost. That is why long-context models need tricks: FlashAttention (better memory access), sliding windows, sparse or linear attention.
• KV caching stores each generated token's Key and Value so decoding new tokens does not recompute the past, which is why generating is cheaper than prompting per token.