Inference and Decoding

By this point, training is finished. Pretraining gave the model broad language ability. Fine-tuning taught assistant behavior. Preference optimisation made preferred responses more likely than rejected ones.

How does a trained language model turn a probability distribution into an actual response?

Given a prompt, the model produces probabilities over vocabulary tokens. It has not yet written a response. Something must choose one token, append it to the context, and repeat the process. That decision process is decoding.

Inference is the full process of using the trained model to generate outputs. Decoding is the part of inference that chooses tokens from model probabilities.

During training, gradient descent changes the model parameters. During inference, the weights are frozen. The model does not call loss.backward(), does not calculate parameter gradients, and does not learn from the current conversation just because it has seen it.

The subscript $\theta$ remains because the probabilities depend on learned weights. But $\theta$ no longer changes.

Autoregressive Generation

Most LLMs generate autoregressively: each new token depends on the tokens before it. If the model chooses cheese after The best burger topping is, the future context is different than if it chose onion or subjective.

During training, the next tokens come from the dataset. During inference, the model must live with its own choices. One early token can change every later probability distribution.

The transformer usually outputs logits first: one real-valued score per vocabulary token.

Logits are not probabilities. They can be negative and do not need to add to one. Softmax converts them into a probability distribution:

For logits $[2,1,0]$, exponentiating gives $[7.39,2.72,1]$. The total is $11.11$, so the probabilities are approximately:

A response contains many tokens, so its probability is a product of conditional probabilities. Systems usually work with log probabilities because products become sums:

Greedy Decoding

Greedy decoding chooses the highest-probability token at every step:

It is deterministic, cheap, and easy to understand. But it chooses the best local token, not necessarily the best full sequence. A first token with slightly lower probability may lead to a much stronger continuation.

Sampling

Sampling draws from the probability distribution. A token with probability $0.30$ appears about 30 percent of the time across many generations. Sampling preserves uncertainty and variation, but unrestricted sampling can select strange low-probability tokens from the long tail.

Temperature

Temperature reshapes logits before softmax:

Low temperature sharpens the distribution and makes the top token more dominant. High temperature flattens the distribution and gives lower-probability tokens more chance. Temperature does not change token rankings for positive $T$. It changes the gaps, not the order.

Top-k And Top-p

Top-k keeps the $k$ highest-probability tokens and renormalises within that set. Top-p, or nucleus sampling, keeps the smallest set of highest-probability tokens whose cumulative probability reaches at least $p$.

Top-p is not an individual token threshold. With $p=0.9$, we do not require every token to have probability above 0.9. We keep tokens until their total probability mass reaches 0.9.

Beam Search

Beam search keeps several candidate sequences at once. Beam width $B=1$ is greedy search. Larger beams retain more possible continuations.

Beam search compares accumulated log probabilities, often with length normalisation:

Without length correction, raw probabilities can prefer short outputs because every extra token multiplies by another probability less than or equal to one.

Generation stops when the model selects an end-of-sequence token, reaches a maximum token limit, emits a stop string, calls a tool, or is interrupted by an external controller.

Repetition controls modify logits to discourage repeated tokens. A presence penalty depends on whether a token appeared at all:

A frequency penalty depends on how often it appeared:

These controls can reduce loops, but aggressive penalties can damage legitimate repetition. They change generation behavior; they do not give the model deeper understanding.

A simplified autoregressive generation loop looks like this:

tokens = tokenizer.encode(prompt)

for _ in range(max_new_tokens):
    logits = model(tokens)
    next_token_logits = logits[-1]
    probabilities = softmax(next_token_logits)
    next_token = decode(probabilities)
    tokens.append(next_token)

    if next_token == eos_token_id:
        break

response = tokenizer.decode(tokens)

The model usually produces logits for every input position. For generation, we care about the final position because it predicts the next token.

A naive inference loop recomputes old tokens again and again. After generating one new token, the earlier prompt tokens have not changed, but the model would still recompute their attention keys and values.

In self-attention:

During autoregressive decoding, previous keys and values remain useful to future tokens. The KV cache stores them for each layer. A new token computes its new query, key, and value, then attends to cached keys and values.

KV caching speeds decoding, but it costs memory:

Longer context and more simultaneous users mean larger caches. KV caching should not intentionally change the answer. It changes how efficiently the same logits are computed.

Speculative decoding uses a small draft model to propose several tokens and a larger target model to verify them. If many draft tokens are accepted, the system advances several positions after one target-model verification step. If the draft model is often wrong, the overhead can erase the speed-up.

Latency and throughput are different. Latency is how long one request waits. Throughput is how much total work the system completes over time. Batching and continuous batching improve hardware utilisation, but they are serving strategies, not changes to the model's conceptual generation process.

Decoding chooses from the probabilities the model supplies. If the model assigns almost no probability to the correct answer, sampling controls may not rescue it. A sensible top-k or top-p filter may remove the correct token entirely.

Determinism is not truth. Greedy decoding can produce the same false answer every time. Randomness is not creativity. High temperature can produce surprising text, but it can also produce nonsense.

Decoding is a control surface. It is not a replacement for better training, better context, retrieval, tools, or evaluation.

Greedy Decoding

import torch

def greedy_next_token(logits: torch.Tensor) -> int:
    if logits.ndim != 1:
        raise ValueError("Expected one-dimensional logits.")
    return int(torch.argmax(logits).item())

Softmax is not required before argmax because softmax preserves ranking.

Top-p Sampling

def sample_top_p(logits, p=0.9, temperature=1.0):
    adjusted_logits = logits / temperature
    probabilities = torch.softmax(adjusted_logits, dim=-1)

    sorted_probs, sorted_indices = torch.sort(
        probabilities,
        descending=True,
    )
    cumulative = torch.cumsum(sorted_probs, dim=-1)

    remove_mask = cumulative > p
    remove_mask[1:] = remove_mask[:-1].clone()
    remove_mask[0] = False

    filtered = sorted_probs.masked_fill(remove_mask, 0.0)
    filtered = filtered / filtered.sum()

    sampled_position = torch.multinomial(filtered, num_samples=1)
    return int(sorted_indices[sampled_position].item())

The mask shift keeps the first token that crosses the cumulative threshold, so the retained set actually reaches the requested probability mass.

Generation Function

@torch.inference_mode()
def generate(model, input_ids, eos_token_id, max_new_tokens=100):
    generated = input_ids.clone()

    for _ in range(max_new_tokens):
        outputs = model(input_ids=generated)
        next_token_logits = outputs.logits[0, -1]
        next_token_id = sample_top_p(
            next_token_logits,
            p=0.9,
            temperature=0.8,
        )
        next_token = torch.tensor(
            [[next_token_id]],
            dtype=generated.dtype,
            device=generated.device,
        )
        generated = torch.cat([generated, next_token], dim=-1)

        if next_token_id == eos_token_id:
            break

    return generated

@torch.inference_mode() tells PyTorch not to track gradients. Production systems also add KV caching, batching, streaming, mixed precision, memory management, and request scheduling.