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

> Discover the implementation differences between nucleus sampling and top-k sampling for language model decoding. Learn how nucleus sampling adapts to probability distributions for better text generation.

- Repository: [labml.ai/annotated_deep_learning_paper_implementations](https://github.com/labmlai/annotated_deep_learning_paper_implementations)
- Tags: deep-dive
- Published: 2026-03-04

---

**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`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/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`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/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`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/sampling/__init__.py) and can be chained with other components like temperature scaling.

```python
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`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/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`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/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`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/sampling/top_k.py) and `NucleusSampler` in [`labml_nn/sampling/nucleus.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/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.