Top-k Sampling

What Is Top-k Sampling?

Top-k Algorithm

import numpy as np

def softmax(logits, temperature=1.0):
    scaled = logits / temperature
    e = np.exp(scaled - np.max(scaled))
    return e / e.sum()

def top_k_sample(logits, k=50, temperature=1.0):
    probs = softmax(logits, temperature)

    # Find top-k indices
    top_k_indices = np.argsort(probs)[::-1][:k]
    top_k_probs = probs[top_k_indices]

    # Renormalize
    top_k_probs = top_k_probs / top_k_probs.sum()

    # Sample
    chosen = np.random.choice(top_k_indices, p=top_k_probs)
    return chosen

# With a 10-token vocabulary:
logits = np.random.randn(10)
print('Chosen token:', top_k_sample(logits, k=3))

k=1: Greedy Decoding

logits = np.array([3.0, 2.0, 1.0, 0.5, -1.0])

# k=1: greedy
top_1 = top_k_sample(logits, k=1)
print(f'k=1 always picks: {np.argmax(logits)}')  # index 0, the highest logit
print(f'top_k_sample result: {top_1}')  # always 0

# Multiple runs
for _ in range(5):
    print(top_k_sample(logits, k=1), end=' ')
# Output: 0 0 0 0 0 — perfectly deterministic

Typical Creative Range: k=50

import openai
client = openai.OpenAI(api_key='sk-...')

# Note: OpenAI API does not expose top_k directly in chat completions.
# Top-k is primarily a parameter in Hugging Face Transformers and Anthropic's API.

# Hugging Face example:
from transformers import pipeline
generator = pipeline('text-generation', model='gpt2')
output = generator(
    'Once upon a time',
    max_new_tokens=100,
    do_sample=True,
    top_k=50,
    temperature=1.0
)
print(output[0]['generated_text'])

Top-k in Anthropic's Claude API

import anthropic
claude = anthropic.Anthropic(api_key='sk-ant-...')

# top_k limits the number of tokens considered
message = claude.messages.create(
    model='claude-opus-4-5',
    max_tokens=256,
    temperature=1.0,
    top_k=50,  # sample from top 50 most probable tokens
    messages=[{
        'role': 'user',
        'content': 'Write a short poem about debugging code.'
    }]
)
print(message.content[0].text)

The Fixed-k Problem

# Illustrating the fixed-k problem
logits_confident = np.array([5.0] + [0.1] * 9)  # model is very sure
logits_uncertain = np.array([1.0] * 10)           # model has no idea

probs_conf = softmax(logits_confident)
probs_unc = softmax(logits_uncertain)

print('Confident — top 3 tokens cover:', np.sort(probs_conf)[::-1][:3].sum().round(3))
# ~0.998 — k=50 is extremely wasteful, includes near-zero probability tokens

print('Uncertain — top 3 tokens cover:', np.sort(probs_unc)[::-1][:3].sum().round(3))
# ~0.30 — k=50 may actually be needed to cover a reasonable nucleus

Top-k at Distribution Tails

# Tail behavior comparison
import numpy as np

# Highly skewed distribution (one dominant token)
skewed_logits = np.array([4.0, 2.0, 1.5, 1.0, 0.5,
                           0.1, 0.0, -0.1, -0.5, -1.0])
probs = softmax(skewed_logits)

print('Probability of tokens 6-9 (tail):')
for i in range(6, 10):
    print(f'  Token {i}: {probs[i]:.4%}')
# These tokens are very unlikely but are included in top-k=10
# Top-p=0.9 would exclude them entirely

Combining Top-k and Top-p

def top_k_top_p_sample(logits, k=50, p=0.9, temperature=1.0):
    probs = softmax(logits, temperature)

    # First apply top-k
    top_k_indices = np.argsort(probs)[::-1][:k]
    top_k_probs = probs[top_k_indices]

    # Then apply top-p within top-k
    sorted_k = np.sort(top_k_probs)[::-1]
    cumulative = np.cumsum(sorted_k)
    nucleus_size = np.searchsorted(cumulative, p) + 1

    final_indices = top_k_indices[:nucleus_size]
    final_probs = top_k_probs[:nucleus_size]
    final_probs = final_probs / final_probs.sum()

    return np.random.choice(final_indices, p=final_probs)

When Top-k Is Preferred Over Top-p

# Constrained output with top-k=4
# For a multiple choice task (A, B, C, D)
# If A/B/C/D tokens have indices 32, 33, 34, 35
# top-k=4 with those as the top-4 logits forces selection from those 4 only

multiple_choice_prompt = (
    'Answer with only A, B, C, or D.\n'
    'What is the capital of France?\n'
    'A) Berlin\n'
    'B) Paris\n'
    'C) Rome\n'
    'D) Madrid\n'
    'Answer:'
)

# With temperature=0, top_k=1: always picks the highest logit token

Practical Top-k Guidelines

TOP_K_GUIDELINES = {
    'factual_qa': 1,           # greedy
    'code_generation': 10,     # near-greedy, correct syntax
    'summarization': 20,       # slightly varied but focused
    'chat': 50,                # natural variation
    'creative_writing': 100,   # wider vocabulary exploration
    'poetry': 200,             # unusual word choices encouraged
}

def call_with_top_k(task, prompt, model='claude-opus-4-5'):
    k = TOP_K_GUIDELINES.get(task, 50)
    claude = anthropic.Anthropic(api_key='sk-ant-...')
    return claude.messages.create(
        model=model,
        max_tokens=512,
        top_k=k,
        messages=[{'role': 'user', 'content': prompt}]
    )

Top-k and Temperature Together

from transformers import pipeline

generator = pipeline('text-generation', model='gpt2')

# Combined top-k + temperature
output = generator(
    'The future of AI is',
    max_new_tokens=80,
    do_sample=True,
    top_k=50,
    temperature=0.9
)
print(output[0]['generated_text'])

# Compare: top-k=1 (greedy)
greedy_output = generator(
    'The future of AI is',
    max_new_tokens=80,
    do_sample=False  # greedy, equivalent to top_k=1
)
print(greedy_output[0]['generated_text'])

Knowledge Check

Recap: Top-k Sampling