The loop, sampling, and stopping

The loop

In the previous part, we left off with the model turning an input into a probability distribution over the whole vocabulary, its guess for the next token. But you don’t experience models producing one token. You experience whole paragraphs, streaming out a piece at a time. So how do we get from a single prediction to a full response?

The answer is almost disappointingly simple: we run the model again. We take the token we just picked, append it to the input, and feed the whole thing back in. Now the input is one token longer, and the model gives us a distribution for the next next token. Append that too, feed it back, and repeat. The paragraph you see is just this loop running: predict, append, predict, append.

We’re using GPT-2 because it’s small enough to run locally and easy to inspect. Here it is in code, picking the token with the highest probability each step and running for 50 iterations. First, the setup:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
lm = AutoModelForCausalLM.from_pretrained("gpt2", attn_implementation="eager")
lm.eval()
prompt = "the river bank was"

Then the loop:

ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
for _ in range(50):
    with torch.no_grad():
        logits = lm(ids, use_cache=False).logits
    next_id = torch.argmax(logits[0, -1])  # pick the most probable token
    ids = torch.cat([ids, next_id.reshape(1, 1)], dim=1)
print(tokenizer.decode(ids[0]))

Every iteration runs the model, takes the last position’s prediction, and appends the chosen token ID to the input. We keep the IDs and decode them for display. Decoding each token into text and tokenizing it again can change the token sequence. GPT-2 returns predictions at every input position, but the last position is the one we need to continue the text.

There is no separate “generation” machinery in the weights. The model predicts the next token, and the loop turns those predictions into text. The weights determine what the model is likely to write. In the loop, we choose a token and decide whether to keep going. When you later call model.generate(), this is the basic mechanism running underneath.

You may now think that this loop is trivial, boring even. But it’s where we get to do something with the model’s prediction. As you might have figured, always feeding the whole input with the newly appended token results in a lot of unnecessary calculations. We can save some of that work (spoiler: KV cache). The loop is also the point where we can directly interact with the probability distribution. We do not have to simply take the token with the highest probability. We can play around and see what happens if we don’t. We can even mask certain tokens by setting their probability to zero (spoiler: structured output).

So even if it’s trivial, a lot that actually matters lives inside this loop.

How you pick

Always picking the token which was assigned the highest probability seems to be the most intuitive approach, right? The model seems to be most confident about this token, so why pick one with a lower probability? Because the most probable next token is not necessarily the one that leads to the best paragraph. Greedy decoding makes a local choice, one token at a time.

Always picking the most probable token can lead to something counterintuitive: text degeneration. With strictly greedy decoding, patterns such as repetitiveness and, for us humans, weird genericness can emerge. In their experiments with freeform text generation, Holtzman et al. found that maximizing likelihood can produce text that is repetitive and bland. Greedy decoding can work well for other tasks. Here, though, we want it to write a paragraph. The Curious Case of Neural Text Degeneration

With our loop, the repetition is immediately visible:

the river bank was a major source of water for the city.
The city’s water supply was also a major source of food and medicine.
The city’s water supply was also a major source of food and medicine.
The city’s water supply

One explanation is that repetition can reinforce itself. Xu et al. found that the more often a sentence was repeated in the context, the more likely their models became to repeat it again. So once a loop gets started, each repetition can make the next one more likely. Learning to Break the Loop

If you want to look at how models pick up and continue patterns, see Olsson et al.’s In-context Learning and Induction Heads.

As for the genericness, human language is very different from always picking the most probable continuation. A distribution can contain several plausible ways to continue, and greedy decoding throws all but one away.

Temperature

To give those other continuations a chance, we can sample from the distribution. Temperature lets us control how concentrated that distribution is. It is a scaling factor $1/T$ applied to the logits before softmax. Here we can see the probability assigned to token $i$ with corresponding logit $z_i$:

Setting $T=1$ gives us the normal softmax. A temperature between 0 and 1 sharpens the distribution: the strongest candidates get a larger share. Setting $T>1$ flattens it, giving weaker candidates a better chance. What matters is the gap between logits, which gets divided by $T$.

Setting $T=0$ is undefined, because we would divide by zero. Some APIs interpret zero as a request for greedy decoding. In our code, we use argmax for that. As $T$ approaches zero from above, sampling approaches greedy decoding, provided one token has the highest score. If several tokens tie for the highest score, the limiting distribution splits its probability between them.

Here are the five most probable tokens at three temperatures for our first prediction. Their probabilities are calculated over the whole vocabulary. The last row adds up everything else:

So let’s replace the selection line in our loop:

T = 0.8
probs = torch.softmax(logits[0, -1] / T, dim=0)
next_id = torch.argmax(probs)

But we still get exactly the same output. Temperature rescales the distribution. It doesn’t rerank the logits! The order stays exactly the same, and we are still picking the most probable token with argmax. So the next step is to draw from the distribution:

next_id = torch.multinomial(probs, num_samples=1)

Now a token with a probability of 10% has a 10% chance of being picked at this step. Here is the loop with sampling, using the same random seed for each temperature so we can reproduce the comparison:

for T in [0.001, 1.0, 2.0]:
    rng = torch.Generator(device="cpu").manual_seed(42)
    ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
    for _ in range(50):
        with torch.no_grad():
            logits = lm(ids, use_cache=False).logits[0, -1]
        probs = torch.softmax(logits / T, dim=0)
        next_id = torch.multinomial(probs, num_samples=1, generator=rng)
        ids = torch.cat([ids, next_id.reshape(1, 1)], dim=1)
    print(f"T={T}: {tokenizer.decode(ids[0])}")

the river bank was a major source of water for the city.
The city’s water supply was also a major source of food and medicine.
The city’s water supply was also a major source of food and medicine.
The city’s water supply

the river bank was solid and empty a few days in advance."
Does the project need a new permit?
Can’t find any recent maps that ACGS has published of Salerno. The CHANGE Alexandria cost $65,600, has an estimated completion

the river bank was solid fairy fountains commissioned in dependencies more abound Zurich, Vice La Gear Glad Inspective Publishing happens bacon farewell Yamato Midnight)." Takes Once It ACHT WORLD American loneliness Sal Meielve Dal mogocent Alexandria cost cambpp North bands sm--- hisicism

At $T=0.001$, this run matches the greedy output exactly. At $T=1$, it escapes the repeated sentence but wanders between unrelated topics. At $T=2$, we get fragments like “solid fairy fountains commissioned in dependencies”. It is sampling from the entire vocabulary, and the weaker candidates collectively have a lot of probability mass.

We got rid of the repeated sentence, but the text still wanders off. With this model and prompt, high temperature without a filter quickly gives us garbage. These are just individual runs, so they tell us little about which temperature works best in general.

Top-k

Only keep the $k$ largest logits, mask the rest, and sample from the remaining distribution. For inputs where some tokens clearly dominate the distribution, and $k$ is much larger than that count, the resulting distribution can still include a long tail of weak candidates. If $k$ is too small, we can cut off useful alternatives instead.

Top-p (nucleus sampling)

First sort the tokens by probability and keep the smallest group whose probabilities add up to at least $p$. That group is the nucleus. Renormalize its probabilities and sample. The number of tokens we keep can vary per input. There might be an input where $p=0.8$ is reached by two tokens, and another where it needs 20. This is the nucleus sampling introduced by Holtzman et al. in The Curious Case of Neural Text Degeneration.

Here is a small implementation. We mask excluded logits with negative infinity, so softmax gives those tokens zero probability. The token that takes us across the threshold stays in the nucleus:

def top_p_filter(logits, p):
    sorted_logits, indices = torch.sort(logits, descending=True)

    # Running probability total, from most to least likely.
    cumulative = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
    remove = cumulative >= p

    # Shift the mask right so the token that reaches p stays in.
    # clone() avoids reading and writing overlapping tensor slices.
    remove[1:] = remove[:-1].clone()
    remove[0] = False  # Always keep the most likely token.

    filtered = logits.clone()
    # indices maps the sorted mask back to the original token positions.
    # Negative infinity becomes zero probability after softmax.
    filtered[indices[remove]] = -float("inf")
    return filtered

Temperature and top-k or top-p are often combined. In our loop, temperature first rescales the logits, then top-p selects the nucleus and cuts off the tail. Order matters here: temperature can change how many tokens are needed to reach $p$. For our first prediction, a nucleus of $p=0.9$ contains 35 tokens at $T=0.5$, 1,279 at $T=1$, and 14,905 at $T=2$. At that temperature, even the filtered distribution gives us thousands of candidates. The filter only sees their probabilities.

Watch out when switching frameworks, because they don’t necessarily apply these operations in the same order. Hugging Face Transformers 4.57.6 applies temperature before top-k and top-p, while llama.cpp’s documented default sampler chain puts temperature after them. Neither ordering is universally required, but the same parameter values can produce different candidate sets if the order changes. Transformers implementation, llama.cpp defaults

Sampling also explains why the same prompt can give different answers. The three runs below use $T=1$, $p=0.9$, and different seeds. Once one token differs, the model also sees a different context for the next prediction. Fixing the seed lets us reproduce these examples in the same environment. Different backends can still produce different output.

the river bank was actually able to avoid being blown out by the flood."
They added that the 3-miles across and along the stream led to huge trenches that coated the stream before it drained away.
The excavation site was discovered as part of a

the river bank was solid and empty). The activity in this county, though, is notable: once again, we find it an intoxicating mixture of eighteenth-century cities, American Indian villages and American farming communities and American statesmen. You will notice a separation in both

the river bank was 18 inches deep, but it was saved by 26 investors.
Xenfie and her father, Albert, bought a 5 million square foot house at Camp Baals in 1952 and had to sell to pay back the investments.<|endoftext|>

How it stops

In the first example, we break out of the loop after 50 iterations. The model makes 50 predictions. But how can we be sure that after 50 iterations we are not in the middle of a sentence or breaking a chain of thought? Well, we can’t. Counting tokens alone doesn’t tell us when ending the generation makes sense semantically. That’s why we need the model to be able to tell us when it thinks it’s done. But how? It is just predicting the next token over and over again.

When training a model, we can include “special” tokens in the vocabulary and insert them into the training data. An EOS (End of Sequence) token marks boundaries where generation may end. GPT-2 already has one: <|endoftext|>. We don’t need to add it ourselves.

eos = tokenizer.eos_token_id
print(tokenizer.eos_token, eos)

<|endoftext|> 50256

That ID comes from GPT-2’s configuration. A training sequence might look like this:

...end of article A <|endoftext|> start of article B...

The model gets a training signal that associates those boundaries with the EOS token. It can learn when that token is likely to come next, just as it learns about other tokens. Adding a new special token to an already trained model, by itself, would not teach it that behavior.

Chat models are commonly trained further on conversations with an explicit message structure, something like <user> question <assistant> a complete helpful answer <end-of-turn>. The actual markers vary between models. Ending a turn and ending a document can use different tokens. Chat templates

But predicting a special token does not stop the program. The loop has to check for it and break. Here is our sampled loop with both top-p filtering and that check:

T, p = 1.0, 0.9
rng = torch.Generator(device="cpu").manual_seed(43)
ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
for _ in range(50):
    with torch.no_grad():
        logits = lm(ids, use_cache=False).logits[0, -1]
    filtered = top_p_filter(logits / T, p)
    probs = torch.softmax(filtered, dim=0)
    next_id = torch.multinomial(probs, num_samples=1, generator=rng)
    ids = torch.cat([ids, next_id.reshape(1, 1)], dim=1)
    if next_id.item() == eos:
        break
print(tokenizer.decode(ids[0], skip_special_tokens=True))

This run selects EOS as its 47th new token, so it exits before the limit of 50 tokens. The other two sampled runs reach the limit. We still need that cap, because the model is not guaranteed to select EOS in time. EOS can still come after nonsense or an incomplete answer. It’s a learned prediction like any other token.

Putting it together

You built the basic mechanism behind generate() by hand: loop, pick, stop. With Hugging Face, our original greedy example becomes a single generation call:

inputs = tokenizer(prompt, return_tensors="pt")
output = lm.generate(**inputs, do_sample=False, max_new_tokens=50,
                     pad_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The library takes care of the loop, the configured EOS check, and optimizations such as caching. Here, do_sample=False tells it to use greedy decoding, and max_new_tokens=50 sets our limit.

Some hosted models don’t expose all of these controls. For example, Anthropic restricts temperature and top-k changes during thinking on older Claude models, and some newer models restrict sampling parameters regardless of thinking mode. Claude’s sampling restrictions

We’ve seen how the model’s predictions become text: pick a token, append it, repeat, and check when to stop. But what if you need to force it to only pick tokens that form valid JSON? That’s constrained decoding, and it’s next.

Reproduction note: the recorded examples use GPT-2, PyTorch 2.8.0, Transformers 4.57.6, CPU float32, and eager attention. The full evidence script pins the checkpoint and records its environment and seeds in evidence.json. Generated text is shown as model output, not as factual information.