Top-K

How This Lesson Fits the Module & Volume

Pure sampling—even with temperature—can still draw from a long low-probability tail. Top-k truncates the candidate set to the k highest-scoring tokens before renormalizing and sampling.

It is the fixed-count cousin of top-p (nucleus) sampling. Together they are the workhorse filters in modern LM inference APIs.

Learning Objectives

By the end of this lesson, students should be able to:

Define top-k sampling and the truncate–renormalize–sample pipeline.
Implement top-k filtering on logits in PyTorch.
Choose k relative to vocabulary size and task risk.
Contrast top-k with temperature-only and with top-p.
Explain failure modes when k is too small or too large.
Combine temperature and top-k in one decode step.

Definition

Top-k sampling keeps only the k tokens with the largest logits (or probabilities), sets all other logits to −∞, applies softmax to renormalize over the kept set, then samples. The support size is fixed at k every step (or fewer if V < k).

Algorithm Steps

Score

Compute logits

Select

Keep top k

Mask

Others → −∞

Sample

Softmax + draw

k	Behavior	Risk
k = 1	Equivalent to greedy	No diversity
k = 10–50	Common creative defaults	May still miss mass if flat
k = 100+	Closer to full sampling	Tail noise returns
k = V	No truncation	Same as pure sampling

Top-K vs Top-P

Fixed candidate count.
Simple to reason about.
Ignores how peaked p is.

Top-p

Adaptive set by mass.
Small set when peaked.
Larger set when uncertain.

Both

Often combined in APIs.
Apply after temperature.
Then multinomial sample.

Code: Top-K Filter + Sample

import torch
import torch.nn.functional as F

def top_k_logits(logits, k):
    # logits: (B, V)
    if k <= 0 or k >= logits.size(-1):
        return logits
    values, _ = torch.topk(logits, k)
    min_keep = values[:, -1].unsqueeze(-1)
    return torch.where(logits < min_keep, torch.full_like(logits, float("-inf")), logits)

def sample_top_k(logits, k=50, temperature=1.0):
    if temperature > 0:
        logits = logits / temperature
    filtered = top_k_logits(logits, k)
    probs = F.softmax(filtered, dim=-1)
    return torch.multinomial(probs, num_samples=1)

logits = torch.randn(1, 1000)
torch.manual_seed(0)
print(sample_top_k(logits, k=10, temperature=0.8))

Strengths and Tradeoffs

Strengths

Hard ceiling on how weird a token can be.
Easy hyperparameter with clear meaning.
Cheap with torch.topk.

Tradeoffs

Fixed k is awkward when the distribution is very peaked or very flat.
Too-small k kills legitimate diversity.
Does not adapt per step like top-p.

Common Misconception

“Top-k returns the k tokens as the output text.” Top-k only restricts the candidate distribution for one generation step. You still sample (or argmax) a single next token, append it, and repeat.

Related module pages: Temperature, Top-P, Sampling, Logits, Inference.

Knowledge Check

Short Answer: What happens to tokens outside the top k? Answer: Their logits are set to −∞ (probability 0 after softmax).
True/False: Top-k with k = 1 is greedy decoding. Answer: True.
Multiple Choice: After masking, you must: (a) renormalize via softmax, (b) skip softmax, (c) delete the model. Answer: (a).
Short Answer: Name one weakness of a fixed k. Answer: It does not adapt when the distribution is peaked vs flat.
True/False: Top-k outputs k tokens of final text each step. Answer: False—still one next token.
Multiple Choice: Typical order is: (a) truncate then temperature, (b) temperature then truncate, (c) beam then dropout. Answer: (b) is common (T then filter).
Short Answer: Which method keeps tokens until cumulative probability ≥ p? Answer: Top-p / nucleus sampling.
True/False: torch.topk helps implement the filter. Answer: True.
Multiple Choice: Very large k behaves like: (a) pure sampling, (b) beam search, (c) tokenization. Answer: (a).
Short Answer: What lecture covers adaptive nucleus truncation? Answer: Top-P.

Key Takeaways

Top-k keeps only the k best logits, then samples.
It blocks the long tail with a fixed candidate budget.
k = 1 is greedy; k = V disables truncation.
Often paired with temperature; compare with top-p.
Next: Top-P.

Trainer’s Guide

Hands-on idea: On a peaked and a flat logit vector, print how much probability mass the top-10 tokens cover; discuss why fixed k feels different in each case.

Discussion prompt: For code completion, would you rather tighten k or lower temperature first?

Recap: Top-k truncates to a fixed shortlist before sampling the next token. Continue with Top-P.