Tokens and Vocabularies in LLMs
Greedy decoding uses torch.argmax to always select the highest-probability token, which makes the next-token choice deterministic.
Why the Next Word Is a Choice
A language model does not write a sentence all at once. It produces one token at a time, and at each step it assigns a score to every entry in its vocabulary. Those scores have to be turned into one actual word, and the simplest way to do that is to take the highest-scoring entry every time. Before we look at the mechanics, try a small guess.
What do you think happens?
A tiny vocabulary holds the words the, cat, sat, on, and mat. The model gives the word on the highest score at the current step. Which word does greedy decoding output?
Reveal answer
on
Greedy decoding always takes the single highest-scoring entry. In the small example used in this article, that entry is on, which sits at index 3. Every run on the same scores prints on again.
Tracing One Step from Logits to Word
The next-token pipeline has three stages. First, the logits are the raw scores the model assigns to each vocabulary entry for the current context. Second, softmax converts those logits into probabilities. Third, argmax selects the position with the largest probability, and an inverse vocabulary converts that position back into a word. Follow one step through this path, and notice that every stage passes along a position number rather than a word until the very end.
One greedy step on a five-word vocabulary
Vocabulary positions 0 to 4 hold the words the, cat, sat, on, and mat. The logits for the current context are 0.2, 1.5, -0.3, 2.8, and 0.9. Trace which word greedy decoding generates.
- 1Read the logits
Each logit sits at the same position as its vocabulary entry: position 0 is the (0.2), position 1 is cat (1.5), position 2 is sat (-0.3), position 3 is on (2.8), and position 4 is mat (0.9).
- 2Apply softmax
Softmax turns the five logits into probabilities. Because softmax keeps the ordering of the values, the largest logit becomes the largest probability. The entry at position 3 therefore has the highest probability in the distribution.
- 3Apply argmax
Argmax returns the position of the largest value. Here that position is 3. The other four candidates are discarded at this point.
- 4Apply the inverse vocabulary
The inverse vocabulary maps position 3 back to its word, which is on.
The generated token is on. Running this step again on the same logits produces on every time.
Reading the Winning Index
Python indexes from zero, so the first entry in a vocabulary sits at position 0, and the largest value's position is the index that determines the generated word. The argmax result is only a number. It becomes a word only when the inverse vocabulary is used to look that number up. The lookup is the last place a mistake can hide, because a correct probability computation followed by a wrong lookup still produces the wrong word.
The logits and the probabilities line up position by position. The entry at index 3 in the logits is the same entry at index 3 in the probabilities, so checking the index in one array tells you where to look in the other.
What Temperature Scaling Replaces
Greedy decoding is useful for tracing because it is fixed, but that fixedness is also a limitation when you want output with more variety. Temperature scaling is a technique that adds a probabilistic selection process to the next-token step. Its key change is in the selection step: temperature scaling replaces argmax with a function that samples from the probability distribution. The probability scores the model generates for each vocabulary entry are what the sampler draws from. Because sampling is random, the chosen token is no longer forced to be the single highest-scoring entry, so repeated runs can produce different words from the same context.
Where a Result Can Diverge
When a generated word is not what you expected, check three places in order: the index chosen by argmax, the probabilities that produced that index, and the vocabulary lookup that turned the index into a word. The notes below show where each kind of error tends to enter.
Common mistake Treating the argmax result as the word itself
What it looks likeReading index 3 and reporting 3 as the output token.
Why it failsArgmax returns a position, not a vocabulary entry. The output word exists only after the inverse vocabulary maps that position back to a word.
Do this insteadAlways pass the argmax index through the inverse vocabulary before reading the result.
Common mistake Reading the wrong position in a different array
What it looks likeTaking the largest logit's position and then looking up a word in a vocabulary with a different ordering.
Why it failsThe logits and the probabilities line up position by position, and the vocabulary lookup depends on that same position. A mismatch between arrays moves the answer to a different word.
Do this insteadConfirm that the index you read in the logits, the probabilities, and the vocabulary refers to the same entry.
Common mistake Expecting greedy output to vary between runs
What it looks likeRunning the same context twice under greedy decoding and waiting for a different word.
Why it failsGreedy decoding always takes the top entry, so the same context always produces the same next word. Variation is not expected from this rule.
Do this insteadIf you need variety, change the selection step to sampling, as temperature scaling does.
Common mistake Assuming sampling always returns the top word
What it looks likeExpecting sampling to behave exactly like argmax because on has the highest probability.
Why it failsSampling draws from the distribution, so lower-probability entries such as mat or cat can be selected on some runs.
Do this insteadExpect the highest-probability word to be chosen most often, but not on every run.
Practice Tracing a Step
Using the five-word vocabulary the, cat, sat, on, mat with logits 0.2, 1.5, -0.3, 2.8, and 0.9, state the argmax index, the generated word under greedy decoding, and the single step where a wrong answer could appear if the vocabulary were stored in a different order than the logits.
Hints
- Find the largest logit first, then note its position counting from zero.
- The inverse vocabulary must match the same ordering that the logits use.
- If the vocabulary order changes but the logits do not, the same index points to a different word.
Key Takeaways
- A model assigns a score to every vocabulary entry at each step, and the next token is chosen from those scores one step at a time.
- The pipeline runs logits to softmax probabilities, then argmax to a position, then the inverse vocabulary to a word.
- Greedy decoding uses argmax, which always picks the single highest-probability entry, so the same context always gives the same next word.
- Temperature scaling replaces argmax with sampling from the probability distribution, which adds variety across repeated runs.
- When a result is unexpected, check the argmax index, the probabilities at that position, and the vocabulary lookup in that order.