Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Top-k Sampling › Combining Top-k with Temperature

Michael BrenndoerferUpdated July 29, 20255 min read

Language AI HandbookAll chapters

Part 28 of 67GPT Architecture

Chapter 203Top-k Sampling › Combining Top-k with Temperature

Reading section 7 of 13 · Combining Top-k with Temperature · 5 minSection 7 of 13 · 5 min

Section 7 of 13 · 5 min read

Top-k Sampling › Combining Top-k with TemperatureOne section · Sections in this chapter7 / 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

Top-k sampling and temperature scaling are two different tools that address the same underlying problem from different angles. Temperature reshapes the entire distribution before any truncation occurs, while top-k determines the size of the sampling space after reshaping. Used together, they give you two independent dimensions of control over generation quality.

Temperature scaling, which we covered in the previous chapter, works by dividing all logits by a temperature parameter $T$ before applying softmax. When $T < 1$ , the distribution becomes more peaked, exaggerating the difference between high and low logits. When $T > 1$ , the distribution flattens, compressing those differences. The key point is that temperature operates on the full distribution and affects every token's probability.

Top-k, in contrast, operates after temperature has been applied. It looks at the shaped distribution and says: "From this reshaped distribution, I will only consider the top $k$ candidates." The combination means you first decide what shape of distribution you want (via temperature), then decide how many candidates to sample from (via $k$ ). These are separate controls: you can have low temperature with high $k$ , high temperature with low $k$ , or any combination.

Think of the combination as a two-stage filter. Temperature is the first stage: it determines how spread out or concentrated the probability distribution looks. Top-k is the second stage: it draws a boundary around the highest-probability region of that distribution and refuses to sample outside it. Together, they give you fine-grained control that neither parameter alone provides.

Out[18]:

Visualization

Horizontal bar chart of top-20 tokens at temperature 0.5 showing concentrated probability mass in the top few tokens. — Low temperature (T=0.5) concentrates probability in top tokens, making k=20 effectively sample from fewer options.

Horizontal bar chart of top-20 tokens at temperature 1.5 showing more evenly spread probability across all kept tokens. — High temperature (T=1.5) flattens the distribution, spreading probability more evenly across all 20 kept tokens.

The figure shows how temperature reshapes the distribution within the top- $k$ candidates. At low temperature, the top token dominates even among the kept tokens, meaning the effective sampling space is much smaller than $k$ would suggest. At high temperature, probability spreads more evenly across all $k$ tokens, giving each candidate a meaningful chance.

Temperature-First vs Top-k-FirstLink Copied

The order of operations matters. Standard practice applies temperature first, then top-k:

Temperature scaling: Divide logits by temperature
Top-k selection: Keep the $k$ highest values
Softmax: Convert to probabilities
Sample: Draw from the truncated distribution

In[19]:

Code

def top_k_sample_with_temperature(logits, k, temperature=1.0):
    """
    Top-k sampling with temperature applied first.

    This is the standard order: temperature reshapes the distribution,
    then top-k truncates, then we sample.
    """
    # Step 1: Temperature scaling
    scaled_logits = logits / temperature

    # Step 2: Get top-k
    top_k_vals, top_k_idx = torch.topk(scaled_logits, k)

    # Step 3: Softmax over top-k only
    top_k_probs = torch.softmax(top_k_vals, dim=0)

    # Step 4: Sample
    sample_idx = torch.multinomial(top_k_probs, num_samples=1)

    return top_k_idx[sample_idx].item()

Why temperature first? Temperature affects which tokens qualify as "top-k." When you apply high temperature, the differences between logits are compressed, which can change which tokens make the cut near the boundary. A token at rank 11 with the original distribution might move to rank 9 after high-temperature scaling because its neighbors' scores were squeezed more. By applying temperature before the top-k cut, you ensure that the set of candidates you are sampling from reflects the distribution you want to sample from, not the original distribution.

Applying temperature after top-k selection would mean truncating based on the original distribution, then reshaping the truncated distribution's probabilities. This gives slightly different results and is less theoretically clean. The standard order, temperature first, is conceptually cleaner: it says "here is the distribution I want, now take the top $k$ of it."

Out[20]:

Console

Prompt: 'The best programming language is'

Top-10 tokens at different temperatures:
--------------------------------------------------
T=0.5: [' Java', ' Python', ' the', ' a', ' not', ' C', ' one', ' JavaScript', ' Haskell', ' probably']
T=1.0: [' Java', ' Python', ' the', ' a', ' not', ' C', ' one', ' JavaScript', ' Haskell', ' probably']
T=2.0: [' Java', ' Python', ' the', ' a', ' not', ' C', ' one', ' JavaScript', ' Haskell', ' probably']

The output shows that the top- $k$ tokens remain mostly stable across temperatures, with the highest-ranked tokens appearing in all lists. However, tokens near the boundary (rank 8-10) may swap in and out as temperature changes, which can subtly affect sampling behavior. At very low temperature, the ranking becomes almost deterministic because logit differences are amplified. At high temperature, tokens near the boundary are nearly interchangeable, so the specific set included in the top- $k$ becomes somewhat arbitrary.