Weighted Averages and Attention

A weighted average is sum(w_i * x_i) where the weights are chances: they are ≥ 0 and they sum to 1. You already met this in the sums lesson. Softmax builds those weights.

Attention is that idea with extra names:

1. Compare a query to each key. Each comparison is a score (often a dot product).
2. Softmax the scores. Now you have weights.
3. Take the weighted average of the value vectors.

The output is one vector: a mix of the values, with more mix from the keys that matched the query. That is “look at the relevant tokens.” You do not need a GPU to see it. You need a query, some keys, some values, and softmax.

(Transformer internals, one line: real models also scale the dots by sqrt of dimension and run several of these mixes in parallel heads. The object did not change.)

A wrong picture

A wrong picture is: “attention is the model thinking” or “attention means the model understands that span.” Attention means those softmax weights were large on that span’s value. You can log weights in a toy loop. In a giant model you usually cannot — you infer it from behavior.

Another wrong picture is: “the output is always the value with the biggest key.” Only if softmax is extremely peaked (huge score gaps, or tiny temperature). If scores are equal, the output is the mean of the values. If every weight is ~1/n, attention is just a mean. The query is not distinguishing anything. Check the scores.

A third: mixing queries and keys from different spaces (linear-maps lesson). Dot products then rank noise. Same transform on both, or none.

The formula in words

Score_i = dot(query, key_i). Weights = softmax(scores). Output = sum weight_i * value_i (slot by slot).

Tiny numeric. Query [1, 0]. Keys [1, 0] and [0, 1]. Values [10, 0] and [0, 10]. First score is 1, second is 0. Softmax puts most weight on the first. Output looks like [10, 0] mixed with a little of the second. Flip the query to [0, 1] and the output looks like the second value.

If both scores were 0, weights 50/50, output [5, 5] — the mean.

Moving parts

If scores are equal, weights are 1/n and the output is the mean of the values. Attention is then just an average. The query is not distinguishing anything.

A second walkthrough (three slots)

Query [1, 0]. Keys [1, 0], [0.7, 0.7], [0, 1]. Scores: 1, 0.7, 0.

Max is 1. Shifted: [0, -0.3, -1]. Exp ≈ [1, 0.741, 0.368]. Sum ≈ 2.109.

Weights ≈ [0.474, 0.351, 0.174]. Most mass on slot 0, but not a hard pick. Softmax of 1 vs 0.7 vs 0 is still a mix.

Values [8, 0], [0, 8], [4, 4]:

• x = 0.4748 + 0.3510 + 0.174*4 ≈ 3.792 + 0.696 = 4.488
• y = 0.4740 + 0.3518 + 0.174*4 ≈ 2.808 + 0.696 = 3.504

Output about [4.49, 3.50]. Flip the query to [0, 1] and the scores become 0, 0.7, 1 — mass moves toward the last value. Same mixer, different mix.

Huge score gap, say [10, 0, 0]: weights ≈ [1, 0, 0], output ≈ first value. Tiny temperature on the scores (divide scores by T < 1 before softmax) peaks the same way.

A Friday ticket

Friday 15:20. A postmortem said “the model attended to the refund policy.” Someone had dumped three attention weights: 0.34, 0.33, 0.33. That is a mean, not a look. The scores were nearly equal. The mixer was averaging the refund FAQ with a password snippet and a cafeteria line. The generator then blended them.

They started logging attn_max_weight, entropy of the weights, and the score list. A rule: if max weight < 0.45 on a 3-slot memory, treat the mix as confused — ask a clarifying question instead of quoting the average.

A 2-token picture

Two memory slots. A query that looks like slot 0 should put most weight on value 0.

Run to execute this in your browser. Nothing is sent to a server.

First query sits on the first key: weights about [0.731, 0.269], output about [7.31, 2.69] — mostly the first value [10, 0], some mix of [0, 10]. Second query sits on the second key: weights flipped, output about [2.69, 7.31]. Softmax of scores 1 and 0 is not a hard 1.0 — exp(1) vs exp(0) is about 2.718 vs 1, so 73% / 27%. If you wanted a harder pick, larger score gaps (or lower T on the scores) would peak the weights.

If the scores were equal, the output would be the mean of the values. That is attention with a confused query.

What goes wrong

• Equal scores: weights 1/n, output is the mean. Logging “attended to chunk 1” is a lie if weight is 0.34 of 0.33. Print the weights.
• Score overflow: same as softmax. Subtract max(score) before exp. A huge dot from an unnormalized key can nan the mix.
• Mismatched spaces: query from embedder A, keys from embedder B. Dots rank noise. Same transform on both, or none.
• Mixing query/key/value lengths: zip truncates. Assert one dimension for keys vs query, and all values the same width.
• Hard select vs mix: top-1 of values is not attention unless softmax is extremely peaked. If you wanted a hard select, say so. If you wanted a mix, check entropy of the weights.

Production logs: scores, weights (they are few in a toy memory), output magnitude, max(weight), entropy of weights. Assert weights sum to 1, no nan, and that values have equal dimension. Sampling happens after this mix, on token logits — attention itself is deterministic given query, keys, values.

How agents use this

Long context is “which past tokens should this step mix?” Retrieval outside the net is the same geometry: neighbors in, then a mix or a hard top-k. The generator cannot recover a chunk that never made the candidate set.

• Tokens: inside a model, this mix is how past tokens influence the next. You do not implement it here. You do implement the same mix for memories you own: weighted average of retrieved vectors.
• Ranking: scores are dots (or cosines). Softmax turns them into mix weights. Top-k without softmax is a hard select, not a mix.
• Loss: if you train a tiny mixer, CE or a ranking loss on the output vector is legal. Giant-model attention weights are usually not your training handle.
• Sampling: attention is deterministic given query, keys, values. Sampling happens after, on token logits.

When a paper says “the model attended to the tool result,” they mean those softmax weights were large on that span. Check the numbers. If every weight is ~1/n, the query is not distinguishing anything.

Tip:If every weight is ~1/n, attention is just a mean. The query is not distinguishing anything. Check the scores.