Stochastic Sampling Methods

At temperature 0.0, I am a boring but reliable narrator. At temperature 2.0, I am a jazz musician who has lost the sheet music.

Greedy, Temperature-Unstable AI Agent

1. Pure Random Sampling

The most direct form of stochastic decoding is ancestral sampling: at each step, sample the next token from the full probability distribution. If the model says "the" has probability 0.15, "a" has 0.10, "quantum" has 0.0001, and so on across the entire 50,000-token vocabulary, you sample according to those exact probabilities.

This produces maximally diverse output, but the quality is often poor. The long tail of the vocabulary contains thousands of tokens that are individually very unlikely but collectively hold significant probability mass. Even if each improbable token has only a 0.001% chance, with 50,000 tokens in the vocabulary, sampling from the full distribution occasionally draws rare and contextually inappropriate words, derailing the generation.

Nucleus sampling controls which tokens are eligible for selection, but it does not change the relative probabilities among them. To control how sharply the model distributes probability across candidates, we need a different knob: temperature.

2. Temperature Scaling

Temperature is the most fundamental control knob for stochastic sampling. Before applying softmax, we divide the logits by a temperature parameter T. You will encounter temperature again as a practical API parameter in Chapter 10 and as a training hyperparameter for knowledge distillation in Chapter 16:

The effect is intuitive: Code Fragment 5.2.1 below puts this into practice.

• T = 1.0: The original distribution (no modification)
• T < 1.0: Sharpens the distribution, making high-probability tokens even more dominant. At T → 0, sampling becomes greedy decoding.
• T > 1.0: Flattens the distribution, giving low-probability tokens a better chance. At T → ∞, all tokens become equally likely (uniform sampling).

Figure 5.2.2: Temperature controls the "peakiness" of the distribution. Lower temperatures concentrate probability on top tokens; higher temperatures spread it more evenly.

# Temperature scaling: divide logits by T before softmax.
# Low T sharpens the distribution; high T flattens it toward uniform.
import torch
import torch.nn.functional as F

# Simulating temperature effect on a small vocabulary
logits = torch.tensor([5.0, 3.5, 2.0, 1.0, 0.5, 0.1, -1.0, -2.0])
tokens = ["the", "cat", "dog", "it", "my", "old", "an", "..."]

for temp in [0.3, 0.7, 1.0, 1.5, 2.0]:
 probs = F.softmax(logits / temp, dim=-1)
 top_prob = probs[0].item()
 entropy = -(probs * probs.log()).sum().item()
 print(f"T={temp:.1f} | P('the')={top_prob:.3f} | entropy={entropy:.3f} | dist={[f'{p:.3f}' for p in probs.tolist()]}")

# Nucleus (top-p) sampling: sort tokens by probability, accumulate until
# the cumulative mass reaches p, then zero out all remaining tokens.
def top_p_sampling(logits, p=0.9, temperature=1.0):
 """Apply nucleus (top-p) filtering then sample."""
 scaled_logits = logits / temperature
 probs = F.softmax(scaled_logits, dim=-1)

 # Sort probabilities in descending order
 sorted_probs, sorted_indices = torch.sort(probs, descending=True)
 cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

 # Find the cutoff: first index where cumulative prob exceeds p
 # We keep tokens up to (but not including) this cutoff
 sorted_mask = cumulative_probs - sorted_probs > p
 sorted_probs[sorted_mask] = 0.0

 # Renormalize
 sorted_probs /= sorted_probs.sum()

 # Sample from filtered distribution
 sampled_index = torch.multinomial(sorted_probs, num_samples=1)
 return sorted_indices[sampled_index]

# Demonstrate adaptive behavior
confident_logits = torch.tensor([8.0, 4.0, 1.0, 0.5, 0.1, -1.0, -2.0, -3.0])
uncertain_logits = torch.tensor([2.0, 1.8, 1.6, 1.4, 1.2, 1.0, 0.8, 0.5])

for name, logits in [("Confident", confident_logits), ("Uncertain", uncertain_logits)]:
 probs = F.softmax(logits, dim=-1)
 sorted_probs, _ = torch.sort(probs, descending=True)
 cumsum = torch.cumsum(sorted_probs, dim=-1)
 nucleus_size = (cumsum < 0.9).sum().item() + 1
 print(f"{name}: nucleus size = {nucleus_size} tokens for p=0.9")
 print(f" Probs: {[f'{p:.3f}' for p in sorted_probs.tolist()]}")
 print(f" Cumsum: {[f'{c:.3f}' for c in cumsum.tolist()]}\n")

3. Top-k Sampling

Top-k sampling (Fan et al., 2018) restricts sampling to the k most probable tokens at each step. All other tokens have their probability set to zero, and the remaining probabilities are renormalized to sum to 1.

This eliminates the long tail problem: no matter how flat the distribution is, only k tokens are ever considered. However, top-k has a significant limitation: the optimal value of k varies depending on the context. When the model is very confident (e.g., after "The capital of France is"), even k=10 might include irrelevant tokens. When the model is uncertain (e.g., after "I enjoy"), k=10 might be too restrictive, cutting off perfectly valid continuations. Code Fragment 5.2.2 below puts this into practice.

# Top-k sampling: keep only the k highest-scoring tokens,
# set the rest to -inf, then sample from the truncated distribution.
def top_k_sampling(logits, k=50, temperature=1.0):
 """Apply top-k filtering then sample from the result."""
 # Apply temperature
 scaled_logits = logits / temperature

 # Find the k-th largest value as threshold
 top_k_values, _ = torch.topk(scaled_logits, k)
 threshold = top_k_values[..., -1, None]

 # Zero out everything below threshold
 filtered = scaled_logits.masked_fill(scaled_logits < threshold, float('-inf'))

 # Convert to probabilities and sample
 probs = F.softmax(filtered, dim=-1)
 return torch.multinomial(probs, num_samples=1)

# Example: sampling with different k values
logits = torch.tensor([5.0, 3.5, 2.0, 1.0, 0.5, 0.1, -1.0, -2.0])
tokens = ["the", "cat", "dog", "it", "my", "old", "an", "..."]

for k in [2, 4, 6]:
 filtered = logits.clone()
 threshold = torch.topk(filtered, k).values[-1]
 filtered[filtered < threshold] = float('-inf')
 probs = F.softmax(filtered, dim=-1)
 active = [f"{tokens[i]}({probs[i]:.3f})" for i in range(len(tokens)) if probs[i] > 0]
 print(f"k={k}: {', '.join(active)}")

4. Nucleus (Top-p) Sampling

Nucleus sampling (Holtzman et al., 2020) addresses top-k's fixed-size problem with an elegant idea: instead of keeping a fixed number of tokens, keep the smallest set of tokens whose cumulative probability exceeds a threshold p. This adapts automatically to the shape of the distribution.

When the model is confident, the nucleus might contain only 2 or 3 tokens. When the model is uncertain, it might contain 100 or more. This adaptivity is what makes top-p the most widely used sampling method in production systems. Code Fragment 5.2.3 below puts this into practice.

Figure 5.2.3: Top-p sampling adapts the number of candidate tokens to model confidence. When confident, few tokens suffice; when uncertain, the nucleus expands.

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
inputs = tokenizer("The meaning of life is", return_tensors="pt")

output = model.generate(
 **inputs, max_new_tokens=50,
 do_sample=True, temperature=0.8, top_p=0.92, top_k=50,
 repetition_penalty=1.2
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

# Production equivalent using model.generate()
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
inputs = tokenizer("The future of AI", return_tensors="pt")
output = model.generate(
 **inputs, max_new_tokens=50,
 temperature=0.8, top_k=50, top_p=0.95,
 repetition_penalty=1.2, do_sample=True,
)

5. Min-p Sampling

Min-p sampling is a newer technique that takes a different approach to adaptive filtering. Instead of specifying a cumulative probability threshold, min-p sets a minimum relative probability: any token whose probability is less than min_p × max_probability is discarded.

This is conceptually simple and has appealing properties. When the model is very confident (top token at 0.95), even a min_p of 0.1 only keeps tokens above 0.095, resulting in a tiny nucleus. When the model is uncertain (top token at 0.05), the threshold drops to 0.005, allowing many tokens through. The behavior adapts naturally without the cumulative probability bookkeeping of top-p. Code Fragment 5.2.4 below puts this into practice.

# Min-p sampling: discard any token whose probability falls below
# min_p times the top token's probability, adapting the cutoff dynamically.
def min_p_sampling(logits, min_p=0.1, temperature=1.0):
 """Apply min-p filtering then sample."""
 scaled_logits = logits / temperature
 probs = F.softmax(scaled_logits, dim=-1)

 # Threshold: min_p * max probability
 max_prob = probs.max()
 threshold = min_p * max_prob

 # Zero out tokens below threshold
 filtered_probs = probs.clone()
 filtered_probs[probs < threshold] = 0.0

 # Renormalize and sample
 filtered_probs /= filtered_probs.sum()
 return torch.multinomial(filtered_probs, num_samples=1)

# Compare min-p behavior
for name, logits in [("Confident", confident_logits), ("Uncertain", uncertain_logits)]:
 probs = F.softmax(logits, dim=-1)
 max_p = probs.max().item()
 threshold = 0.1 * max_p
 kept = (probs >= threshold).sum().item()
 print(f"{name}: max_p={max_p:.3f}, threshold={threshold:.4f}, kept={kept} tokens")

6. Typical Sampling

Typical sampling (Meister et al., 2022) takes an information-theoretic approach. The idea is that humans tend to produce words that are neither too predictable nor too surprising. Formally, typical sampling keeps tokens whose information content (negative log-probability) is close to the entropy of the distribution (the expected information content).

A token with probability 0.9 carries very little surprise (low information). A token with probability 0.0001 carries enormous surprise. Typical sampling favors the middle ground: tokens that are about as surprising as you would expect on average. This tends to produce text that feels natural and avoids both boring and incoherent extremes.

All the sampling strategies we have covered so far control which tokens are considered and how probabilities are distributed. Yet even with well-tuned sampling, language models have a persistent tendency to fall into repetitive loops. The next family of techniques tackles this problem directly by penalizing tokens that have already appeared.

7. Repetition Penalty, Frequency Penalty, Presence Penalty

Even with good sampling strategies, language models tend to repeat themselves. Several penalty mechanisms address this:

Repetition Penalty

Introduced by Keskar et al. (2019), repetition penalty directly modifies the logits of tokens that have already appeared in the generated text:

Here θ > 1 reduces the probability of repeated tokens. A value of θ = 1.0 means no penalty; values of 1.1 to 1.3 are common.

Frequency and Presence Penalties

Popularized by the OpenAI API, these work by subtracting from logits based on token counts: Code Fragment 5.2.5 below puts this into practice.

• Frequency penalty: Subtracts α × count(token) from the logit. Penalizes tokens proportionally to how often they have appeared. Good for reducing word-level repetition.
• Presence penalty: Subtracts β from the logit if the token has appeared at all (regardless of count). Encourages the model to explore new topics.

# Repetition and frequency/presence penalties: scale down logits of
# already-generated tokens to discourage the model from repeating itself.
def apply_repetition_penalty(logits, generated_ids, penalty=1.2):
 """Apply repetition penalty to logits for already-generated tokens."""
 for token_id in set(generated_ids.tolist()):
 if logits[token_id] > 0:
 logits[token_id] /= penalty
 else:
 logits[token_id] *= penalty
 return logits

def apply_frequency_presence_penalty(logits, token_counts,
 freq_penalty=0.5, presence_penalty=0.5):
 """Apply OpenAI-style frequency and presence penalties."""
 for token_id, count in token_counts.items():
 logits[token_id] -= freq_penalty * count
 logits[token_id] -= presence_penalty # flat penalty if present
 return logits

# Demonstration
logits = torch.tensor([4.0, 3.0, 2.5, 2.0, 1.5])
tokens = ["the", "cat", "sat", "on", "mat"]
generated = torch.tensor([0, 1, 2]) # "the", "cat", "sat" already generated

original_probs = F.softmax(logits, dim=-1)
penalized = apply_repetition_penalty(logits.clone(), generated, penalty=1.3)
penalized_probs = F.softmax(penalized, dim=-1)

print("Token | Original | Penalized")
for i, t in enumerate(tokens):
 marker = " *" if i in generated.tolist() else ""
 print(f"{t:9s} | {original_probs[i]:.4f} | {penalized_probs[i]:.4f}{marker}")

Notice how the penalty shifts probability mass from already-generated tokens ("the," "cat," "sat") toward new tokens ("on," "mat"), encouraging the model to avoid repetition.

8. Combining Sampling Methods

In practice, these methods are often combined. A typical pipeline might apply transformations in this order:

1. Repetition penalty on the raw logits
2. Temperature scaling
3. Top-k filtering (if used)
4. Top-p filtering
5. Sample from the remaining distribution

9. Lab: Visualizing Sampling Distributions

This lab builds an interactive visualization that compares how temperature, top-k, and top-p reshape the token probability distribution.

# Side-by-side comparison: run the same logit vector through greedy,
# top-k, top-p, and min-p to show how each strategy reshapes the output.
import torch
import torch.nn.functional as F

# Simulate a realistic token distribution from a language model
torch.manual_seed(42)
logits = torch.randn(100) # 100 tokens for visualization
logits[0] = 5.0 # make a few tokens clearly dominant
logits[1] = 3.5
logits[2] = 3.0

methods = {
 "Original (T=1.0)": F.softmax(logits, dim=-1),
 "T=0.5": F.softmax(logits / 0.5, dim=-1),
 "T=1.5": F.softmax(logits / 1.5, dim=-1),
}

# Top-k=10
top_k_logits = logits.clone()
threshold = torch.topk(top_k_logits, 10).values[-1]
top_k_logits[top_k_logits < threshold] = float('-inf')
methods["Top-k=10"] = F.softmax(top_k_logits, dim=-1)

# Top-p=0.9
probs = F.softmax(logits, dim=-1)
sorted_p, sorted_i = torch.sort(probs, descending=True)
cumsum = torch.cumsum(sorted_p, dim=-1)
mask = cumsum - sorted_p > 0.9
sorted_p[mask] = 0
sorted_p /= sorted_p.sum()
top_p_probs = torch.zeros_like(probs)
top_p_probs.scatter_(0, sorted_i, sorted_p)
methods["Top-p=0.9"] = top_p_probs

for name, probs in methods.items():
 nonzero = (probs > 1e-6).sum().item()
 top1 = probs.max().item()
 entropy = -(probs[probs > 0] * probs[probs > 0].log()).sum().item()
 print(f"{name:20s} | active tokens: {nonzero:3d} | top-1 prob: {top1:.4f} | entropy: {entropy:.3f}")

This output reveals the key differences. Low temperature (T=0.5) makes the distribution very peaked, with the top token getting 81% of the probability mass. Top-p=0.9 is the most restrictive here, keeping only 5 tokens and achieving the lowest entropy. These numbers help you develop intuition for how each method reshapes the probability landscape.