samankeon.com

From Greedy to Nucleus, How LLMs Choose the Next Token

· #llm

When I was implementing Gemma3-1b on MintEngine, I was struggling with what to use to select the token among the sea of logits coming out of the model. Just to resolve the moment’s need I resorted to use Argmax. Argmax selects the token with maximum probability. By chance my choice of token selection method was right for the task of numerical debugging of a language model and finally I got to write a report that I’m proud of: Gemma3, Architecture and Mathematical Foundations. But promised myself to get back to Sampling and do a deep dive on the subject. So here is what I learned.

Note: this article is written by human (me), might seem dense. Feel free to use AI for the concepts I didn’t explain clearly.

How decoding algorithm impacts the eval score. Image from “A Thorough Examination of Decoding Methods in the Era of LLMs”
How decoding algorithm impacts the eval score. Image from “A Thorough Examination of Decoding Methods in the Era of LLMs”

What I should call ‘em? Decoding or Sampling algorithms? Deterministic or stochastic?

I was in the camp of people who used the terms interchangeably. Did some search but didn’t get to any definitive definition of the terms in our context.

But using some unsupervised learning on myself, I got to find what’s what. All of these are part of decoding methods. Basically, they are used in decoders.

Then we have “deterministic” and “stochastic” methods that are under umbrella of decoder methods. Any algorithm that introduces randomness is called sampling or stochastic; others with low or no randomness go to the deterministic camp.

Deterministic algorithms

Deterministic algorithms are great for these two scenarios:

  • Accuracy: where we want to get most probably result. Greedy chooses the most probable result in the current token logprobs, and beam search does it over a sequence.

  • Consistency: when we want to generate same output, given the same input to the system. No surprises.

Greedy algorithm

In each generation step, Greedy algorithm chooses the token with maximum probability. It applies Argmax over the generated output (which has vocabsize length), and selects the top logprob.

Let the model define a conditional distribution over the vocabulary V:

Greedy decoding selects:

The generated sequence is:

with the constraint that the argmax is applied locally at each step, not globally over sequences.

Greedy vs Beam search. Image from “Greedy Search vs Beam Search Decoding: Concepts, Examples”
Greedy vs Beam search. Image from “Greedy Search vs Beam Search Decoding: Concepts, Examples”

It seems Beam search doesn’t have a clear birth date as in published through a paper. But found that folks started widely using it in speech recognition back in 1970s at IBM and Bell labs. In that time they called it Beam pruning, to limit the hypothesis space. And that’s the history, let’s go to practice.

Beam search is usually used with Beam width. In the above image beam width = 2, which means keep top 2 probable sequences. This doesn’t guarantee globally most probable sequence, but pretty much approximates it.

Stochastic algorithms

Stochastic algorithms introduce some variability to the system.

Temperature

Temperature is a simple but fundamental stochastic control applied to the model’s logits before sampling. It does not change the model, it reshapes the probability distribution used for token selection.

Let the model output logits be at step .

Standard softmax:

Temperature-scaled softmax introduces a scalar T>0:

Sampling is then performed from Pt​.

What temperature does mathematically

  • T=1:
    No change, original distribution.

  • T<1:
    Distribution becomes sharper. High-probability tokens dominate more strongly.

  • T>1:
    Distribution becomes flatter. Lower-probability tokens become more likely.

Two limits are worth noting:

This recovers greedy decoding.

This approaches pure randomness.

Typical usage ranges in practice

  • T=0.0 or omitted with greedy decoding

  • T=0.2to 0.5, focused and factual

  • T=0.7 to 1.0, balanced

  • T>1.0, creative but unstable

when implemented with code, it looks like this:

probs = torch.softmax(logits / temperature, dim=-1)

Top-K

top-k chooses the top tokens based on the probability and discards the rest. TopK can be applied after temperature step. That’s why when doing inference with vLLm or other inference engines you can set the topk alongside the temperature.

The high-level pipeline looks like this:

At generation step 𝑡:
- Model produces logits
- Apply temperature (optional)
- Convert to probabilities
- Apply top-k filter
- Renormalize probabilities
- Categorical sampling selects the token

So how we renormalize probabilities? Imagine we are doing Top2, selecting top 2 tokens with these probs: A=0.5, B=0.2. To renormalize, we divide each token prob by the sum of the probs, so:

Top-P (nucleus) sampling

Top-p keeps the smallest set of tokens whose cumulative probability mass is at least p, then samples from that set.

Unlike top-k, the size of the candidate set is adaptive.

Let tokens be sorted by descending probability:

Define the nucleus set Vp​ as the smallest prefix such that:

Define the truncated distribution:

For modern models, temperature + top-p is way to go. Temperature controls how strongly probabilities compete, while top-p controls which tokens are allowed to compete, and together they provide stable, high-quality sampling across a wide range of prompts and model sizes.

Homework

Everything we talked in the stochastic algorithms were about changing the probability and leaving some k tokens. We didn’t talk about how to exactly choose these tokens. The method is called categorical sampling, but I leave it to you to read about it. :)

Outro

I’ve covered most frequently used decoding methods. Both deterministic and stochastic. Looking at any inference engine again, you won’t be lost setting the parameters.

Looking at the image I shared at the beginning of the article, you can see there is no size fits all. Some evals, respond better to specific decoding methods.

How decoding algorithm impacts the eval score. Image from “A Thorough Examination of Decoding Methods in the Era of LLMs”
How decoding algorithm impacts the eval score. Image from “A Thorough Examination of Decoding Methods in the Era of LLMs”