Nucleus Sampling vs Top-k Sampling for Language Model Decoding: Implementation Differences

Nucleus sampling dynamically selects the smallest token set whose cumulative probability exceeds a threshold p, while top-k sampling selects exactly k tokens with the highest individual probabilities regardless of their cumulative mass, making nucleus sampling more adaptive to varying probability distributions.

Language model decoding strategies determine the quality and diversity of generated text. In the labmlai/annotated_deep_learning_paper_implementations repository, both nucleus (top-p) and top-k sampling are implemented as modular components in the labml_nn/sampling/ package that filter logits before stochastic selection.

How Candidate Token Sets Are Defined

The fundamental difference between these methods lies in how they construct the candidate vocabulary for each decoding step.

Top-k Sampling: Fixed-Size Ranking

Top-k sampling restricts generation to the k highest-probability tokens based on their individual ranks. According to the implementation in labml_nn/sampling/top_k.py, the TopKSampler class uses torch.topk to identify the k largest logits, then masks all other positions to -inf before passing the filtered logits to an inner sampler.

This approach maintains a fixed candidate set size regardless of the probability distribution's shape. If the distribution is sharply peaked, top-k may include low-probability tail tokens; if flat, it may exclude viable candidates.

Nucleus Sampling: Dynamic Probability Mass

Nucleus sampling (also called top-p sampling) selects the smallest set of tokens whose cumulative probability exceeds a threshold p. As implemented in labml_nn/sampling/nucleus.py, the NucleusSampler first applies softmax to obtain probabilities, sorts them in descending order, and computes a cumulative sum. It then builds a boolean mask where the cumulative sum is less than p, ensuring at least one token is always retained to avoid empty sets.

This creates a dynamic candidate set that expands or contracts based on the model's confidence. For peaked distributions, nucleus sampling may select only 2-3 tokens; for uniform distributions, it may include hundreds.

Implementation Comparison in PyTorch

Both samplers inherit from the abstract Sampler class defined in labml_nn/sampling/__init__.py and can be chained with other components like temperature scaling.

import torch
from labml_nn.sampling import Sampler
from labml_nn.sampling.temperature import TemperatureSampler
from labml_nn.sampling.top_k import TopKSampler
from labml_nn.sampling.nucleus import NucleusSampler

# Example logits from a language model

logits = torch.randn(1, 10000)

# Top-k sampling: exactly 50 candidates

topk_sampler = TopKSampler(k=50, sampler=TemperatureSampler(temperature=0.8))
topk_token = topk_sampler(logits)

# Nucleus sampling: candidates covering 95% probability mass

nucleus_sampler = NucleusSampler(p=0.95, sampler=TemperatureSampler(temperature=0.8))
nucleus_token = nucleus_sampler(logits)

In labml_nn/sampling/top_k.py, the TopKSampler.__call__ method uses torch.topk(logits, self.k) to extract indices, then scatters these into a -inf initialized tensor. Conversely, labml_nn/sampling/nucleus.py computes torch.cumsum(torch.sort(softmax(logits)) to determine the dynamic cutoff, handling edge cases where the cumulative sum exactly equals p by ensuring at least one token remains.

When to Use Each Method

Use top-k sampling when you need predictable computational costs or want to strictly limit vocabulary diversity to a fixed number of candidates. This is useful in constrained generation tasks where you want to exclude all but the most likely k options regardless of probability mass distribution.

Use nucleus sampling when you want adaptive quality control that responds to the model's confidence. Because nucleus sampling guarantees retaining probability mass p, it automatically adjusts to peaked distributions (fewer candidates) and flat distributions (more candidates), typically producing higher-quality, less degenerate text than fixed-size alternatives.

Summary

• Top-k sampling selects exactly k tokens with the highest individual probabilities, creating a fixed-size candidate set that may exclude significant probability mass or include low-quality tail tokens depending on the distribution shape.
• Nucleus sampling selects the smallest token set whose cumulative probability exceeds p, dynamically adapting the candidate set size to the distribution's peakedness while guaranteeing retention of at least p probability mass.
• Both methods are implemented in the labmlai/annotated_deep_learning_paper_implementations repository as modular samplers (TopKSampler in labml_nn/sampling/top_k.py and NucleusSampler in labml_nn/sampling/nucleus.py) that mask logits to -inf before delegating to an inner sampler.

Frequently Asked Questions

What is the main advantage of nucleus sampling over top-k sampling?

The main advantage is adaptability. Nucleus sampling adjusts the candidate set size based on the probability distribution's shape, ensuring high-probability mass coverage without manually tuning k for different contexts. When the model is confident (peaked distribution), nucleus sampling automatically uses fewer candidates; when uncertain (flat distribution), it includes more, whereas top-k sampling would use the same fixed number regardless.

Can top-k and nucleus sampling be used together?

Yes, they can be combined sequentially. In the labmlai/annotated_deep_learning_paper_implementations repository, both TopKSampler and NucleusSampler accept an inner sampler parameter, allowing you to chain them. For example, you could first apply top-k to limit to 100 candidates, then apply nucleus sampling with p=0.9 to further filter based on cumulative probability, before finally applying temperature scaling.

Why do both methods mask logits to negative infinity instead of zeroing probabilities?

Masking logits to -inf (using torch.full initialized to -inf and scattering valid logits) ensures that after the softmax operation in downstream samplers, the probability of masked tokens becomes exactly zero. If probabilities were zeroed directly after softmax, the remaining probabilities would need renormalization. By operating on logits before softmax, the implementation maintains numerical stability and allows seamless composition with other logit-based transformations like temperature scaling.

Which hyper-parameter values are typically used for nucleus and top-k sampling?

For top-k sampling, common values range from k=40 to k=100, with k=50 being a frequent default in many language generation pipelines. For nucleus sampling, typical p values range from 0.9 to 0.95, where p=0.9 retains the smallest set covering 90% of probability mass. Lower values (e.g., p=0.5) produce more focused, deterministic output, while higher values (e.g., p=0.99) allow greater diversity but may include more low-quality tokens.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}