Logits and Unembedding

After the last transformer block you still have vectors, one per token. Generation cares about the last vector (the position you are predicting).

Unembedding is a linear map: multiply that vector by a big matrix with one column (or row) per vocab item. You get logits: raw scores, one per token id. They can be negative. They do not sum to 1. They are not chances yet.

Softmax turns logits into chances that are positive and sum to 1. Decoding (next lesson) picks an id from those chances, or just takes the argmax of the logits. The embedding table went id → vector. Unembedding goes vector → id scores. Some models use the same matrix both ways (tied weights).

If you already have log-chances from a stack, do not softmax twice. If you have logits, do not treat them as chances that already sum to 1. Those two bugs show up in homemade samplers and in confused evals.

A wrong picture

A wrong picture is: “a logit is a probability.” A logit is a raw score. It can be 12.4 or -3.1. Softmax is the map to chances. People say “logit” loosely. Your code should not.

Another wrong picture is: “we need softmax to pick a winner.” Argmax of logits is the same winner as argmax of softmax chances. Softmax is for sampling and for the training loss, not for picking a winner you already know.

A third wrong picture is: “the last-token vector of a chat model is a great embedding for retrieval.” Sometimes people pool it anyway. Chat models were not always trained for that. Dedicated embedders are. Mixing those vectors is a space bug.

The map in words

Hidden vector h of width d. Unembedding matrix U with V rows (or columns) of width d. Logit for vocab item k is the dot product of h with row k of U.

That is one number per tile in the vocab. Softmax:

1. Subtract the max logit (stable).
2. Exp each.
3. Divide by the sum.

The largest logit becomes the largest chance. Temperature (next lesson) divides logits before this softmax.

Tied weights: U is the embedding table (transposed). Then “tokens that sit nearby as rows” are also “tokens that are easy to predict from nearby hidden states.” A tokenizer change breaks both ends at once.

A tiny example in words

Vocab: search, sql, finish. Hidden h = [0.8, 0.1]. Rows of U:

• search: [1.0, 0.0]
• sql: [0.2, 1.0]
• finish: [0.0, 0.3]

h is closer to search, so search should win both as logit and as chance. If you bump the sql row toward h, sql can overtake. That bump is what training does, slowly, on millions of tokens.

Dots to logits to chances

Lists of numbers. Print logits, chances, and argmax. Use math.exp, not extra libraries.

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

h is closer to the search row, so search wins. sum probs should print 1.0 (tiny float noise is fine). Argmax from logits and from chances should match.

Change h to [0.0, 1.0]. Sql should win. That is the unembedding: direction of the last vector vs direction of each vocab row.

If you softmax the already-printed probs again, you will sharpen them and they will no longer mean “the model’s chances.” One softmax after logits. Stop.

Confidence is a shape, not a vibe

If the top chance is 0.92 and the rest is a thin tail, the unembedding is peaked. If the top chance is 0.34 and the rest is a pile, the unembedding is saying “I don’t know” in distribution form. Argmax still returns some id. A peaked wrong id is a confident mistake. A flat distribution is an uncertain mistake. Treat them differently in a router.

Logits on tool names (a small subset of the vocab) are a router you can threshold. You do not need the whole 100k-way softmax to notice the model is torn between search and sql. You do need to remember that the rest of the vocab still ate some mass if you computed softmax over all tiles.

The last-token vector is also what some tools pool for embeddings. Chat models were not always trained for that; dedicated embedders are. Do not mix.

Common mistakes

Softmax twice: you took chances, then treated them as logits, then softmax again. The peak gets fake-sharp. Log the raw logits or the first chances. Not both stacked.

Comparing logit 4.2 from model A to logit 4.2 from model B. Scales differ. Unembedding matrices differ. Compare chances inside one model, or compare winners, not raw scores across stacks.

Reading logits at a pad index after a clumsy collator. You will get a confident pad or a nonsense tile. The padding lesson is this bug. Unembedding is innocent.

Using the last hidden state of a chat decoder as a memory vector for retrieval without checking neighbors. Sometimes it works. Often it clusters by length or by “assistant tone.” Dedicated embedders exist because this shortcut is unreliable.

Thresholding the whole vocab’s top chance when you only care about three tool names. Mass leaked to “the”, “,” and random tiles. Restrict the view to the legal action set, then threshold.

How agents use this

Logits on tool names are a router you can threshold. If the top chance is 0.34 and the rest is a pile, the unembedding is saying “I don’t know.” Escalate. Do not pretend argmax is a plan.

When you log a decision, log the top ids and their chances (or logits), not only the decoded string. “It called sql” and “it called sql at 0.41 vs search at 0.39” are different incidents.

Do not softmax twice. Do not treat logits as percents. Do not compare logits from two different models as if they share a scale.

• Logits: raw scores, one per vocab id.
• Softmax once: then chances.
• Argmax: same winner with or without softmax.
• Flat top: escalate; do not “just decode.”
• Tied tables: vocab change hits both ends.

Tip:Argmax of logits is the same as argmax of softmax chances. Softmax is for sampling and for loss, not for picking a winner you already know.