Top-k Sampling › The Problem with Full Vocabulary Sampling
Language AI HandbookAll chapters Top-k Sampling › The Problem with Full Vocabulary SamplingOne section · Sections in this chapter2 / 13 ↓
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
Before diving into top-k, it is worth spending time understanding exactly why sampling from the full vocabulary distribution causes problems. The issue is not immediately obvious, because any individual low-probability token should be selected very rarely. The problem only becomes clear when you think statistically about what happens over many generation steps.
A trained language model assigns a non-zero probability to every token in its vocabulary, even tokens that would be semantically absurd given the context. This is an unavoidable consequence of training: the model is trained on a finite corpus and learns a smooth probability distribution over the vocabulary. The smoothness prevents it from assigning exactly zero probability to any token, because that would create sharp discontinuities that generalize poorly. The result is a long tail of tokens with tiny but non-zero probabilities.
Think of full vocabulary sampling as a lottery where most tickets have a small but real chance of winning. In a vocabulary of 50,000 tokens, even if each "bad" token has probability 0.001%, collectively thousands of such tokens might represent 10-20% of the total probability mass. When you sample hundreds of tokens to generate a paragraph, the cumulative probability of at least one bad token slipping through becomes substantial. Just one incoherent token can derail the generation, because subsequent tokens are conditioned on the previous context, and an incoherent token shifts the distribution toward other incoherent tokens.
The key insight is that the model's probability distribution is shaped like a power law: a small number of tokens hold most of the probability mass, and the rest is distributed across a very long tail. The top 50 tokens might hold 95% of the probability, while the remaining 49,950 tokens share the other 5%. When we sample, we almost always get one of those top 50 tokens, but occasionally we get something from the long tail. That occasional bad draw is enough to degrade generation quality significantly, especially in longer sequences where errors can compound.
def get_next_token_distribution(model, tokenizer, prompt):
"""Get the probability distribution over next tokens."""
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
outputs = model(inputs["input_ids"])
logits = outputs.logits[0, -1, :]
probs = torch.softmax(logits, dim=0)
return probsPrompt: 'The capital of France is' Vocabulary size: 50,257 Tokens with P > 1%: 14 Tokens with P > 0.1%: 109 Top 5 predictions: 1. ' the': 0.0846 2. ' now': 0.0479 3. ' a': 0.0462 4. ' France': 0.0324 5. ' Paris': 0.0322
Notice how only a handful of tokens have probability above 1%, and only a few dozen have probability above 0.1%. The remaining 50,000+ tokens collectively hold a small but non-trivial share of the probability mass. This observation motivates truncation: if we only sample from the top few hundred tokens, we capture essentially all the probability that matters while guaranteeing we never produce true garbage.
The figure illustrates why the long tail is problematic. After the first few tokens, probability drops rapidly, but thousands of tokens retain some small probability. Summed together, this tail can be significant. If we sample proportionally, low-probability tokens occasionally win, creating outputs like "The capital of France is ????????" or "The capital of France is asdf". Even rarer, a sequence of low-probability choices can cascade into completely incoherent text that never recovers. Top-k sampling is our circuit breaker: once we identify the drop-off point in the probability ranking, we simply refuse to sample below it.