Decoding Strategies and Output Control

A language model does not write text directly. Instead, it returns logits for the next token. The decoding algorithm decides how to turn those logits into a token, and repeating this decision produces the output text.

The decoding algorithm affects the behavior of the model. Greedy decoding is deterministic and stable, but it can be dull. Sampling introduces some randomness, which can produce more diverse text but may also produce mistakes. Beam search can be useful for some constrained tasks but is usually not the best default for chat-style generation. Output constraints can make the model produce JSON or stop at a specific marker.

In this chapter, you will learn about:

Greedy decoding
Temperature sampling
Top-k and nucleus sampling
Repetition penalties
Stop conditions
Beam search
Structured output constraints

Let’s get started.

Overview

This chapter is divided into nine parts; they are:

Reading Logits from a Model
Greedy Decoding
Temperature Sampling
Top-$k$ Sampling
Nucleus Sampling
Repetition Penalties
Beam Search
Stop Conditions
Structured Output Constraints

Reading Logits from a Model

The model returns a vector of logits for every position in the input sequence. For generation, you normally use only the last position because it predicts the next token.

The following example uses the Hugging Face transformers library with a small GPT-2 style model. The checkpoint is small enough for local experimentation, but the same logic applies to larger models.

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "sshleifer/tiny-gpt2" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained(model_name) model.eval() prompt = "A language model is" input_ids = tokenizer(prompt, return_tensors="pt").input_ids with torch.no_grad(): outputs = model(input_ids) logits = outputs.logits next_token_logits = logits[:, -1, :] print(next_token_logits.shape) |

The output shape is:

| 1 | [batch_size, vocab_size] |

The logits are not probabilities. To turn logits into probabilities, use softmax:

| 1 | probs = torch.softmax(next_token_logits, dim=-1) |

However, you often do not need to compute probabilities explicitly. Greedy decoding only needs the index of the largest logit, which is the same as the token with the highest probability.

| 1 2 | next_token = next_token_logits.argmax(dim=-1, keepdim=True) print(tokenizer.decode(next_token[0])) |

This is the simplest decoding strategy.

Greedy Decoding

Greedy decoding always chooses the token with the highest score. A complete greedy decoding function can be written as follows:

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 | torch.no_grad() def greedy_decode(model, tokenizer, prompt, max_new_tokens=30): input_ids = tokenizer(prompt, return_tensors="pt").input_ids for _ in range(max_new_tokens): outputs = model(input_ids) next_token_logits = outputs.logits[:, -1, :] next_token = next_token_logits.argmax(dim=-1, keepdim=True) input_ids = torch.cat([input_ids, next_token], dim=1) if next_token.item() == tokenizer.eos_token_id: break return tokenizer.decode(input_ids[0], skip_special_tokens=True) |

Greedy decoding is deterministic. Given the same model and prompt, it returns the same output. This is useful for debugging and for tasks where variation is undesirable.

The weakness is that the best local token is not always the best continuation. Greedy decoding can repeat itself, choose common phrases too often, and miss more interesting continuations.

Temperature Sampling

Temperature sampling draws from a probability distribution obtained by scaling the logits with a temperature parameter.

The figure below shows how temperature changes the probability distribution without changing the underlying logits. The same ten token scores are converted to probabilities three times: once with temperature 0.5, once with temperature 1, and once with temperature 2.

Sampling chooses the next token randomly from the model’s probability distribution. Temperature controls how sharp or flat that distribution is. Given logits $\mathbf{z}$ and temperature $T$, temperature sampling uses:

\mathbf{p} = \operatorname{softmax}(\mathbf{z} / T)

A low temperature makes the distribution $\mathbf{p}$ sharper. A high temperature makes it flatter. If the temperature approaches zero, sampling behaves like greedy decoding, provided one token has a uniquely highest logit. If the temperature is too high, differences between the logits become less important, and the model may choose unlikely tokens too often.

A sampling loop using temperature looks like this:

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | @torch.no_grad() def temperature_decode(model, tokenizer, prompt, temperature=0.8, max_new_tokens=30): input_ids = tokenizer(prompt, return_tensors="pt").input_ids assert temperature > 0, "temperature must be positive" for _ in range(max_new_tokens): outputs = model(input_ids) # apply temperature to the logits for the next token logits = outputs.logits[:, -1, :] / temperature # convert logits to probabilities and sample from the distribution probs = torch.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) # append the next token to the input for next iteration input_ids = torch.cat([input_ids, next_token], dim=1) if next_token.item() == tokenizer.eos_token_id: break return tokenizer.decode(input_ids[0], skip_special_tokens=True) |

Temperature is not a quality knob by itself. It changes the amount of randomness. The right value depends on the task. A factual extraction task usually wants a lower temperature. Brainstorming and creative writing may benefit from a higher temperature.

Top-$k$ Sampling

In the figure above, a 10-token distribution is shown as an example. An actual model may have a vocabulary of hundreds of thousands of tokens, including many that have extremely low probability in a given context.

Top-$k$ sampling keeps only the $k$ highest-scoring tokens and removes all other tokens from consideration. Its primary purpose is to prevent the model from sampling extremely unlikely tokens. It does not avoid computing logits over the full vocabulary, but it does reduce the number of candidates you sample from.

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 | @torch.no_grad() def top_k_sample(logits, k): assert k > 0, "k must be positive" assert k <= logits.size(-1), "k must not exceed the vocabulary size" # Get top-k logits and their indices values, indices = torch.topk(logits, k) # Convert logits to probabilities over top-k candidates probs = torch.softmax(values, dim=-1) # Sample from top-k indices according to their probabilities sampled = torch.multinomial(probs, num_samples=1) # Recover actual token ids using gathered top-k indices next_token = indices.gather(-1, sampled) return next_token |

Top-$k$ is easy to understand, but it uses a fixed number of candidates. Sometimes the model is very confident and only a few tokens matter. Sometimes many tokens are plausible, in which case a fixed top-$k$ cutoff may be inappropriate. This motivates nucleus sampling.

Nucleus Sampling

Nucleus sampling, also called top-$p$ sampling, keeps the smallest set of tokens whose cumulative probability is at least $p$. For example, with $p=0.9$, it keeps the most likely tokens that together account for 90 percent of the probability mass.

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | torch.no_grad() def top_p_sampling(logits, temperature=1.0, k=0, p=0.9): """ Apply temperature scaling, optional top-k filtering, and top-p filtering. Accept a 1D tensor of logits and return one sampled token ID. """ assert logits.dim() == 1, "logits must be a 1D tensor" assert 0 < p <= 1, "p must be in (0, 1]" vocab_size = logits.size(0) # Apply temperature logits = logits / temperature # Optional top-k filtering if k > 0 and k < vocab_size: topk_vals, topk_idx = torch.topk(logits, k) # Create a mask filled with -inf, put top-k logits at their indices new_logits = torch.full_like(logits, float('-inf')) new_logits[topk_idx] = topk_vals logits = new_logits # Top-p (nucleus) filtering sorted_logits, sorted_indices = torch.sort(logits, descending=True) sorted_probs = torch.softmax(sorted_logits, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) # Tokens to remove, but keep at least one token remove = cumulative_probs > p remove[1:] = remove[:-1].clone() remove[0] = False sorted_logits = sorted_logits.masked_fill(remove, float('-inf')) # Sampling final_probs = torch.softmax(sorted_logits, dim=-1) sampled = torch.multinomial(final_probs, num_samples=1) next_token = sorted_indices.gather(-1, sampled) return next_token |

The function above combines temperature sampling, optional top-$k$ filtering, and top-$p$ filtering. Combining these techniques is common. Their order matters because temperature scaling and filtering affect the distribution from which the next token is sampled. Top-$p$ is adaptive: it may keep only a handful of tokens when the model is confident and many tokens when the distribution is broad.

Repetition Penalties

Autoregressive models can fall into loops in which a pattern of tokens repeats itself. Adding a repetition penalty reduces the scores of tokens that have already appeared so that those tokens are less likely to be chosen again.

One simple version divides positive logits by the penalty and multiplies negative logits by the penalty:

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | @torch.no_grad() def apply_repetition_penalty(logits, generated_ids, penalty=1.1): assert logits.dim() == 2 and logits.size(0) == 1, ( "logits must have shape [1, vocab_size]" ) assert generated_ids.dim() == 2 and generated_ids.size(0) == 1, ( "generated_ids must have shape [1, sequence_length]" ) assert penalty >= 1.0, "penalty must be at least 1" if penalty == 1.0: return logits logits = logits.clone() token_ids = set(generated_ids[0].tolist()) for token_id in token_ids: token_logit = logits[0, token_id] logits[0, token_id] = torch.where( token_logit > 0, token_logit / penalty, token_logit * penalty, ) return logits |

This function is intentionally simple and assumes a batch size of one. For example, multiple occurrences of the same token do not increase the penalty. The caller also decides whether generated_ids includes prompt tokens, generated tokens, or both. If you use repetition penalties with top-$k$ or nucleus sampling, apply the penalties first. Production implementations usually handle larger batches and may also distinguish frequency penalties from presence penalties.

Repetition penalties can help, but they can also harm quality. Some words should repeat. Code, names, citations, and structured formats often require exact repetition. Use this control only when repetition is a real problem.

Beam Search

Greedy decoding keeps only one candidate sequence. Beam search keeps several candidates. At each step, it expands each candidate with possible next tokens and keeps the best-scoring sequences.

Beam search is useful when there is a well-defined sequence-level objective, such as translation in older sequence-to-sequence systems. For open-ended chat generation, beam search often produces generic text because it favors high-probability continuations.

A minimal beam search loop looks like this:

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | @torch.no_grad() def beam_search(model, tokenizer, prompt, num_beams=3, max_new_tokens=20): input_ids = tokenizer(prompt, return_tensors="pt").input_ids beams = [(0.0, input_ids)] # Each iteration adds one token to each beam for _ in range(max_new_tokens): candidates = [] # Expand each beam with its num_beams highest-scoring next tokens for score, token_ids in beams: outputs = model(token_ids) logits = outputs.logits[:, -1, :] log_probs = torch.log_softmax(logits, dim=-1) values, indices = torch.topk(log_probs, num_beams, dim=-1) for value, token_id in zip(values[0], indices[0]): next_ids = torch.cat([token_ids, token_id.view(1, 1)], dim=1) candidates.append((score + value.item(), next_ids)) # Keep only the best num_beams candidates for the next iteration beams = sorted( candidates, key=lambda candidate: candidate[0], reverse=True )[:num_beams] # Return only the best beam as the final output best_score, best_token_ids = beams[0] return tokenizer.decode(best_token_ids[0], skip_special_tokens=True) |

This implementation is deliberately small. A real implementation should normalize scores by sequence length, handle end-of-sequence tokens, and avoid recomputing the whole prefix by using a KV cache.

Beam search is expensive: the loops make generation slower, and the number of beams increases memory usage. If you use four beams, the model tracks four continuations. This increases compute and cache memory compared with ordinary sampling. Therefore, beam search is usually avoided in LLM services.

Stop Conditions

Generation must stop at some point. The simplest stop condition is a maximum number of new tokens. Another common condition is the model’s end-of-sequence token. Usually the vocabulary in a language model contains some special tokens. The end-of-sequence token is one of them.

The greedy decoding example above can be modified to accept an arbitrary stop token:

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 | @torch.no_grad() def greedy_decode_with_stop(model, tokenizer, prompt, stop_token_id, max_new_tokens=30): input_ids = tokenizer(prompt, return_tensors="pt").input_ids for _ in range(max_new_tokens): outputs = model(input_ids) next_token_logits = outputs.logits[:, -1, :] next_token = next_token_logits.argmax(dim=-1, keepdim=True) input_ids = torch.cat([input_ids, next_token], dim=1) if next_token.item() == stop_token_id: break return tokenizer.decode(input_ids[0], skip_special_tokens=True) |

This function checks whether the next token is the stop token. If it is, the loop ends and the function returns the generated text before reaching the maximum number of new tokens. This implementation does not handle batched inputs; it assumes a single prompt. With batched inputs, different sequences may stop at different times, in which case more sophisticated handling is needed.

Structured Output Constraints

Some applications need the model to produce a format such as JSON, a SQL query, or a value from a fixed list. One approach is to prompt the model and hope that it follows the format. A stronger approach is constrained decoding.

The idea is to mask out tokens that would make the output invalid. For example, if the output must be one of three labels, you can score only those labels:

| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | @torch.no_grad() def choose_label(model, tokenizer, prompt, labels): assert labels, "labels must not be empty" # Run the prompt to obtain logits over the vocabulary input_ids = tokenizer(prompt, return_tensors="pt").input_ids outputs = model(input_ids) logits = outputs.logits[:, -1, :] # Score each label, assuming that it is exactly one token in this context label_scores = [] for label in labels: # Include any required leading whitespace in the label string label_ids = tokenizer.encode(label, add_special_tokens=False) assert len(label_ids) == 1, f"{label!r} must encode to exactly one token" label_scores.append(logits[0, label_ids[0]].item()) best_score, best_label = max(zip(label_scores, labels)) return best_label |

This example handles only labels that encode to one token after the prompt. Tokenization can depend on context, including preceding whitespace, so callers must construct the labels accordingly. Multi-token labels require scoring complete token sequences or constraining each decoding step. Structured decoding turns output requirements into token constraints. More advanced systems use grammars, tries, or finite-state machines to decide which tokens are valid at each step.

Constrained decoding can improve reliability, but it can also slow inference. The system must compute and apply token masks at each step. As with every inference technique, you should measure both quality and performance.

Summary

In this chapter, you learned that decoding is the process of choosing tokens from logits. Greedy decoding is deterministic and simple. Temperature sampling, top-$k$ sampling, and nucleus sampling introduce controlled randomness. Beam search tracks multiple candidates but increases inference cost. Repetition penalties and stop conditions help control output length and behavior. Structured output constraints can make model outputs easier to use in applications.

In the next chapter, you will learn how to measure inference performance so that these choices can be compared with real numbers instead of intuition.

Newspaper