Temperature and top-p filtering
Medium15 min12 hidden tests
A writing assistant samples each next token from the model's probabilities. Before sampling, it applies a temperature and then top-p (nucleus) filtering. Write top_p_filter(logits, temperature, top_p) that divides the logits by the temperature, takes the softmax, keeps the smallest set of most likely tokens whose probabilities add up to at least top_p, sets every other token to 0 and renormalizes what is left. In the first example, the top two tokens hold 0.5 + 0.3 = 0.8 of the mass, which is past top_p = 0.75, so the other two are dropped.
Set the temperature and top-p
keep the top tokens until their total reaches top_pAt temperature 1.00 the ranked chances are help 0.42, time 0.25, reply 0.14, patience 0.09, message 0.06, note 0.04. The total of the tokens above each one: 0, 0.42, 0.67, 0.81, 0.9, 0.96. Keep while that total is below 0.75: help, time, reply, rescaled to 0.52, 0.31, 0.17.
Kept tokens: help, time, reply.
Examples
Tokens 2 and 0 hold 0.8 of the mass
top_p_filter(np.log([0.3, 0.1, 0.5, 0.1]), 1.0, 0.75)
returns [0.375, 0, 0.625, 0]
Temperature 0.5 sharpens the distribution first
top_p_filter([2.0, 1.0, 0.0], 0.5, 0.9)
returns [0.880797, 0.119203, 0]
Requirements
logitsis a 1-D list of floats,temperatureis above 0, andtop_pis in (0, 1]. Logits can be as large as 1000.- Rank tokens by probability, highest first, ties to the lower index. Walk down the ranking and keep each token while the total of the tokens ranked above it is below
top_p. The top token is always kept. - Return a 1-D float array of probabilities in the original token order: 0 for dropped tokens, and the kept ones rescaled to sum to 1.
The idea behind it
Language models produce a probability for every token in the vocabulary, and a sampler picks one. Two knobs shape that choice. Temperature divides the logits before the softmax. Below 1 it sharpens the distribution toward the top token. Above 1 it flattens it. Top-p, or nucleus sampling, drops the long tail of unlikely tokens. It keeps the smallest set of the most likely tokens whose total probability reaches p, then renormalizes. When the model is confident the nucleus is one or two tokens, and when it is unsure it widens on its own, which fixed top-k filtering cannot do. The token that pushes the total past p belongs to the set. Otherwise a confident model with one token above p would keep nothing. Fixing the tie order makes the filter reproducible, which matters when a team tests a sampler.
Tested in interview case A reproducible sampler with top-k, top-p and speculative acceptance.