← Ashish's notebook/How Top-K Sampling Works in LLMs

Notebook

How Top-K Sampling Works in LLMs

The model ranks thousands of words to pick just one. Top-K keeps a few and drops the rest.

07 SEP 2026

In the last notebook we learned about temperature.

Temperature changes the probabilities of the possible next tokens which can change the personality and output of the model.

Now we will learn Top-K. What is Top K? Top-K is a sampling technique of LLM which is done on the probabilities of the next token. Basically we can choose how many of the highest-probability tokens the model should consider for the next token.

You will understand better by seeing it, so let's get right into it.

We will take the model apart again and actually see this happen.

Input → Tokenization → Model → Logits → Softmax → Probabilities → Top-K filtering → Sampling → Next token

What we will learn

1. Where Top-K Fits in Text Generation
2. What K means
3. How to Get the Model's Raw Logits
4. How Softmax Converts Logits into Probabilities
5. How to See the Top Candidate Tokens
6. How Renormalization Changes the Probabilities
7. How Different K Values Change the Output
8. How Different K Values Change the Generated Text
9. Problems with Top-K sampling.

Setup — Load the Model and Tokenizer

Same two steps as before: the tokenizer, which turns text into numbers, and the model itself.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "Qwen/Qwen3-0.6B"

# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
)

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

Loading weights:   0%|          | 0/311 [00:00<?, ?it/s]

Let's take top_k=5 for now

from transformers import GenerationConfig

model.generation_config = GenerationConfig(
    do_sample=True,
    temperature=0.3,    
    top_k=5,            #"k" is the integer representing how many of the highest-probability tokens to keep
    max_new_tokens=10,
)

Now we write the prompt and print the output of next predictied words.

return_tensors="pt" tells the tokenizer to return PyTorch tensors (vectors) instead of Python lists.

There are some special tokens in llms. For example - "<|endoftext|>", it tells the llm the end of text. We don't want these in the output , that's why we use skip_special_tokens=True it skips all the special tokens from the output.

prompt = "Can you bring me"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

output = model.generate(**inputs)


print(tokenizer.decode(output[0], skip_special_tokens=True))

Can you bring me a list of all the possible numbers that are divisible

Chapter 1 — Get the logits

Before Top-K can happen, the model first produces logits.

The logits contain one score for every token in the vocabulary.

import torch

with torch.no_grad():
    outputs = model(**inputs)

logits = outputs.logits[0, -1, :]

print(logits.shape)

torch.Size([151936])

Chapter 2 — Convert logits to probabilities

Softmax converts the logits into probabilities.

Now every token has a probability between 0 and 1, and all probabilities add up to 1.

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

print(probs.shape)

torch.Size([151936])

Chapter 3 — Visualize the Top-K tokens

This is the graph we want to see.

Let's set top_k to 10, it should now give us top 10 probabilities of the next output.

import matplotlib.pyplot as plt

k = 10

topk_probs, topk_idx = torch.topk(probs, k)

tokens = [
    tokenizer.decode([idx.item()]).strip() or "␣"
    for idx in topk_idx
]

plt.figure(figsize=(10, 5))

plt.bar(
    tokens,
    topk_probs.float().detach().cpu().numpy()
)

plt.xlabel("Candidate next token")
plt.ylabel("Probability")
plt.title(f"Top-{k} tokens for: {prompt!r}")
plt.xticks(rotation=45)

plt.show()

Now lets try by setting top_k to 20

import matplotlib.pyplot as plt

k = 20

topk_probs, topk_idx = torch.topk(probs, k)

tokens = [
    tokenizer.decode([idx.item()]).strip() or "␣"
    for idx in topk_idx
]

plt.figure(figsize=(10, 5))

plt.bar(
    tokens,
    topk_probs.float().detach().cpu().numpy()
)

plt.xlabel("Candidate next token")
plt.ylabel("Probability")
plt.title(f"Top-{k} tokens for: {prompt!r}")
plt.xticks(rotation=45)

plt.show()

Chapter 4 — Original vs Top-K

This part is important. We will learn about renormalization which is an important part.

Imagine the original probabilities of the Top-5 tokens are:

0.40
0.25
0.15
0.10
0.05

Their total is:

0.95

Top-K removes the other tokens.

For sampling, the remaining probabilities are renormalized so that they add up to 1:

0.40 / 0.95
0.25 / 0.95
0.15 / 0.95
0.10 / 0.95
0.05 / 0.95

So:

Top-K picks the tokens. Renormalization gives those tokens their new probabilities.

k = 5

topk_probs, topk_idx = torch.topk(probs, k)

renormalized_probs = topk_probs / topk_probs.sum()

tokens = [
    tokenizer.decode([idx.item()]).strip() or "␣"
    for idx in topk_idx
]

for token, original, new in zip(
    tokens,
    topk_probs,
    renormalized_probs
):
    print(
        f"{token!r:20} "
        f"original={original.item():.4f}   "
        f"after Top-K={new.item():.4f}"
    )

print("Sum after renormalization:", renormalized_probs.sum().item())

'a'                  original=0.1895   after Top-K=0.3066
'the'                original=0.1670   after Top-K=0.2715
'up'                 original=0.1152   after Top-K=0.1865
'some'               original=0.1016   after Top-K=0.1650
'help'               original=0.0422   after Top-K=0.0684
Sum after renormalization: 1.0

Chapter 5 — Compare different values of K

Now let's actually see what changes when we use:

K = 1
K = 5
K = 10
K = 20

Remember:

• K=1 → only one candidate
• small K → narrow sampling pool (safer output)
• large K → larger sampling pool (more variety)

ks = [1, 5, 10, 20]

fig, axes = plt.subplots(
    1,
    len(ks),
    figsize=(20, 4),
    sharey=True
)

for ax, k in zip(axes, ks):

    topk_probs, topk_idx = torch.topk(probs, k)

    tokens = [
        tokenizer.decode([idx.item()]).strip() or "␣"
        for idx in topk_idx
    ]

    ax.bar(
        range(k),
        topk_probs.float().detach().cpu().numpy()
    )

    ax.set_title(f"K = {k}")
    ax.set_xticks(range(k))
    ax.set_xticklabels(
        tokens,
        rotation=90,
        fontsize=8
    )

    ax.set_xlabel("Token")

axes[0].set_ylabel("Probability")

plt.tight_layout()
plt.show()

Chapter 6 — Output generation with different values of k.

Here we are setting different values of k and seeing how the output changes based on that

from transformers import GenerationConfig

for k in [1, 5, 20]:

    model.generation_config = GenerationConfig(
        do_sample=True,
        temperature=0.3,
        top_k=k,
        max_new_tokens=10,
    )

    output = model.generate(**inputs)

    text = tokenizer.decode(
        output[0],
        skip_special_tokens=True
    )

    print(f"K={k}: {text}")

K=1: Can you bring me a list of all the possible combinations of 3
K=5: Can you bring me some information on the history of the United States Constitution
K=20: Can you bring me a list of all the possible combinations of the numbers

Chapter 7 — Problems with Top-K sampling.

The main problem with Top-K sampling is that K is a fixed number. It doesn't matter how confident the model is already.

For example: The probabilites are

A: 0.90
B: 0.04
C: 0.02
D: 0.01
E: 0.01
etc

If K = 5, Top-K keeps A–E. But most of the probability mass is concentrated on A, so others becomes nearly extra noise.

Now the question is When can we use K sampling?.And the answer is when -

1. You want to remove very unlikely tokens.
2. You want predictable control over randomness.
3. You want a simple way to control sampling.

This problem is exactly what Top-P (nucleus) sampling fixes. Instead of keeping a fixed number of tokens, it keeps as many tokens as needed to reach a certain probability mass. So when the model is very confident it keeps fewer tokens, and when it is unsure it keeps more.

We will learn about Top-P in the next one.

Summary

Top-K is a sampling method that limits the model to the K most likely next tokens and removes the rest. We went through the process step by step: the model first produces logits, softmax converts them into probabilities, and torch.topk selects the top K tokens. Since the remaining probabilities no longer add up to 1, they are renormalized so the model can sample from them properly.

We also saw how different K values affect the output. A small K, such as K=1, makes the output more predictable, while a larger K gives the model more choices and more variety. Comparing K=1, 5, 10, and 20 showed how K controls the number of possible tokens the model can choose from at each step.

← All notebooks

GitHub LinkedIn