Softmax

A model does not emit chances directly. It emits logits: real numbers, one per token (or per tool). They can be negative. They do not sum to 1.

Softmax converts a list of logits z into a categorical distribution:

p_i = exp(z_i) / sum_j exp(z_j)

The exp is huge for large logits, so in code you subtract max(z) first. That does not change p (it cancels) and it stops overflow.

After softmax:

• every p_i is at least 0
• the list sums to 1
• a bigger logit becomes a bigger chance, but not in a straight line

Argmax of logits is the same as argmax of softmax chances. Softmax is for sampling and for loss (cross-entropy next), not for picking a winner you already know.

A wrong picture

A wrong picture is: “logits are already probabilities.” They are not. They can be 1000, or -3, and they need not add to 1. Mixing raw logits with softmax chances on a dashboard makes the dashboard lie. Ask the vendor: is this already softmax, or a raw logit?

Another wrong picture is softmax twice. You squash an already-valid chance list toward a new, usually peakier, list that is not the model’s distribution. If you already have log-chances, you still exp and normalize — same function, once.

A third: skipping subtract-max. exp(1000) overflows. Some environments print inf then nan after you divide. Subtract max(logit) every time. Algebra: exp(z_i - m) / sum exp(z_j - m) equals exp(z_i)/sum exp(z_j) because exp(-m) cancels.

Softmax does not “add information.” It is a map from a list of reals to a list of chances. Temperature (next lesson) rescales logits before this map.

The formula in words

Exp every logit (after subtracting the max). Add those exps. Divide each exp by the total. You now have a categorical: weights that sum to 1.

Tiny numeric. Equal logits [1, 1, 1] → equal chances [1/3, 1/3, 1/3]. Clear winner [4, 1, 0]: the 4 dominates, but the others are not zero. Sum is 1. Huge logits [1000, 999, 998]: subtract 1000 first, then exp of [0, -1, -2], still a valid list — the same shape as softmax of [0, -1, -2].

Winner: argmax of [2.0, 0.1, -1.0] is index 0, same before and after softmax.

Moving parts

Only differences between logits matter. Adding 10 to every logit does not change p. Dividing by temperature (next lesson) does, because it changes differences.

A second walkthrough (tied losers)

Tool logits [2.0, 0.5, 0.5] for search, sql, finish. Max is 2. Shifted: [0, -1.5, -1.5].

Exp: [1, exp(-1.5), exp(-1.5)] ≈ [1, 0.2231, 0.2231]. Sum ≈ 1.446.

Chances ≈ [0.691, 0.154, 0.154]. Search wins, but the two tied losers share the leftover mass equally. Softmax does not pick a unique second place when logits tie. Greedy still picks search. Sampling (next lesson) will draw sql and finish equally often among the tail.

One logit: softmax([7.0]) is always [1.0]. Two equal logits, any size: always [0.5, 0.5]. Huge gap [20, 0]: the winner is ~1.0 for any practical printout; a tiny tail remains in theory.

A Friday ticket

Friday 17:10. A guardrail used cutoff 0.35 on a column named probability. The vendor was sending raw logits. Search sat at 4.2, sql at 1.1. Both passed 0.35. The bot called two tools every turn. After they ran softmax, search was about 0.95 and sql about 0.05. Only search passed the cutoff.

The ticket was titled “duplicate tool calls.” The math was: logits are not chances. They asserted abs(sum(p) - 1) < 1e-6 after softmax, labeled the log column p_tool, and refused to compare a cutoff meant for chances against a logit.

Implement it

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

Equal logits → [0.333, 0.333, 0.333]. Clear winner → about [0.936, 0.047, 0.017] — most of the mass, not all of it. Sum prints 1.0. Stable huge prints the same shape as softmax of [0, -1, -2], about [0.665, 0.245, 0.090]. Argmax prints search. The huge-logit line still prints numbers because we subtracted 1000 first.

Never softmax twice. Never treat logits as if they already sum to 1. In tests, check sum(p) == 1 within a tiny tolerance, and check that a huge logit does not become nan.

Empty list or a single logit: one logit softmax is always [1.0]. Two equal logits always split 50/50, regardless of how large they both are — only differences matter.

Hand-check [4, 1, 0] without a computer. Max is 4. Shifted: [0, -3, -4]. Exp: [1, exp(-3), exp(-4)] ≈ [1, 0.050, 0.018]. Sum ≈ 1.068. Chances ≈ [0.936, 0.047, 0.017] (the box rounds a bit differently if it does not subtract 4 first — algebra says the same p). The winner is not 100%. Softmax always leaves a tail unless the logit gap is huge.

Logits are not embeddings. Softmax is not a normalize-to-length-1 step (that is v / |v| from the vectors lesson). Confusing those two maps is how people “softmax an embedding” and then wonder why ranking broke. Use cosine or dot product for ranking lists of floats that stand for text. Use softmax when you need chances that sum to 1 so you can sample or take -log q_true.

What goes wrong

• Overflow: exp(1000) is inf. Then inf/inf is nan. Subtract max(z) every time. The box already does this. Copy that pattern into production.
• Softmax twice: you squash an already-valid chance list toward a peakier list that is not the model’s distribution. If you already have chances, stop. If you have log-chances, exp and normalize once.
• Empty list: no logits, no distribution. Raise. Do not return [] and later divide by sum = 0.
• Ties: equal logits share mass. Argmax needs a documented tie-break (first index, or a reserved finish). Silent “first in the list” is a policy. Write it down.
• Cutoff on the wrong column: 0.35 on a logit is not 0.35 on a chance. Assert the column sums to 1 if you claim it is a chance.

Production logs: the full small catalog of p_i (tools are few), or at least argmax name, argmax p, and entropy of p. Assert no nan, no negative p, unit sum. A test vector [1000, 999, 998] must return finite chances, same shape as softmax of [0, -1, -2].

How agents use this

Tool-calling models score tools with the same machinery as tokens. Softmax is the last map before a draw. If you log “probability” from a vendor, ask: is that already softmax, or a raw logit?

• Tokens: vocab-sized logit list → softmax → chances. Loss is -log of the chance on the true token (next lesson). Sampling draws from this list (lesson after).
• Ranking: you can softmax scores over a small candidate set to get a distribution over chunks. Entropy of that list is retrieval confusion. This is not required for top-k; it is useful when you need weights that sum to 1 (attention, mixture).
• Loss: softmax + cross-entropy is the usual classification loss. Numerically, people use log-softmax (subtract max, then log of normalized exp) so they never exp then log.
• Temperature: divide logits by T before softmax. Same function, different input. Argmax of logits does not change when T > 0.

Subtract max(logit) inside softmax every time. Overflow is a silent NaN in some environments and a crash in others. The algebra is the same; the float is not.

Watch out:Subtract max(logit) inside softmax every time. Overflow is a silent NaN in some environments and a crash in others.