Softmax that never overflows
Easy8 min11 hidden tests
A ticket router scores every support ticket against each team. The scores are logits: raw numbers where higher means a better match. Write softmax_rows(logits) that turns each row of logits into probabilities that sum to 1. Some rows hold logits near 1000, where np.exp overflows to infinity, so the function must stay finite there. A logit of -inf marks a team that cannot take the ticket and must get probability 0.
Drag the logits, then shift them all
softmax(z) = softmax(z − max z)softmax = [0.09, 0.245, 0.665]. Each step of 1 between two logits multiplies the ratio of their chances by e.
Probabilities 0.09, 0.245, 0.665.
Examples
Each step of 1 in a logit multiplies the probability by e
softmax_rows([[1.0, 2.0, 3.0]])
returns [[0.0900306, 0.244728, 0.665241]]
Large logits stay finite, and equal logits split evenly
softmax_rows([[1000.0, 1001.0], [-5.0, -5.0]])
returns [[0.268941, 0.731059], [0.5, 0.5]]
Requirements
logitsis a 2-D list or array of shape (n, k): n tickets, k teams. Each row has at least one finite logit.- Softmax each row on its own:
exp(z_j) / sum_i exp(z_i)over the k entries of that row. - Return a float array of shape (n, k) with no
nanorinf, even for logits of 1000 and above. A-inflogit gets exactly 0.0.
The idea behind it
Softmax turns a vector of logits into a probability distribution: p_j = exp(z_j) / sum_i exp(z_i). Larger logits get larger shares, and the shares always sum to 1. In float64, exp(z) overflows to infinity once z passes about 709, and infinity divided by infinity is nan. The fix uses one fact: softmax ignores a constant shift. Multiply the top and bottom by exp(-c) and the c cancels. Choosing c as the row's max makes every exponent 0 or less, so each term lies between 0 and 1, and the largest is exactly 1. The sum is then at least 1, so the division is safe. Classifiers, attention layers and samplers all use this. Attention also sets blocked positions to -inf, and exp(-inf) is exactly 0.