Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Top-k Sampling › Implementing Top-k Sampling

Michael BrenndoerferUpdated July 29, 20258 min read

Language AI HandbookAll chapters

Part 28 of 67GPT Architecture

Chapter 203Top-k Sampling › Implementing Top-k Sampling

Reading section 5 of 13 · Implementing Top-k Sampling · 8 minSection 5 of 13 · 8 min

Section 5 of 13 · 8 min read

Top-k Sampling › Implementing Top-k SamplingOne section · Sections in this chapter5 / 13 ↓

Sections in this chapter13

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

With the mathematics established, let us translate the algorithm into code. The implementation is straightforward, but walking through it step by step reveals how each line corresponds to a piece of the formula.

The implementation relies on two key PyTorch operations. First, torch.topk(tensor, k) efficiently finds the $k$ largest values in a tensor and returns both those values and their original positions (indices) in the tensor. We need both the values (to compute probabilities) and the indices (to identify which vocabulary tokens were selected). Second, torch.multinomial(probs, num_samples=1) draws a random sample from a discrete probability distribution, returning the index of the sampled element.

Think of the implementation as a pipeline: raw logits flow in one end, get temperature-scaled, get truncated to the top $k$ , get converted to probabilities via softmax, and then a single sample emerges from the other end. Each stage is a simple operation, and the pipeline is easy to read and debug.

Building the Core SamplerLink Copied

Our implementation needs to accomplish four things:

Apply temperature scaling to the logits
Find the $k$ highest-scoring tokens
Convert those scores to a valid probability distribution
Sample one token from that distribution

The key insight is that we do not need to explicitly zero out the non-top-k tokens. Instead, we can compute softmax over only the top-k logits, which automatically gives us a properly normalized distribution over just those tokens. This is equivalent to the renormalization step in our formula and is numerically cleaner than setting probabilities to zero and then renormalizing.

In[9]:

Code

def top_k_sample(logits, k, temperature=1.0):
    """
    Sample from top-k truncated distribution.

    Args:
        logits: Raw logits from model, shape (vocab_size,)
        k: Number of top tokens to keep
        temperature: Temperature for scaling before truncation

    Returns:
        Sampled token index
    """
    # Apply temperature scaling
    scaled_logits = logits / temperature

    # Get top-k values and indices
    top_k_values, top_k_indices = torch.topk(scaled_logits, k)

    # Convert to probabilities
    top_k_probs = torch.softmax(top_k_values, dim=0)

    # Sample from truncated distribution
    sample_idx = torch.multinomial(top_k_probs, num_samples=1)

    # Map back to original vocabulary index
    return top_k_indices[sample_idx].item()

Let's trace through each step:

Temperature scaling divides each logit by the temperature, controlling how peaked or flat the distribution becomes before truncation. Dividing by a value less than 1 sharpens the distribution; dividing by a value greater than 1 flattens it.
torch.topk efficiently finds the $k$ largest values and their positions in the vocabulary, returning both the values and their indices. The indices tell us which vocabulary tokens these are.
Softmax over top-k values computes $e^{z_i} / \sum_j e^{z_j}$ using only the kept tokens, which is equivalent to the renormalization step in our formula. We get a valid probability distribution automatically.
torch.multinomial samples according to the probability distribution, returning an index into our top-k list (a number from 0 to $k-1$ , not the vocabulary index).
Index mapping converts the sampled position (0 to $k-1$ ) back to the actual vocabulary token ID using top_k_indices[sample_idx].

Seeing It in ActionLink Copied

Let's sample multiple times from the same context to see the variety top-k produces:

Out[10]:

Console

Prompt: 'The capital of France is'

Sampling with k=5 (10 samples):
  Sample 1: ' the'
  Sample 2: ' France'
  Sample 3: ' the'
  Sample 4: ' now'
  Sample 5: ' the'
  Sample 6: ' the'
  Sample 7: ' the'
  Sample 8: ' a'
  Sample 9: ' now'
  Sample 10: ' a'

Notice how all samples are reasonable completions. With $k=5$ , we are sampling from only the five most likely tokens, which for this factual prompt are all sensible choices. The variation comes from the probabilistic sampling, but every option is plausible. If we had used $k=1$ , we would always get the same token; if we had used $k=50000$ , we might occasionally get something bizarre.

The output also demonstrates one of the practical advantages of top-k sampling for interactive applications: you can run the same prompt multiple times and get different, interesting responses, unlike greedy decoding which always gives you the same deterministic output.

From Single Tokens to Full Text GenerationLink Copied

A single sample is interesting, but the real power of top-k sampling emerges when we generate entire sequences. Each token we generate becomes part of the context for the next token, creating a chain of sampling decisions. The model sees the prompt plus all previously generated tokens, creating a new distribution, from which we sample again. This auto-regressive loop continues until we hit a stopping condition.

The generation loop follows a simple pattern: get logits, sample a token, append it to the sequence, repeat. At each step, top-k ensures we only consider reasonable continuations. The accumulated context shapes which tokens are plausible at each new step, so the generated text remains coherent even as it explores different directions.

In[11]:

Code

def generate_with_top_k(
    model, tokenizer, prompt, max_tokens=30, k=50, temperature=1.0
):
    """Generate text using top-k sampling."""
    input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
    generated_ids = input_ids.clone()

    for _ in range(max_tokens):
        with torch.no_grad():
            outputs = model(generated_ids)
            next_token_logits = outputs.logits[0, -1, :]

        # Sample using top-k
        next_token_id = top_k_sample(
            next_token_logits, k=k, temperature=temperature
        )

        # Append and continue
        generated_ids = torch.cat(
            [generated_ids, torch.tensor([[next_token_id]])], dim=1
        )

        # Stop if EOS token
        if next_token_id == tokenizer.eos_token_id:
            break

    return tokenizer.decode(generated_ids[0], skip_special_tokens=True)

The Effect of Different k ValuesLink Copied

Now we can see how $k$ shapes the character of generated text. Smaller $k$ means stricter filtering, keeping only the most confident predictions. Larger $k$ allows more exploration of the probability space, creating text with more varied word choices and sentence structures.

The relationship between $k$ and text quality is not monotonic. At $k=1$ we get greedy decoding: deterministic but repetitive. As $k$ increases toward 10-20, we gain diversity while maintaining high quality. As $k$ grows past 100-200, we start including tokens that are implausible for the context, and quality begins to degrade. The sweet spot depends on the task and the model, but values in the range 40-100 work well for most general text generation applications.

Out[12]:

Console

Prompt: 'Artificial intelligence will'

Generated continuations:
--------------------------------------------------

k=10: Artificial intelligence will also be more complex, and more diverse, than ever before. For example, artificial intelligence, which is already a key technology for the world's financial

k=50: Artificial intelligence will also create even more exciting challenges for our nation's workers, and many of the benefits will come from improving our ability to build our own factories and reduce

k=200: Artificial intelligence will become so advanced that, like our pets, children, pets will be doing almost any form of AI-related activity. Companies are taking our money.

The outputs reveal distinct personalities. With $k=10$ , the text tends toward common, expected phrasings. With $k=50$ , there's more room for varied word choices while maintaining fluency. With $k=200$ , the model has significant freedom, though even here, we've eliminated the large majority of the vocabulary's long tail.

Out[13]:

Visualization

Heatmap showing sampling counts across token ranks for different k values. — Empirical sampling frequency across 100 samples for different k values. With k=5, nearly all samples come from the top 2-3 tokens. As k increases, the distribution spreads more evenly across the kept tokens, though higher-probability tokens still dominate.

The heatmap reveals how $k$ shapes sampling behavior empirically. With $k=5$ , the dark cells concentrate in the leftmost columns, meaning the top-ranked tokens receive nearly all samples. As $k$ increases, the distribution spreads rightward, but higher-ranked tokens still dominate because they have higher probability even after renormalization. Increasing $k$ does not magically make all kept tokens equally likely; it simply allows lower-ranked tokens to have a non-zero chance of being selected.