Top-k Sampling: Why a Fixed Shortlist Ignores How Sure the Model Is?

A language model (the AI that writes text one word at a time) scores every possible next word, then picks one at random, weighted by those scores. Each score becomes a percentage chance, and the chances add up to 100%. Many systems first cut that list down to a shortlist and pick only from the shortlist. The simplest cut is top-k sampling: keep the k most likely words, throw away everything else, and divide each kept chance by the kept total so the kept chances add up to 100% again.

Top-k is popular for a good reason: the shortlist is bounded and predictable. But the same rule that makes it predictable also makes it blind. The k is a number you choose ahead of time. A k of 5 means 5 words whether the model is nearly certain or nearly guessing. This post shows what that blindness does, using a made-up list of 12 possible words and real, rerunnable code.

The Rule

Top-k takes the model's chances for each word, sorts them from highest to lowest, keeps the first k, drops the rest, and rescales the kept chances so they add up to 100% again. Why sort at all? Because the rule needs to know which words rank highest, not their exact scores. The rule never looks at how spread out the chances are. It only looks at rank.

The small made-up example asks one question: "How many days do I have to recover a deleted file?" The correct answer is 30. The made-up list has 12 possible words. It includes wrong answers that sound reasonable (14 and 90) and three words that don't answer the question at all (immediately, never, unsure). The raw scores (called logits) are hand-picked; every chance and percentage below comes from running the code.

The Model's Own Chances: Five Words for One Clear Answer

Here the model's own chances are used as they are (T=1.0, the default setting). The correct answer, 30, gets 53.12% of the kept chances after rescaling. Top-k with k=5 keeps five words: 14 (23.87%), 90 (16.00%), 60 (4.36%), and 45 (2.64%). These five shares add up to 100%. The last two together are about 7%, yet the model can still pick them, even though they are unlikely.

Why does that matter? Because the cutoff never asked whether the model was confident. It counted to five, and it would count to five on any prompt.

The Same k With Flatter Chances

At T=1.0, the same k=5 keeps the same five words and throws away about 5% of the chances. Now spread the chances out with T=2.0, a setting that flattens them. The same k=5 still keeps 30, 14, 90, 60, and 45, but now throws away about 21%. The same code ran; only that setting changed. The amount thrown away grew more than four times.

Is that a failure? Not in this example. At T=2.0, the seven discarded words are 7, 10, 365, 0, and the three words that don't answer the question (immediately, never, unsure). So top-k removed all three non-answers. That is the honest result for this example. Had the correct answer landed in the discarded part, a fixed cutoff would have removed it. This example does not show that case. Top-k cannot tell which situation it is in, because it never looks at the chances it throws away.

What Top-k Is Good For

Doesn't a fixed cutoff just make results worse? Not always. The upside is predictable behavior. The model only considers k words, so words ranked below k can never be picked, however unlikely they are. For a real model with tens of thousands of tokens (word fragments), that is a practical safety limit.

The price is that k is one setting for every word the model writes. A value that feels right for an open-ended story may be far too wide for a factual answer, and a value tuned for facts may choke a creative prompt. The number has to be picked for the task, not for the chances.

Get the Code

The softmax function (it turns raw scores into chances that add up to 100%) and the top-k cut, with the T=1.0 and T=2.0 runs:

# Top-k sampling -- a fixed candidate count that can't see how confident the model is.
import numpy as np

np.random.seed(42)

VOCAB = ["30", "14", "90", "60", "45", "7", "10", "365", "0",
         "immediately", "never", "unsure"]

# Raw logits for: "How many days do I have to recover a deleted file?"
LOGITS = np.array([
    4.0, 3.2, 2.8, 1.5, 1.0, 0.6, 0.4, -0.5, -0.8, -1.0, -1.5, -2.0,
])

WRONG_TOKENS = ["immediately", "never", "unsure"]
K = 5


def softmax(logits):
    shifted = logits - np.max(logits)
    exp = np.exp(shifted)
    return exp / exp.sum()


def apply_temperature(logits, T):
    return softmax(logits / T)


def top_k_filter(probs, k):
    idx_sorted = np.argsort(probs)[::-1]
    kept_idx = idx_sorted[:k]
    discarded_mass = probs[idx_sorted[k:]].sum()
    kept_probs = probs[kept_idx]
    return kept_idx, kept_probs / kept_probs.sum(), discarded_mass


probs_T10 = apply_temperature(LOGITS, 1.0)
probs_T20 = apply_temperature(LOGITS, 2.0)

kept_T10, renorm_T10, discarded_T10 = top_k_filter(probs_T10, K)
kept_T20, _, discarded_T20 = top_k_filter(probs_T20, K)

print(f"=== top_k={K} at T=1.0 ===")
print(f"keeps: {[VOCAB[i] for i in kept_T10]}")
print(f"renormalized share of kept tokens: {[f'{p:.4%}' for p in renorm_T10]}")
print(f"discarded probability mass: {discarded_T10:.4%}")

print(f"\n=== top_k={K} at T=2.0 ===")
print(f"keeps: {[VOCAB[i] for i in kept_T20]}")
print(f"discarded probability mass: {discarded_T20:.4%}")

wrong_idx = {VOCAB.index(t) for t in WRONG_TOKENS}
ungrounded_kept_T20 = [VOCAB[i] for i in kept_T20 if i in wrong_idx]
print(f"\nungrounded tokens still kept at T=2.0: {ungrounded_kept_T20 or 'none'}")
discarded_T20_words = [VOCAB[i] for i in range(len(VOCAB)) if i not in set(kept_T20)]
print(f"discarded at T=2.0: {discarded_T20_words}")
print(f">> Same k={K}, same code path: discarded mass goes from "
      f"{discarded_T10:.2%} to {discarded_T20:.2%}.")

Running it prints:

=== top_k=5 at T=1.0 ===
keeps: ['30', '14', '90', '60', '45']
renormalized share of kept tokens: ['53.1238%', '23.8701%', '16.0006%', '4.3607%', '2.6449%']
discarded probability mass: 4.7243%

=== top_k=5 at T=2.0 ===
keeps: ['30', '14', '90', '60', '45']
discarded probability mass: 21.3310%

ungrounded tokens still kept at T=2.0: none
discarded at T=2.0: ['7', '10', '365', '0', 'immediately', 'never', 'unsure']
>> Same k=5, same code path: discarded mass goes from 4.72% to 21.33%.

Summary

Top-k keeps the k most likely words and drops the rest. Its strength is a bounded, predictable shortlist. Its weakness is that k is fixed. On the same made-up list of 12 words, k=5 kept a word holding 2.64% of the kept chances at T=1.0. At T=2.0 it threw away about 21% of all chances, and in this example that discarded 21% was all non-answers. The risk is a case where the right answer gets discarded, which this example does not test. Top-k is a useful safety limit, but you set k once for each task, and it stays the same for every word the model writes.

Should you drop top-k? Not necessarily. It is a cheap, predictable safety limit. The trade-off is that its cutoff never adapts to how sure the model is, so the k you choose has to fit the task.

Comments

Post a Comment

Popular posts from this blog

All about READ in RPGLE & Why we use it with SETLL/SETGT?

READ READ is one of the most used Opcodes in RPGLE. As the name suggests main purpose of this Opcode is to read a record from Database file. What are the different READ Opcodes? To list, Below are the five Opcodes.  READ - Read a Record READC - Read Next Changed Record READE - Read Equal Key Record READP - Read Prior Record READPE - Read Prior Equal Record We will see more about each of these later in this article. Before that, We will see a bit about SETLL/SETGT .  SETLL (Set Lower Limit) SETLL accepts Key Fields or Relative Record Number (RRN) as Search Arguments and positions the file at the Corresponding Record (or Next Record if exact match isn't found).  SETGT (Set Greater Than) SETGT accepts Key Fields or Relative Record Number (RRN) as Search Arguments and positions the file at the Next Record (Greater Than the Key value). Syntax: SETLL SEARCH-ARGUMENTS/KEYFIELDS FILENAME SETGT  SEARCH-ARGUMENTS/KEYFIELDS FILENAME One of the below can be passed as Search...
Read more

Why LoRA Works: The SVD Math Behind Efficient Fine-Tuning

Image
We keep talking about Gradient Descent when it comes to training a machine learning model. We calculate the gradient, adjust weights to minimize the loss. That sounds simple enough. But what does it mean to train a large language model with billions of parameters? Gradient descent still works exactly as it should — you compute gradients with respect to every single parameter and store optimizer states for all of them. What's the problem then? It works perfectly, but requires a lot of compute power and tens of gigabytes just to hold the training state. It requires GPUs, possibly a connection to a data center. What if you are a student or a researcher with just a laptop? Do you give up on fine-tuning a model with billions of parameters? What if you could train a model where, instead of modifying billions of weights, you only tune a few parameters — and still get nearly the same result? The answer is Low Rank Approximation . And to understand it, ...
Read more

Greedy vs. Sampling: Why the Same Prompt Can Give You a Different Answer Every Time

Image
Ask an LLM (large language model — the kind of AI behind tools like ChatGPT) the exact same question twice. Sometimes you get the exact same answer back, word for word. Other times you get two answers that take different angles on the same question, worded completely differently. Both are normal — and both come down to one setting (a parameter passed into the code — not to be confused with a model's parameters, the separate term for its internal size), hiding in plain sight in the same line of code everyone copies from a tutorial: outputs = model.generate( **inputs, max_new_tokens=40, # caps the reply at 40 tokens (word pieces) do_sample=False, ) That's a simplified version of the standard Hugging Face (a popular platform and code library for running AI models) transformers convention for generating text — the full version used for every result below pins a couple of extra settings for consistency, shown later. do_sample decides which of two en...
Read more