Tutorial 14: Sampling Strategies — Temperature, Top-p, Top-k¶
The model outputs logits — raw scores for every token in the vocabulary. To generate text, we must convert these logits into a token. Sampling strategies control how we pick the next token, balancing creativity and coherence.
1. From Logits to Probabilities¶
The model outputs a logit vector of size vocab_size (e.g., 151,936 for
Qwen3). To turn these into probabilities:
This is the log-softmax: it normalizes the logits into log-probabilities that sum to 1 in probability space. We work in log-space for numerical stability.
2. Greedy Decoding (Temperature = 0)¶
The simplest strategy: always pick the token with the highest probability.
This is deterministic — given the same input, you always get the same output. It produces coherent but repetitive text.
3. Temperature Sampling¶
Temperature controls the "sharpness" of the probability distribution:
# Before sampling, scale the logprobs
scaled_logprobs = logprobs / temp
next_token = mx.random.categorical(scaled_logprobs, axis=-1)
| Temperature | Effect | Output |
|---|---|---|
| 0.0 | Picks the most likely token | Deterministic, repetitive |
| 0.5 | Sharp distribution, but allows some randomness | Focused, slightly creative |
| 1.0 | Original distribution | Balanced |
| 1.5 | Flatter distribution | More random, less coherent |
| 2.0 | Nearly uniform | Almost random |
The Math¶
Temperature scales the logits before softmax:
- T < 1: Amplifies differences — high-probability tokens become even more likely. The distribution sharpens.
- T = 1: No change.
- T > 1: Reduces differences — all tokens become more equally likely. The distribution flattens.
At T → 0, the distribution approaches a spike at the max. At T → ∞, it approaches uniform.
4. Top-k Sampling¶
Top-k limits sampling to the k most probable tokens, setting all others to
negative infinity (zero probability):
if top_k is not None and top_k > 0:
# Find the k-th largest probability threshold
kth = min(top_k, logprobs.shape[-1]) - 1
# Get indices of tokens ranked below top-k
mask_elements = mx.argpartition(-logprobs, kth=kth, axis=-1)[:, kth + 1:]
# Set their logprobs to -inf (zero probability)
logprobs[:, mask_elements] = -mx.inf
Example¶
All 20 logprobs: [-2.3, -1.1, -0.5, -0.8, -3.0, -0.3, -1.5, -4.0, ...]
top_k = 3
After masking: [-inf, -inf, -0.5, -inf, -inf, -0.3, -inf, -inf, ...]
↑ these 3 tokens are the only candidates
argpartition is used instead of argsort because we only need the top-k
elements, not a full sort. This is O(n) instead of O(n log n).
5. Top-p (Nucleus) Sampling¶
Top-p dynamically selects the smallest set of tokens whose cumulative probability
exceeds p:
if top_p is not None and top_p > 0:
# Sort by probability (descending)
sorted_idx = mx.argsort(-logprobs, axis=-1)
sorted_logprobs = logprobs[:, sorted_idx]
# Convert to probabilities and compute cumulative sum
sorted_probs = mx.exp(sorted_logprobs)
cumsum = mx.cumsum(sorted_probs, axis=-1)
# Keep tokens where cumulative probability minus their own probability < p
# This means: the token itself is the one that pushes the sum over p
mask_elements = cumsum - sorted_probs < top_p
# Mask out everything after the nucleus
logprobs[:, sorted_idx] = mx.where(mask_elements, sorted_logprobs, -mx.inf)
Why Top-p is Often Better Than Top-k¶
Top-k with k=50 always considers exactly 50 tokens, regardless of the distribution. If the model is very confident (one token has 99% probability), top-k still considers 49 irrelevant tokens. Top-p adapts:
Confident model: P(token_1) = 0.99
top_k=50: considers 50 tokens (mostly noise)
top_p=0.9: considers 1-2 tokens (just the confident one)
Uncertain model: P(top_5) ≈ [0.15, 0.14, 0.13, 0.12, 0.11]
top_k=50: considers 50 tokens
top_p=0.9: considers all 5 (they sum to ~0.65, need all of them)
6. Combining Strategies¶
All three strategies can be applied together. The order matters:
def make_sampler(temp, top_p, top_k):
def sample(logprobs):
# Step 1: Handle greedy decoding
if temp == 0:
return mx.argmax(logprobs, axis=-1)
# Step 2: Top-k filtering (applied first — simplest mask)
if top_k is not None and top_k > 0:
kth = min(top_k, logprobs.shape[-1]) - 1
mask_elements = mx.argpartition(-logprobs, kth=kth, axis=-1)[:, kth + 1:]
logprobs[:, mask_elements] = -mx.inf
# Step 3: Top-p filtering (applied after top-k)
if top_p is not None and top_p > 0:
sorted_idx = mx.argsort(-logprobs, axis=-1)
sorted_logprobs = logprobs[:, sorted_idx]
sorted_probs = mx.exp(sorted_logprobs)
cumsum = mx.cumsum(sorted_probs, axis=-1)
mask_elements = cumsum - sorted_probs < top_p
logprobs[:, sorted_idx] = mx.where(mask_elements, sorted_logprobs, -mx.inf)
# Step 4: Apply temperature
logprobs = logprobs / temp
# Step 5: Sample from the filtered distribution
return mx.random.categorical(logprobs, axis=-1)
return sample
Common Settings¶
| Use Case | Temperature | Top-p | Top-k |
|---|---|---|---|
| Code generation | 0.0 | — | — |
| Factual Q&A | 0.3 | 0.9 | 50 |
| Creative writing | 0.7-0.9 | 0.95 | — |
| Brainstorming | 1.0-1.2 | 0.95 | — |
7. Log-Space Computation¶
We work in log-space throughout to avoid underflow/overflow:
# WRONG: computing probabilities directly
probs = mx.exp(logits) / mx.sum(mx.exp(logits)) # exp(1000) = inf
# RIGHT: log-sum-exp trick
logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
The mx.logsumexp function computes log(sum(exp(x))) in a numerically stable
way by subtracting the max before exponentiating. This is why our sampler
receives logprobs (log-probabilities), not probs.
8. Summary¶
| Strategy | What It Does | When to Use |
|---|---|---|
| Greedy (temp=0) | Always pick the most likely token | Code, factual tasks |
| Temperature | Scales distribution sharpness | Controls randomness |
| Top-k | Keep only top-k tokens | Limits token pool |
| Top-p | Keep tokens until cumulative prob ≥ p | Adapts to confidence |
Sampling is the final step in the generation pipeline. It transforms the model's raw predictions into the creative, diverse text we see from LLMs.