Choosing the Next Token

Always picking the most likely word sounds sensible. But try it, and your model soon gets stuck, repeating itself forever. Sampling from everything has the opposite problem, and it might pick a bizarre word. Good generation lies in between.

• Top choice repeats forever
• Full sampling gets weird
• Good generation lies between

Top-k and top-p

Here are eight possible next words, with their scores. Plain sampling can pick anything, even the silly words. Top K keeps only the K best choices. Top P is smarter. It keeps the smallest group of words whose chances add up to a target, so that it adapts by itself. When the model is sure, it keeps only a few words. When it's unsure, it keeps more. Run it. Look at which words each method allows. Plain lets rubbish in, while the others stay sensible. Most chatbots use top P.

import torch

torch.manual_seed(0)
logits = torch.tensor([4.0, 3.5, 3.0, 1.0, 0.5, 0.2, 0.1, -1.0])
words = ["mat", "sofa", "bed", "roof", "moon", "car", "tree", "idea"]

def show(name, p):
    pick = torch.multinomial(p, 12, replacement=True)
    print(f"{name:<10}", sorted({words[i] for i in pick.tolist()}))

p = logits.softmax(0)
show("Plain", p)

top = logits.clone()
top[top < top.topk(3).values[-1]] = float("-inf")
show("Top-k 3", top.softmax(0))

order = p.sort(descending=True)
keep = order.values.cumsum(0) - order.values < 0.7
q = torch.zeros_like(p)
q[order.indices[keep]] = p[order.indices[keep]]
show("Top-p 0.7", q / q.sum())

Plain      ['bed', 'mat', 'moon', 'sofa']
Top-k 3    ['mat', 'sofa']
Top-p 0.7  ['mat', 'sofa']

Beam search

When you need the single best answer, as in translation, there's beam search. Instead of picking one word at each step, it keeps several candidate sentences alive, say five, and extends each of them. At the end, it chooses the best full sentence. It's more accurate, but it often sounds dull, so chatbots usually prefer sampling. Translation and speech recognition prefer beams.

Summary