Next-token prediction and sampling interview questions

Logits are the raw scores the LM head gives every vocabulary token, 50,257 for GPT-2, from the last position's vector. Softmax turns them into probabilities: after "Thank you very", " much" got 0.992.

For a sentence, loop: pick a token, append it, run the model again. So 25 new tokens means 25 passes.

Temperature divides every logit by T before softmax: low T sharpens, high T flattens, and T near 0 approaches greedy. Top-k keeps a fixed number of top tokens; top-p keeps the smallest set whose probabilities reach p, so the cut adapts to the distribution's shape.

At top-p 0.9, GPT-2 kept 1 token after "Thank you very" but 417 for an open-ended prompt. No single k suits both: Holtzman et al.'s argument.

It's a feedback loop: once a phrase has appeared, the model rates it more likely to appear again, and greedy always takes the top choice. Continuing "I opened the box and found", GPT-2 small picked " amount" at 0.027 the first time and 0.988 on the repeat.

Holtzman et al. report the same for most phrases they tested: decoding that maximizes probability tends to give bland or looping text.

The test assumes a determinism a hosted API doesn't promise. Temperature 0 is greedy and uses no random numbers, and our sampling script gave byte-identical output across two runs, but on one machine. AI Engineering notes that even with temperature, top-p, top-k and seed fixed, hardware can change the output, and with an API the hardware is the provider's choice.

So stop asserting exact text from live calls: cache known answers where you need exact strings, or parse the output and compare fields, labels or structure. A seed won't fix it.

Against it: a token probability says how typical the text is after the prompt, not whether it's correct. Asked for the largest planet, GPT-2 small rated " Earth" at 0.0282, slightly above " Jupiter" at 0.0270. The number also moves with sampling settings, and the GPT-4 report found the model much less calibrated after post-training.

Instead, use logprobs internally: log them, flag low-probability answers for review, and check on a labelled set that the flag actually catches errors.

Keep the k best partial sentences. Each step, extend every one with its likely next tokens, score by the sum of log-probabilities, and keep the top k. A sentence's probability is the exponent of that sum.

Greedy is beam width 1, so it misses sequences whose first token wasn't the best. Longer sentences collect more negative terms, so normalise by length when comparing finished beams.

Count each word's successors, then sample in proportion to the counts.

import random
from collections import Counter, defaultdict

def train(words):
    counts = defaultdict(Counter)
    for a, b in zip(words, words[1:]):
        counts[a][b] += 1
    return counts

def next_word(counts, w):
    nxt = counts.get(w)
    if not nxt:
        return None              # unseen word
    ws, cs = zip(*nxt.items())
    return random.choices(ws, weights=cs)[0]

Interviewers check the unseen-word case and whether you can sample without random.choices: a cumulative sum plus binary search.

Self-attention draws queries, keys and values from one sequence; cross-attention takes queries from the decoder and keys and values from another source, like an encoder's output.

Decoding is a loop: score the next token, pick one greedily or by temperature and top-p, append, repeat, with a KV cache. RL then tunes those sampled outputs against a reward, held near a reference model. Evaluate with perplexity, task benchmarks and human or judge preferences.

Greedy takes the argmax. Both samplers sort the probabilities, keep a prefix, renormalise it and sample.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()

def pick(logits, rng, mode="greedy", k=5, p=0.9,
         temp=1.0):
    if mode == "greedy":
        return int(np.argmax(logits))
    probs = softmax(logits / temp)
    order = np.argsort(probs)[::-1]   # high to low
    if mode == "top_k":
        keep = order[:k]
    else:                             # top_p
        cum = np.cumsum(probs[order])
        keep = order[:np.searchsorted(cum, p) + 1]
    q = probs[keep] / probs[keep].sum()
    return int(rng.choice(keep, p=q))

Interviewers check the max subtracted in softmax, that top-p keeps the token that crosses p, and the renormalisation before sampling.