Top-k sampling chooses the next token from the K most likely candidates predicted by a language model. It truncates the probability distribution to those candidates, making generation more controlled than sampling from the full vocabulary. section 8.5.2 section 8.6
Compared with unrestricted sampling, top-k sampling removes low-probability choices before sampling. A small K, such as 3 or 5, usually produces more predictable and coherent text, but reduces diversity. Without top-k filtering, K effectively equals the vocabulary size. section 8.5.2 section 8.6
A typical generation step first obtains the model’s next-token logits, optionally divides them by temperature, applies softmax, keeps the top K probabilities, renormalizes them, and samples one token. The new token is appended before repeating the process. section 8.5.2 section 8.5
Generating text with top_k=3
A recurrent model can generate ten one-token continuations of “I'm not going to see” with top_k=3 and temperature=0.5. The outputs repeatedly choose “you,” “the,” or “her,” whereas unrestricted sampling produces a wider range of continuations. A longer prompt, “Anna and the prince,” can produce a complete, though imperfect, paragraph. section 8.5.2 section 8.5
Listing 8.6 Generating text with temperature and top-K sampling
def generate(model, prompt , top_k=None,
length=200, temperature=1):
model.eval()
text = prompt.lower().split(' ')
hc = model.init_hidden(1)
length = length - len(text)
for i in range(0, length):
if len(text)<= seq_len:
x = torch.tensor([[word_to_int[w] for w in text]])
else:
x = torch.tensor([[word_to_int[w] for w in text[-seq_len:]]])
inputs = x.to(device)
output, hc = model(inputs, hc)
logits = output[0][-1]
logits = logits/temperature #1
p = nn.functional.softmax(logits, dim=0).detach().cpu()
if top_k is None:
idx = np.random.choice(len(logits), p=p.numpy())
else:
ps, tops = p.topk(top_k) #2
ps=ps/ps.sum()
idx = np.random.choice(tops, p=ps.numpy()) #3
text.append(int_to_word[idx])
text=" ".join(text)
for m in ",.:;?!$()/_&%*@'`":
text=text.replace(f" {m}", f"{m} ")
text=text.replace('" ', '"')
text=text.replace("' ", "'")
text=text.replace('" ', '"')
text=text.replace("' ", "'")
return text
Inspect the generate() implementation to see temperature scaling, probability computation, top-k selection, and sampling in context. section 8.5.2
