JJoeven

Curriculum/Neural Nets & Transformers

Attention

Each token builds a query, matches keys, softmax to weights, then a weighted sum of values.

intermediate22 min6 / 24

Attention is a routing trick: each token builds a query vector, every token offers a key, the match scores (dot products) become weights via softmax, and the output is a weighted sum of values.

In one sentence: look up relevant context, mix it in.

For self-attention, queries, keys, and values all come from the same sequence (after linear maps). In this lesson we skip the linear maps and use the embeddings themselves so you can see the arithmetic. Production models learn those maps. The recipe does not change.

This is the mechanism people mean when they say “the model looked at the invoice id.” Sometimes the weights really pile on that id. Sometimes they do not, and the copy still happens in a later layer. Attention weights are a clue, not an explanation API.

A wrong picture

A wrong picture is: “attention is understanding.” It is multiply, add, softmax, and a weighted sum of lists of numbers. High weight on a token does not mean the model used that fact correctly. Low weight does not mean the fact was ignored forever — a later layer can still mix it.

Another wrong picture is: “every token can always see every other token.” Causal (decoder) attention hides the future: token i may not look at j > i. That is what makes next-token prediction honest. Bidirectional attention lets everyone see everyone — good for embeddings, not for left-to-right generation.

A third wrong picture is: “if the prompt is long, attention will find the needle.” Softmax over a long row dilutes. Mass spreads. Lost-in-the-middle is this fact in clothing. You cannot fix it by asking the model to “pay attention.” You fix the sequence.

The recipe for one head

For token i:

  1. score_ij = query_i · key_j (often divided by sqrt(d) so dots do not explode)
  2. weight_i = softmax(scores_i)
  3. out_i = sum_j weight_ij * value_j

The dot product is the same multiply-and-add from the math track. Large dots mean “this query matches this key.” Softmax turns the row of scores into a row of weights that are positive and sum to 1. Then you mix the value lists with those weights.

The scale 1/sqrt(d) is not decoration. Large d makes dots huge. Softmax then turns into almost-argmax: one weight near 1, the rest near 0, and the slopes that training needs die. Scaling keeps the scores in a civilized range.

Causal mask: if j > i, set score_ij to a huge negative number before softmax. After softmax that weight is ~0. Token i cannot peek at the future. Bidirectional skips that mask.

Naive attention is O(n²) in sequence length: every query against every key. That quadratic is why long transcripts are expensive and fuzzy. Kernels can be clever. The bill still grows faster than linear.

A tiny example in words

Three tokens as lists of numbers (width 4):

  • invoice: [1.0, 0.0, 0.0, 0.0]
  • please: [0.0, 1.0, 0.0, 0.0]
  • refund: [0.8, 0.2, 0.0, 0.0]

Refund points mostly the same way as invoice. Under causal attention, the last token may look at all three. Weight should pile on invoice and refund more than on please. That is “the model looked at the noun,” as a cartoon.

Token 0 can only look at itself. Token 1 can look at 0 and 1. That triangle is the causal picture.

Causal weights on three toy tokens

Lists of numbers. Softmax by hand. Print the weight rows and the last mixed vector.

Causal attention weights (query rows, key columns)
1000.350.6500.450.120.43invpleaserefund

Future cells are zero. Refund lines up with invoice, so weight piles there more than on please.

Causal attention weights (query rows, key columns)
One attention step
QueryKeysWeightsMix values

Match the query to keys, softmax to weights, then mix the values. That mix is the output.

One attention step
Live PythonOpen full playgroundpython
Output
Run to execute this in your browser. Nothing is sent to a server.

Row q2 can see tokens 0–2. Because refund aligns with invoice, weight piles on 0 and 2 more than on please. Future cells on earlier rows should print as 0.000 after softmax (the huge negative did its job).

The last output list is a mix of the three rows, with more invoice-ish first slot than please-ish second slot. Bidirectional last row may look similar here because token 2 already sees everyone under causal. The difference shows up on earlier queries, which bidirectional lets look forward.

Read the weight matrix whenever you debug a toy. Rows are queries. Columns are keys. That picture is the whole mechanism.

Dilution, copy, and honesty

If you append twenty copies of a stack trace, each new row of softmax has more keys to share mass with. The invoice id still sits there. Its weight often shrinks. That is why editing the sequence (summarize, drop stale tools, put the goal at the end) is the real “attention control” you have from outside the weights.

Copying a UUID can happen in one hop: the query at the tool-argument position matches the key at the UUID position, the value carries those dimensions, the next layer reads them. If it fails, the UUID was too far, too drowned, or split into ugly tiles — not “attention is broken.” Fix the sequence and the tokenizer.

Causal masking is honesty for training: you cannot predict token i+1 using token i+1. At generation time the future does not exist yet, so the mask matches reality. Bidirectional models that see the whole sentence are for understanding jobs (classify, embed), not for writing the next id.

How agents use this

When a 20-step loop transcript is dumped raw, attention can see the first tool error in theory. In practice the mass spreads over repeated stack traces and the model re-commits the error. You help attention by editing the sequence: summarize, drop stale tools, put the goal and the latest observation near the end.

Do not prompt “look carefully at every token.” Softmax cannot give every token a large weight. The weights sum to 1. If you need a fact to win, make it short, unique, and well placed, not buried in a 40k-token dump.

  • Edit: you cannot set the weights; you can change the keys they see.
  • Dilute: long rows spread mass; shorten.
  • Mask: future stays hidden in decoder stacks; do not expect bidirectional tricks from a chat generator.
  • Explain: do not ship “attention heatmap” as proof the agent used a policy.
  • Copy: ids that must be copied need a clear, nearby home in the sequence.
Tip:Softmax over a long row dilutes. Lost-in-the-middle is this fact in clothing. You cannot fix it by asking the model to “pay attention.”

Check your understanding

In causal self-attention, why are some scores set to a huge negative number?