Nucleus (top-p) sampling

Overview

A language model gives every token in its vocabulary a probability at each step. Nucleus sampling, also called top-p sampling, does not sample from all of them. It keeps the smallest set of most likely tokens whose probabilities add up to at least p, and draws the next token from that set [arxiv:2206.04624 §5.1]. A small p keeps the text close to the model's top choices; a large p lets it vary more. The text paper uses p = 0.9 as its top-p baseline 2206.04624 §4 2206.04624 §5.2.

The shared idea of every variant is the same: the set of allowed tokens is recomputed at each step from the model's own probabilities. Variants differ in the value of p at each step. The text paper's question is whether lowering p as a sentence goes on can reduce invented facts, given that top-p is reported to hurt factuality 2206.04624 §4 2206.04624 §5.1. Its card also covers a training change that is separate from sampling. Top-p also appears in protein generation, where it is one of two sampling knobs 2206.13517 §3.4. A third paper uses it only to draw a pool of candidate statements, which a separate scorer then ranks 2211.15006 app. C.1.

Variants

Top-p sampling (baseline)

A fixed p for the whole continuation. On the 1.3B model with p = 0.9, the paper reports a named-entity error of 52.4% on factual prompts and an entailment ratio of 2.9%, with diversity 0.88 and repetition 0.2% 2206.04624 tab. 4. Diversity is high, but the paper reports that factuality is poor 2206.04624 §4. Card: Factuality Enhanced Language Models for Open-Ended Text Generation.

Factual-nucleus sampling

The nucleus size for the t-th token of a sentence is

p_t = max(ω, p × λ^(t−1))

so it shrinks by a factor λ at each token and never goes below a floor ω. The schedule restarts at each sentence, detected when the previous token is a full stop 2206.04624 §5.1. The setting the paper carries into its training results is p = 0.9, λ = 0.9 and ω = 0.3 2206.04624 tab. 5. On the 1.3B model this gives a named-entity error of 42.1%, entailment of 10.1%, diversity of 0.55 and repetition of 7.1% 2206.04624 tab. 4. Other settings in the same table trade these numbers against each other: a stronger decay gives lower named-entity error at the cost of diversity and repetition 2206.04624 tab. 4. The paper also tests floors of ω = 0.2 and ω = 0.7 2206.04624 §5.2. A high floor keeps repetition low: with ω = 0.7 repetition is 1.2%, but named-entity error is 46.2%, close to top-p's 52.4% 2206.04624 tab. 4.

The three parts can be ablated separately. Adding decay alone (p = 0.9, λ = 0.9, no reset) leaves repetition at 30.7%. Adding the sentence reset to that decay brings repetition down to 10.3%, with named-entity error roughly unchanged at 41.1% and 41.5% 2206.04624 tab. 4. A stronger decay (λ = 0.5) with ω = 0.2 gives a named-entity error of 39.3%, but diversity drops to 0.38 2206.04624 tab. 4.

Factual-nucleus sampling with factuality-enhanced training

The paper also changes training: each sentence is prefixed with its article title, and the loss ignores the first half of the sentence [arxiv:2206.04624 §6.1] 2206.04624 §6.2. Decoding with factual-nucleus sampling is then applied to that model. On the 1.3B model, with both training changes, the named-entity error falls from 42.1% to 23.6% and entailment rises from 10.1% to 17.4%, while diversity falls to 0.38 2206.04624 tab. 5. At 530B the best reported row is 14.5% named-entity error and 25.5% entailment, against 33.3% and 11.8% for top-p on the vanilla model; diversity drops from 0.90 to 0.33 2206.04624 tab. 5. The 530B row combines training and decoding, so the paper's numbers do not separate their effects at that size.

Top-p with temperature in protein language models

ProGen2 samples protein sequences with a temperature and a top-p value set together. Its universal sampling grid is temperature {0.2, 0.4, 0.6, 0.8, 1.0} by top-p {0.5, 0.7, 0.9, 1.0}, and its fine-tuning grid uses top-p {0.7, 0.9, 1.0} 2206.13517 §3.4. The paper reports sweeps over these values, not one headline setting 2206.13517 §3.4. Its code default of top-p 0.95 is not one of the grid values sample.py.

Higher temperature or top-p gives more sequence diversity [arxiv:2206.13517 fig. 2]. For antibody samples, the paper reports lower aggregation propensity and higher solubility with less restrictive sampling across the temperature and top-p grid 2206.13517 §4.2 2206.13517 fig. 3. These are measured on the paper's generated sequences, not on lab-tested proteins. Card: ProGen2: Exploring the Boundaries of Protein Language Models.

Nucleus candidates with a reranker

Here nucleus sampling does not choose the output. It draws a pool of candidates, and a separate scorer picks one. The consensus paper draws 16 statements per group with cut-off 0.8 and temperature 1, reshuffling the group's opinions in each prompt 2211.15006 app. C.1. A reward model predicts each member's agreement with each candidate, and the candidate with the highest Utilitarian welfare is shown 2211.15006 §3.3. Reranking under three welfare functions showed no reliable difference in mean or minimum agreement 2211.15006 §5.1 2211.15006 app. D.5. The sampler is therefore one part of a pipeline that also includes fine-tuning. This page has no separate result for the sampling step alone. Card: Fine-tuning language models to find agreement among humans with diverse preferences.

When to use which

The first five rows measure factuality against Wikipedia, so they are about fewer unsupported named entities and entailed claims, not about truth in general. The sixth row concerns protein sequences and uses different metrics. The seventh row is about candidate pools for a reranker, where the paper's metrics are agreement ratings from human raters. Use the table only where your deployment has similar constraints.

The text paper runs most decoding and training ablations on 1.3B. Only the best setting is shown at 530B, so the ranking at other sizes is not established 2206.04624 §5.2 2206.04624 §6.3.

Implementation sketch

Language-agnostic pseudocode of the common core. The top-p rule is the same in every variant. The variants differ only in the value of p_t, which is fixed for top-p and scheduled for factual-nucleus.

# Nucleus sampling with the factual-nucleus schedule  [§5.1]
# Inputs: prompt tokens; p (0.9), lam (0.9), omega (0.3)
# With lam = 1 and omega = 0 this reduces to plain top-p  [§5.1]
# Output: generated tokens, ending at end-of-text or max_len (150)  [app. C]

def factual_nucleus_generate(model, prompt, p, lam, omega, max_len):
    tokens = list(prompt)
    # position of the next token inside the current sentence  [§5.1]
    position = 1
    while len(tokens) < max_len:
        # the nucleus shrinks with position, never below omega  [§5.1]
        p_t = max(omega, p * lam ** (position - 1))
        # keep the smallest set of tokens whose mass reaches p_t  [§5.1]
        probs = model.next_token_probs(tokens)  # (V,)
        next_token = sample_top_p(probs, p_t)
        tokens.append(next_token)
        if next_token == END_OF_TEXT:
            break
        # a full stop starts a new sentence, so the schedule restarts  [§5.1]
        position = 1 if next_token == FULL_STOP else position + 1
    return tokens

Notes on the sketch:

• The paper does not say how a full stop is identified in a subword vocabulary 2206.04624 §5.1. Check the token strings of your tokenizer before relying on this.
• The floor is reached early for strong decay. Computed from the formula 2206.04624 §5.1: with λ = 0.9 and ω = 0.3 it applies from about token 12 of each sentence; with λ = 0.5 from about token 3. For most of a sentence, λ = 0.5 therefore behaves like a fixed ω.
• The sampler the paper ran is not in its public repository. It points to a Megatron-LM nucleus sampler for this step [code: README.md#L88-L89] 2206.04624 §5.1.
• The protein paper's sampler is the plain top-p rule with a temperature applied first 2206.13517 §3.4. It has no per-step schedule.
• The consensus paper's candidate sampler is the plain top-p rule with p = 0.8 and temperature 1 2211.15006 app. C.1. Its public code defaults to temperature 0.8 and its API call passes no top-p value, so the nucleus step is not set in the released code Consensus paper card. Set p and temperature yourself when following that paper.

Contradictions

None between these three cards on nucleus sampling: they study different settings (factual text, protein sequences and consensus statements) and do not report the same claim. Within the text paper, §4 gives a 63.99% named-entity error for 126M, while its Table 3 gives 63.69% 2206.04624 §4 2206.04624 tab. 3. This page uses no value from that row, so the difference does not change any number here.

Papers

• Factuality Enhanced Language Models for Open-Ended Text Generation: defines factual-nucleus sampling, measures it against top-p and greedy decoding, and pairs it with training changes that lower named-entity error.
• ProGen2: Exploring the Boundaries of Protein Language Models: uses temperature and top-p sweeps to generate protein sequences and reports how diversity and developability change with them.
• Fine-tuning language models to find agreement among humans with diverse preferences: draws 16 nucleus samples per group and reranks them with a reward model and a welfare function; sampling is one step of a larger pipeline.