Top-k sampling generates text by restricting each next-token choice to the K most probable candidates, then sampling from that smaller set. It sits between greedy decoding, which always chooses one token, and unrestricted random sampling, which considers every token. Unlike temperature, top-k sampling gives tokens outside the cutoff zero probability; both controls can be combined. section 16.2.4
Changing K changes the output
With the same prompt, K=5 produces relatively plausible but repetitive prose, while K=20 allows more varied continuations. Combining K=5 with temperature 0.5 produces more focused, plan-like repetition. As a rule of thumb, smaller K narrows the model's choices; larger K preserves more variety. section 16.2.4
A practical implementation follows three steps:
- Scale the predicted logits by the selected temperature. section 16.2.4
- Use keras.ops.top_k to retain the K highest-scoring candidates and their indices. section 16.2.4
- Sample from the retained scores, then map the sampled position back to its original token index. section 16.2.4
def top_k(preds, k=5, temperature=1.0): preds = preds / temperature top_preds, top_indices = ops.top_k(preds, k=k, sorted=False) choice = keras.random.categorical(top_preds[None, :], num_samples=1)[0] return ops.take_along_axis(top_indices, choice, axis=-1)
The Keras implementation combines temperature scaling, top-k selection, categorical sampling, and index lookup in one decoding function. section 16.2.4
