Module 3 · Training and Generation
Chapter 15 of 17 Run it in your browser

Chapter 15: Greedy and Sampling¶

Name: Teach Your Computer to Write: Build and Train an LLM from Zero
Author: Truong (Jack) Luu

You are here in the big picture Figure 15.1: Where we are: we are ready to generate new text using the trained model.

Our model is trained and loaded, and now we want to use it to generate new text. But the model doesn't give us one character. It produces a list of 65 scores (logits), one for each possible character in the vocabulary, so we need a strategy to pick the winner. In this chapter you will:

Generate text by always choosing the highest-scoring character (greedy decoding).
Generate text by choosing characters randomly based on their probabilities (sampling).
Compare the trade-offs between deterministic and creative text generation.

Words to Know

Greedy Decoding: always picking the single highest-probability token.
Sampling: choosing the next token randomly, giving higher-probability tokens a better chance to be selected.

Theory¶

Greedy Decoding¶

The simplest strategy, called greedy decoding, is to always pick the character with the highest score.

While it seems logical to always pick the "best" answer, greedy decoding has a major flaw, which is that it often gets stuck in loops. Imagine the model learns a common phrase. It predicts the next character, and that character becomes part of the context. The context looks familiar, so it predicts the next character of the phrase, and soon it is repeating the same phrase forever.

Greedy decoding always gives the same result (it is deterministic). If you give it the same prompt, it will always produce the exact same text.

Sampling¶

To avoid repetitive loops, we can use a strategy called sampling. Instead of automatically taking the top character, we turn the raw scores (logits) into percentages (probabilities) and draw a winner randomly.

Argmax taking the top peak vs sampling from the distribution Figure 15.2: Greedy decoding always takes the top peak; sampling draws randomly from the distribution.

If the letter "e" has a 60% probability, it will be chosen 60% of the time. If "x" has a 1% probability, it is rarely chosen, but it still has a chance. This introduces variety and breaks repetitive loops, making the text feel more natural and creative. However, it also means the output is unpredictable, because a different random draw gives a different text. The scripts in this book fix the draw with a seed (torch.manual_seed(42)), so your run prints the same text as the book.

Generation loops the New Text back to the Tokens step on the "where we are" map. The model predicts one character, we add it to the prompt, and the cycle repeats.

In Business

In our house-style assistant, the choice between greedy and sampling depends on the task. If the assistant is answering a factual question from a company manual, you want greedy decoding, which is safe, predictable, and exact. If it is drafting a creative marketing email in the company voice, you want sampling, which is varied and creative and explores new options.

Code¶

We implement both strategies in the same generation loop. The core difference is how next_id is chosen.

src/ch14_generate_greedy.py (excerpt)
def generate_greedy(model, prompt, encode, decode, cfg, max_new_tokens=200):
    # Always picks the token with the highest predicted score
    ids = torch.tensor([encode(prompt)], dtype=torch.long)
    for _ in range(max_new_tokens):
        ctx     = ids[:, -cfg.block_size:]
        logits  = model(ctx)
        next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        ids     = torch.cat([ids, next_id], dim=1)
    return decode(ids[0].tolist())

def generate_sample(model, prompt, encode, decode, cfg, max_new_tokens=200):
    # Picks the next token randomly based on the model's probabilities
    ids = torch.tensor([encode(prompt)], dtype=torch.long)
    for _ in range(max_new_tokens):
        ctx    = ids[:, -cfg.block_size:]
        logits = model(ctx)
        probs  = F.softmax(logits[:, -1, :], dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        ids = torch.cat([ids, next_id], dim=1)
    return decode(ids[0].tolist())

Line 7 picks the character with the single highest score for greedy decoding. Lines 17 and 18 turn the scores into probabilities and draw a winner randomly for sampling.

Code flow: model(ctx) gives logits, argmax or multinomial picks next_id, which is appended to context Figure 15.3: The generation loop predicts one character at a time and appends it to the context.

Notice ids[:, -cfg.block_size:], which keeps only the last block_size tokens of the text so far (the context) before we pass it to the model. The model only learned to read a specific maximum length (our block size of 128) during training, so we must feed it at most that many tokens. We also take logits[:, -1, :] because only the scores at the last position predict the character that comes next.

We can see the two strategies in action with the prompt ROMEO:\n.

Terminal

$ python src/ch14_generate_greedy.py
--- GREEDY (always picks highest-score token) ---
ROMEO:
I will the come the some the stand the son,
And the so shall the see the shall the see the stand
...
--- SAMPLING (picks randomly from distribution) ---
ROMEO:
My must greather, but what it? whom, he ere.

EBRUTUS:
Morcy there bagainVain, I will ever may
To epated as great that the appral tankswain
...
--- Observations ---
Greedy tends to be more repetitive.
Sampling is more varied but can make unexpected choices.

Terminal

$ python diagrams/charts/ch15_probabilities.py
Saved chart: ch15-probabilities.png

Bars comparing greedy and sampling on the same probabilities Figure 15.4: Greedy always picks the tallest bar; sampling can pick any bar based on its height.

What just happened:

The greedy approach quickly got stuck in a repetitive loop ("the shall the shall...").
The sampling approach produced varied, non-repetitive text, but it included some strange spellings ("greather", "bagainVain") because it occasionally picked low-probability characters.

Shape Check¶

Table 15.1 lists the shapes used during the character prediction step.

Table 15.1: Tensors used during the text generation loop.

Tensor	Shape	What it means
`logits[:, -1, :]`	`[1, 65]`	65 vocabulary scores for the final character in the sequence.
`next_id`	`[1, 1]`	1 chosen character ID.

Try It¶

Try It

Different starting prompts create different context for the model. The script below shows how sampling handles new prompts.

src/examples/ch15_different_prompts.py (excerpt)
# ...
encode_fn = lambda s: encode(s, char_to_id)
decode_fn = lambda ids: decode(ids, id_to_char)

print(f"--- Prompt: 'KING:\\n' ---")
print(generate_sample(model, "KING:\n", encode_fn, decode_fn, cfg, 50))

print(f"\n--- Prompt: 'JULIET:\\n' ---")
print(generate_sample(model, "JULIET:\n", encode_fn, decode_fn, cfg, 50))

Lines 6 and 9 generate new text starting from two different prompts.

Terminal

$ python src/examples/ch15_different_prompts.py
--- Prompt: 'KING:\n' ---
KING:
My mine gore to kink on Villence,
For my love grin

--- Prompt: 'JULIET:\n' ---
JULIET:
What reason this bagainVain, I will evers,
As do p

Key Takeaways¶

Greedy decoding always chooses the highest-scoring token, which leads to predictable but repetitive text.
Sampling chooses the next token randomly based on the probability distribution, creating varied and natural text.
The generation loop works by predicting one character, appending it to the context, and repeating the process.
We crop the input context to the model's maximum block_size before each prediction.

Check Your Understanding¶

Why does greedy decoding often get stuck repeating the same phrase?
What PyTorch function do we use to randomly pick a token based on its probabilities?
Why do we slice the logits tensor with [:, -1, :] in the generation loop?

Assignment¶

Assignment 15: Which half changes

About 15 minutes. You change one number and write no new code. Work in this chapter's Colab notebook or on your own laptop.

Run it. Run src/ch14_generate_greedy.py as it is. Copy the text under the GREEDY heading and the text under the SAMPLING heading.
Change one number. In the demo part of the file, find torch.manual_seed(42). The seed fixes the random draws, so every run of the script picks the same characters. Change 42 to any whole number you like and run the script again.
Compare. Which of the two texts changed? Which one is the same, character for character? In the greedy text, which words keep coming back?
Explain it. In two or three sentences a manager could follow, say which rule you would choose for a task that must give the same answer every time, and what you give up by choosing it.

Hand in: both pairs of texts, your answers to step 3, and your sentences.

Chapter 15: Greedy and Sampling¶

Theory¶

Greedy Decoding¶

Sampling¶

Code¶

Shape Check¶

Try It¶

Key Takeaways¶

Check Your Understanding¶

Assignment¶

Further Reading¶