Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Nucleus Sampling › Implementation

Michael BrenndoerferUpdated July 30, 20258 min read

Language AI HandbookAll chapters

Part 28 of 67GPT Architecture

Chapter 204Nucleus Sampling › Implementation

Reading section 6 of 14 · Implementation · 8 minSection 6 of 14 · 8 min

Section 6 of 14 · 8 min read

Nucleus Sampling › ImplementationOne section · Sections in this chapter6 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Now that we understand the algorithm conceptually, let's implement it. We'll build nucleus sampling from scratch, verify it works correctly, then see how to use the production-ready version in Hugging Face's transformers library.

Building the Core AlgorithmLink Copied

The implementation follows our algorithm directly. We'll work with logits (the raw model outputs before softmax), apply optional temperature scaling, convert to probabilities, then perform the nucleus truncation.

In[4]:

Code

import torch
import torch.nn.functional as F


def nucleus_sample(
    logits: torch.Tensor, p: float = 0.9, temperature: float = 1.0
) -> int:
    """
    Sample from the nucleus (top-p) of a probability distribution.

    Args:
        logits: Raw logits from the model (before softmax)
        p: Cumulative probability threshold
        temperature: Temperature scaling factor

    Returns:
        Sampled token index
    """
    # Apply temperature scaling
    scaled_logits = logits / temperature

    # Convert to probabilities
    probs = F.softmax(scaled_logits, dim=-1)

    # Sort probabilities in descending order
    sorted_probs, sorted_indices = torch.sort(probs, descending=True)

    # Compute cumulative probabilities
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

    # Find the cutoff index where cumulative probability exceeds p
    # We shift by 1 to include the token that crosses the threshold
    sorted_indices_to_remove = cumulative_probs > p
    sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
    sorted_indices_to_remove[0] = False

    # Zero out probabilities for tokens outside the nucleus
    sorted_probs[sorted_indices_to_remove] = 0.0

    # Renormalize
    sorted_probs = sorted_probs / sorted_probs.sum()

    # Sample from the filtered distribution
    sampled_idx = torch.multinomial(sorted_probs, num_samples=1)

    # Map back to original vocabulary indices
    return sorted_indices[sampled_idx].item()

One subtlety deserves attention: the shift operation in the cutoff logic. When we compute cumulative_probs > p, the first position where this is True is the token that crosses the threshold. But we want to include that token in the nucleus, not exclude it. The shift ensures we remove tokens after the one that crosses the threshold, keeping the nucleus at exactly the right size.

This shift is critical for correctness. Without it, you would exclude the token that pushes cumulative probability over $p$ , leaving you with a nucleus that covers slightly less than $p$ of the mass. The resulting behavior would be slightly more conservative than intended. With the shift, you always include exactly enough tokens to meet the threshold.

In production code, you would also want to handle edge cases: what if $p = 1.0$ (keep the entire vocabulary)? What if all probabilities are equal (the nucleus becomes the entire vocabulary)? What if the first token alone exceeds $p$ (the nucleus contains exactly one token)? The implementation above handles all these cases correctly because the shift logic works properly at the boundaries.

Verifying Our ImplementationLink Copied

Let's confirm our implementation produces the expected behavior using the "weather" example:

In[5]:

Code

# Create a probability distribution matching our example
token_names = [
    "nice",
    "good",
    "bad",
    "great",
    "terrible",
    "wonderful",
    "cold",
    "hot",
    "other",
]
probs = torch.tensor([0.35, 0.25, 0.15, 0.10, 0.05, 0.04, 0.03, 0.02, 0.01])

# Convert to logits (inverse of softmax, approximately)
logits = torch.log(probs)

# Sample 1000 times to see the empirical distribution
samples = []
for _ in range(1000):
    idx = nucleus_sample(logits, p=0.9, temperature=1.0)
    samples.append(token_names[idx])

Out[6]:

Console

Sample distribution (1000 samples with p=0.9):
  nice: 37.0%
  good: 28.1%
  bad: 18.0%
  great: 11.9%
  terrible: 5.0%
  wonderful: 0.0%

The samples concentrate on the five nucleus tokens, with "wonderful" appearing rarely or never. The empirical frequencies should approximate our renormalized probabilities: "nice" around 39%, "good" around 28%, and so on. Running 1,000 samples provides enough trials to see the distribution clearly, even if individual runs vary due to randomness.

Visualizing Adaptive Nucleus SizeLink Copied

The real insight comes from seeing how nucleus size varies with $p$ . Let's visualize this:

Out[7]:

Visualization

Bar chart showing nucleus size increasing from 1 token at p=0.3 to 7 tokens at p=0.95. — Nucleus size versus probability threshold. Lower p values create smaller nuclei (more focused sampling), while higher p values include more of the tail distribution. The relationship is non-linear because probability mass concentrates in the top tokens.

With our example distribution, $p = 0.3$ requires only 1 token (just "nice" at 35%), while $p = 0.95$ needs 7 tokens. The relationship is non-linear because probability mass concentrates in the top tokens. Moving from $p = 0.3$ to $p = 0.5$ adds just one token, but moving from $p = 0.9$ to $p = 0.95$ adds two because the next tail probabilities are small. Each additional token contributes less marginal probability.

The Adaptive Advantage: Peaked vs. Flat DistributionsLink Copied

The advantage of nucleus sampling becomes clear when we compare it to top-k across different distributional shapes. Consider two scenarios:

Out[8]:

Visualization

Bar chart comparing peaked probability distribution under top-k and nucleus sampling, showing nucleus includes fewer tokens. — Peaked distribution: Model is confident. Top-k=5 includes many low-probability tokens (red), while nucleus sampling with p=0.9 focuses on just 2 tokens.

Bar chart comparing flat probability distribution under top-k and nucleus sampling, showing nucleus includes more tokens. — Flat distribution: Model is uncertain. Top-k=5 misses reasonable options, while nucleus sampling with p=0.9 includes 8 tokens to capture the spread probability mass.

These visualizations show why nucleus sampling outperforms top-k. When the model is confident (peaked distribution), nucleus sampling automatically tightens to just 2 tokens, while top-k=5 wastefully includes 3 near-zero tokens. When the model is uncertain (flat distribution), nucleus sampling expands to 8 tokens to capture the spread probability mass, while top-k=5 arbitrarily excludes 3 reasonable options.

Top-k makes the same decision regardless of context. Nucleus sampling reads the distribution and responds appropriately.

The left panel shows the peaked case: even though top-k=5 is marked (gray dashed line), nucleus sampling determines that only 2 tokens are needed. The right panel shows the flat case: nucleus sampling expands past the top-k=5 mark because the 6th, 7th, and 8th tokens still carry meaningful probability mass. This is exactly the behavior we want, and it happens automatically from a single shared hyperparameter $p$ .

Production Use with Hugging Face TransformersLink Copied

In practice, you won't implement nucleus sampling yourself. The transformers library provides it out of the box. Here's how to generate text with GPT-2 using top-p:

In[9]:

Code

from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.eval()

# Set pad token to avoid warnings
tokenizer.pad_token = tokenizer.eos_token

# Encode prompt
prompt = "The future of artificial intelligence is"
input_ids = tokenizer.encode(prompt, return_tensors="pt")

The generate() method accepts top_p directly. Set do_sample=True to enable sampling, and top_k=0 to disable top-k filtering so nucleus sampling operates alone:

In[10]:

Code

# Generate with nucleus sampling (top_p)
with torch.no_grad():
    outputs = model.generate(
        input_ids,
        max_new_tokens=50,
        do_sample=True,  # Enable sampling
        top_p=0.9,  # Nucleus probability threshold
        top_k=0,  # Disable top-k (set to 0 or very large)
        temperature=1.0,  # No temperature scaling
        pad_token_id=tokenizer.eos_token_id,
        num_return_sequences=3,  # Generate 3 different completions
    )

generated_texts = [
    tokenizer.decode(output, skip_special_tokens=True) for output in outputs
]

Out[11]:

Console

Generated completions with nucleus sampling (p=0.9):

[1] The future of artificial intelligence is a pivotal one. The basic characteristics of robots have not been fully developed for billions of years, so the consequences of human error could evolve in much the same way. Over the past 70 years, we have discovered dozens of glitches in the implementation of artificial

[2] The future of artificial intelligence is uncertain. For one thing, what happens when a computer only learns about features the human eye can't see? As AI develops, each individual person will be deprived of what little information their past activities might enable. As these generations experience more time on Earth

[3] The future of artificial intelligence is likely to become more complex."

What the future of AI will look like can be seen in more detail below:

Each completion takes a different path while remaining coherent. At every generation step, nucleus sampling includes only tokens the model considers plausible, allowing for creativity without introducing obvious errors.