Loading...
Spaceplay← →stepRreset
See how a language model picks the next token with an interactive visualization. Watch logits become probabilities, compare greedy, temperature, top-k, and top-p sampling, and follow the autoregressive loop token by token.
Loading visualization...
Next token prediction is the single operation a language model performs: given a sequence of token ids, produce a probability distribution over the vocabulary for what comes next. Text generation is that operation run in a loop, with a sampling rule deciding which token to append each time.
A language model does not write sentences. It scores every token in its vocabulary as a candidate for the next position, and a decoding loop picks one, appends it, and asks again. Every word of a chat reply came out of this loop, one token at a time.
The forward pass ends with a linear layer (the lm head) that maps the last position's hidden vector to one real-valued score per vocabulary entry. These scores are logits: unbounded and only meaningful relative to each other. Softmax converts them into probabilities by exponentiating each score and dividing by the sum, so larger gaps between logits become larger ratios between probabilities.
Decoding strategies decide what to do with that distribution. Greedy decoding takes the argmax and is deterministic. Sampling draws a token in proportion to its probability, which is why the same prompt gives different answers on different runs. Temperature rescales the logits before softmax, top-k and top-p truncate the tail before drawing, and these are usually combined.
Only the last position's logits are used during generation, yet the model computes hidden states for every position. Inference engines avoid redoing that work with a key-value cache: attention keys and values from earlier positions are stored, so each new step processes only the newest token and attends over the cache. The cost per token is then roughly proportional to the current sequence length rather than to its square.
Interview framing: define Next Token Prediction in one sentence, then explain one concrete runtime behavior and one common pitfall with a short code example.
Logits can be negative, exceed 1, and do not sum to anything in particular. Ranking them is fine, but reading a logit of 5.8 as 'very confident' ignores the values of every other candidate.
Fix: Apply softmax (with the temperature you actually decode at) before interpreting confidence, and compare probabilities or log probabilities, not raw scores.
Any nonzero temperature means a random draw at every step. Even at T = 0.5, tail tokens keep a small probability, and over hundreds of tokens a rare draw becomes likely.
Fix: Use T = 0 (greedy) or a fixed random seed when outputs must be identical across runs, and remember that batching and floating point differences can still change results.
A fixed k keeps too few candidates when the model is unsure and too many junk candidates when it is certain. The right count changes from step to step.
Fix: Prefer top-p, which sizes the candidate set by cumulative probability, or combine a generous top-k with top-p so the count adapts to the model's confidence.
Temperature changes the probabilities that top-p measures against, so applying top-p first and temperature second produces a different candidate set than intended.
Fix: Follow the conventional order: scale logits by temperature, softmax, then top-k and top-p, then draw. Check the order your inference library uses if results look off.
Common interview prompts with concise model answers.
A vector of logits with one entry per vocabulary token, computed from the last position's hidden state. Softmax turns it into a probability distribution, and a decoding rule picks one token from it. The model itself never chooses; the sampler does.
logits = model(ids)[-1] # shape: (vocab_size,)
probs = softmax(logits / T) # sums to 1
next_id = sample(probs) # argmax when T == 0
ids.append(next_id)Temperature divides every logit by T before softmax. Because division preserves order, the ranking of tokens never changes; only the gaps do. T below 1 widens the gaps and concentrates mass on the top token, T above 1 narrows them and spreads mass into the tail. In the limit T approaches 0 the distribution becomes one-hot on the argmax.
Top-k keeps the k most probable tokens regardless of their mass. Top-p sorts tokens by probability and keeps the smallest prefix whose cumulative probability is at least p, so the candidate count grows when the model is uncertain and shrinks when it is confident. Both zero out the rest and renormalize before drawing.
def top_p(probs, p):
order = argsort(probs)[::-1]
cum = cumsum(probs[order])
cutoff = (cum >= p).argmax() + 1
keep = order[:cutoff]
mask = zeros_like(probs)
mask[keep] = probs[keep]
return mask / mask.sum()Attention at each layer needs keys and values for every earlier position, but those do not change once computed because the model is causal. Caching them means a new step runs the transformer only for the newest token and attends over stored tensors. Without the cache each step costs work proportional to the full sequence length squared; with it the per-step cost is closer to linear in the sequence length.
Because sampling draws each token from a probability distribution using a random number. Any token with nonzero probability can be chosen, and over a long completion the sequence of draws diverges quickly. Greedy decoding or a fixed seed removes this source of variation, though hardware nondeterminism in floating point reductions can still cause small differences.
Continue with these concepts to strengthen your mental model.
ids = tokenize("The capital of France is")for step in range(max_new_tokens): logits = model(ids)[-1] # last position probs = softmax(logits) next_id = argmax(probs) # greedy ids.append(next_id) def softmax(z): e = exp(z - max(z)) # shift for stability return e / e.sum()waiting
waiting
waiting
waiting