Chapter 12: Generation and Sampling

From trained model to text

After pre-training, a language model can generate text by sampling one token at a time. The chosen token is appended to the context, and the model runs again to predict the next one. Repeating this produces a sequence.

def generate(prompt, max_length=50):
    tokens = tokenizer.encode(prompt)
    for _ in range(max_length):
        logits = model(tokens)
        probs = softmax(logits[-1])
        next_token = sample(probs)
        tokens.append(next_token)
    return tokenizer.decode(tokens)

Greedy and beam search

Greedy decoding always picks the token with highest probability. It is deterministic but often repetitive and boring. Beam search keeps the top k partial sequences and selects the one with the best overall log-probability.

Beam search reduces obvious mistakes, but it still tends to produce generic and repetitive text. Language models generate more creative and natural output when randomness is introduced.

Sampling with temperature

Temperature scales the logits before softmax. A low temperature makes the distribution sharp, so the model almost always picks the top token. A high temperature flattens the distribution, allowing less likely but more surprising tokens.

Top-k and nucleus sampling

Temperature alone sometimes allows very unlikely tokens, producing nonsense. Top-k sampling restricts the choice to the k most likely tokens. Nucleus or top-p sampling is more flexible: it keeps the smallest set of tokens whose cumulative probability exceeds p.

For example, with top-p = 0.9, we sample only from the tokens that together account for 90% of the probability mass. If the model is confident, this may be a single token; if it is uncertain, it may be many.

def top_p_sampling(logits, p=0.9):
    probs = softmax(logits)
    sorted_probs = np.sort(probs)[::-1]
    sorted_indices = np.argsort(probs)[::-1]
    cumsum = np.cumsum(sorted_probs)
    cutoff = np.searchsorted(cumsum, p) + 1
    kept = sorted_indices[:cutoff]
    probs = probs[kept] / probs[kept].sum()
    return np.random.choice(kept, p=probs)

Repetition and length penalties

Without intervention, models repeat phrases. Repetition penalties subtract from the logits of tokens that have already appeared. Length penalties encourage the model to stop at appropriate points. Stop sequences can also be used to end generation cleanly, for example when a closing tag or final punctuation is produced.

The penalty coefficient alpha is usually small, around 1.1 to 2.0. Too much repetition penalty can push the model into unrelated vocabulary, while too little leaves it looping.

Contrastive search

Contrastive search tries to balance model confidence with diversity. At each step it considers the top-k candidates and penalizes those that are too similar to the recent context. This reduces repetition without requiring a separate repetition penalty.

The idea is to pick a token that is both probable and novel. A penalty term based on the similarity between the candidate's hidden state and the previous tokens' hidden states steers the generation away from reuse.

Speculative decoding

Speculative decoding speeds up generation by using a small, fast draft model to guess several future tokens, then verifying them in parallel with the large target model. If the draft is correct, several tokens are accepted at once; if not, the process falls back to the large model at the first mismatch.

This can reduce wall-clock time by a factor of two or more because the large model can evaluate many tokens in one forward pass, and the small draft model is much cheaper to run. It does not change the output distribution; it only makes generation faster.