Tokens, Embeddings and Probability

Chapter 1 ended with a sentence that looped. Forty greedy steps from the same prompt produce this:

An inference engine takes a trained model and a request, and it must turn them into a model that can be used to generate a model that can be used to generate a model that can be used to generate a model that can be used to generate a model that can be used to generate

We called that a property of the decoding rule and not a defect of the engine. This chapter makes that claim testable. To do so we need to look at the numbers at the end of the forward pass, and at everything that happens to text before it reaches one. That is a short path of four objects: token ids, which come from text; embedding vectors, one per id; logits, one score per vocabulary entry; and probabilities, which a selection rule turns into one chosen token. Each is a place where an engine can be wrong without crashing, and each has a measurement that shows how.

Three findings from this chapter are worth carrying forward. A naive softmax, run in 32-bit floating point on logits from a real model, does not crash: it returns probabilities that sum to one and are wrong by up to 0.30. The number of tokens that top-p 0.9 keeps varies, across six ordinary prompts, from 131 to 14,394, so no fixed truncation setting means the same thing twice. And during the loop above, the model’s confidence in its own repetition climbs from 34 per cent to 99 per cent.

Where this sits

Chapter 2 asked what exactly is being executed, and ended on the observation that the tokeniser is the one part of the artefact that changes a request without touching a weight. This chapter takes up that thread and follows the request through the front and back of the model: text in, tokens through the embedding table, scores out, a choice made.

The Pi series has a volume on embeddings, which deals with models trained to represent whole texts so that similar ones are close together, for matching and retrieval. The embeddings in this chapter are a different object with the same name: the lookup table at the very front of a language model, one vector per vocabulary entry, learned together with everything else and not designed for similarity search. Keeping the two apart matters because the same word will be used for both.

The Inference Lab piece added here is the sampling configuration: a small, validated, fully recorded description of how a token is chosen from scores. Chapter 1’s request record could only say “greedy”, and by implication nothing. From here on, a record that includes generated text must also include the rule that selected it.

Text to ids

A model reads a sequence of integers. The tokeniser is the translation between those integers and text, and it is a learned artefact: for GPT-2 it is a byte-pair encoding (BPE) with a vocabulary of 50,257 entries. BPE starts from bytes and repeatedly merges the most frequent adjacent pair in a training corpus, until it has the vocabulary size it was asked for. Frequent strings end up as single tokens and rare ones are spelled out in pieces. Eleven strings tokenised by the GPT-2 tokeniser show what that does in practice:

All eleven decode back to exactly the text that went in. Several things in the table are worth reading slowly.

A space is part of a token. inference costs two tokens and inference costs one, and they are different inputs with different ids. The tokeniser did not remove the space, it absorbed it. A prompt that ends in a space therefore ends in a different place from one that does not, and a template that concatenates strings without regard to leading spaces can silently change what the model sees. I have not measured how much that changes outputs for this model; exercise 2 asks you to.

Case changes the cost and the input. Three tokens for INFERENCE, two for Inference, two for inference. The model has to learn that these are related, and it has no built-in way to know.

Numbers are chunked arbitrarily. 3.14159 becomes 3, ., 14, 159; the date splits into pairs of digits. The pieces follow corpus frequency and not place value. This is background to a well-known difficulty of arithmetic with such tokenisers. The measurement here shows the splitting, not any consequence for accuracy.

Whitespace is expensive in code. Four spaces of indentation were three tokens and half of a fourth. Source code with deep nesting pays for its indentation in context window and in per-token cost.

Non-Latin text falls back to bytes. The two Japanese characters became six tokens, each a fragment of a UTF-8 sequence that is not itself a valid character. That is the mechanism behind Chapter 1’s finding that a short Japanese sentence needed more than three times as many tokens as an English one. It is a property of this tokeniser’s training data, not of the language.

Nothing in the table is a defect. It is the interface. But it has an engineering consequence that is easy to miss: the token sequence, not the text, is the request. Context limits, prices, cache keys and throughput figures are all counted over token ids. The tokeniser files belong in the artefact manifest from Chapter 2 for the same reason the weights do, and a runtime that tokenises with a different version has run a different request.

Ids to vectors

The first operation inside the model is a table lookup. The embedding table \(E\) has one row per vocabulary entry:

The lookup is context-free. The vector for king is the same whether the sentence is about chess or monarchy; making it depend on context is the work of the twelve layers that follow. What the table does encode is whatever the training process found useful about each token on its own. A nearest-neighbour search by cosine similarity over the rows shows that this is not nothing:

Spelling variants, related words and morphological relatives sit close together. Pairs behave as one would hope: cat and dog have cosine 0.55, cat and the 0.21, king and queen 0.66, king and banana 0.26. These are properties of the input table of one small model, and I would not lean on them as a general claim about meaning. They are enough to show that the table is not random.

One matrix used twice

The second use of this matrix is at the other end. The output projection that turns the final hidden state \(\mathbf{h}\) into scores has the same shape, \(V \times d\), and in GPT-2 it is the same storage, not a copy. This is called weight tying. The score of vocabulary entry \(i\) is the dot product of its embedding row with the hidden state:

Vectors to scores

The model outputs the logit vector \(\mathbf{z} \in \mathbb{R}^V\), one number for each vocabulary entry at the final position. Calling these unnormalised scores undersells how unconstrained they are. The logits for GPT-2 on six ordinary prompts had maxima between −31.9 and −102.1. Every logit on every prompt was negative. The scores are only meaningful as differences.

That is not a quirk of this model, it is a property of the function that consumes them. Softmax maps scores to probabilities:

The softmax that fails quietly

The obvious implementation exponentiates the logits, sums, and divides. It fails in three different ways, and the first two are the ones everyone learns about:

The stable form subtracts the maximum \(m = \max_j z_j\) first, which by the identity above changes nothing mathematically and guarantees the largest exponent is \(e^0 = 1\):

I ran the naive formula in 32-bit arithmetic on real GPT-2 logits for six prompts. None crashed, none produced a NaN, and in every case the probabilities summed to 1 within rounding. The maximum error against the stable result was:

For the last prompt, the correct probability of at least one token was wrong by 0.30. The result still sums to one, because dividing by the denominator renormalises whatever the subnormal arithmetic left behind. A greedy decoder would probably pick the right token. A sampler or a probability-reporting API would silently misreport it, and nothing in the output would say so.

This is the clearest case in the chapter of an optimisation-adjacent bug that testing at a convenient scale will not find: the unit test with logits [2.0, 1.0, 0.5, 0.1] passes, and so does every test with moderate values. It fails only on the magnitudes a real model produces. The chapter’s implementation uses the stable form and computes log-probabilities as x − logsumexp(x), and on all 50,257 real logits for six prompts it agrees with torch.softmax in 64-bit to \(1.3 \times 10^{-15}\). That agreement is a result about this implementation against one reference, on these inputs. It is not a proof of correctness, and Chapter 12 returns to the question when precision is deliberately reduced.

Scores to a distribution

The result of softmax is a probability distribution over the vocabulary. Before any selection rule is applied, the shape of that distribution already tells you a great deal. Six prompts give six quite different pictures:

Entropy, \(H = -\sum_i p_i \ln p_i\), measures the spread of the distribution. A uniform distribution over this vocabulary would have \(H = \ln 50257 \approx 10.8\) nats; a distribution that put all its mass on one token would have \(H = 0\). The prompts here range from 3.9 to 8.8.

Two features of the table matter for everything downstream.

This model is not a good answerer. For The capital of France is, the most probable next token is the, with probability 0.085. Paris is fifth, at 0.032. A greedy decoder produces the. GPT-2 is a 2019 model of 124 million parameters, which is why it is a useful teaching instrument and not a useful assistant. The inference machinery should not be blamed for what it returns, and nothing about better machinery would change the ranking. It is another instance of the stage separation from Chapter 1.

Truncation rules do not mean the same thing across contexts. Top-p keeps the smallest set of tokens whose cumulative probability reaches \(p\), and the sizes of those sets here run from 131 to 14,394 at \(p = 0.9\). For The, which is a single token with nothing before it, the nucleus contained 14,394 of 50,257 tokens. A single-token prompt is unusual for GPT-2, which has no beginning-of-text marker prepended by default, so I would not read much into that particular row. The range across the other five, 131 to 1,956, makes the same point. Top-k with \(k = 40\) captures 82 per cent of the probability mass for the code prompt and 38 per cent for the technical one. A setting that discards 18 per cent of the mass in one context discards 62 per cent in another.

None of this makes top-k or top-p wrong. It means they are not interchangeable and that a number like 40 or 0.9 carries no meaning without the distribution it is applied to.

Temperature

Temperature divides the logits before the softmax:

The same model, the same forward pass, and the nucleus size ranges from one token to nearly twenty thousand. Temperature is not a small knob, and the response is not linear: halving it from 1.0 to 0.5 moves the nucleus from 1,956 tokens to 2. All of this happens after the model has finished computing. Changing the temperature changes the selection rule without touching a single forward pass, which is the cleanest possible demonstration of the separation between computation and choice.

Temperature zero is not defined by the formula, because it divides by zero. Implementations make it mean greedy decoding, and that needs to be a deliberate, named case and not an accident of an if statement. The chapter’s sampler originally raised an error for \(T \le 0\), which is safe but leaves no way to ask for the arg-max through the same interface. The Lab’s SamplingConfig therefore has an explicit greedy flag, defined as the highest logit with ties broken towards the lowest index, which makes the behaviour reproducible instead of dependent on floating-point accident.

From distribution to choice

Selection is a pipeline, and its order matters. The one used here is temperature, then top-k, then top-p, then a draw. Each step produces a new distribution for the next. After top-k removes tokens, the survivors are renormalised, and top-p then measures cumulative mass over the renormalised distribution, not over the original one. Applying top-p before top-k, or before temperature, gives different candidate sets from the same logits. Frameworks implement these rules as a sequence of processors, and the order is part of the specification; check the version you use for the order it applies. A result that says only “top-p 0.9” without the order has under-specified the experiment.

The final draw consumes randomness. The chapter’s sampler takes a seed and draws a single uniform number from a fresh generator for that seed, so each call is a pure function of its inputs. I checked that it samples the distribution it claims to: over 10,000 seeds from the toy logits [2.0, 1.0, 0.5, 0.1], the largest gap between an observed frequency and the true probability was 0.0076, against a largest sampling standard error of 0.0049, about 1.5 standard errors. With top-k = 2 the excluded tokens were drawn zero times in 10,000 and the two survivors appeared with frequencies 0.741 and 0.26 against the renormalised probabilities of 0.731 and 0.269.

That test shows the sampler is not obviously biased. It does not show the generator is statistically strong, and a serving system that uses the same seed for every request or reuses a generator across concurrent requests has problems that this test cannot find. Chapter 39 comes back to reproducibility when the randomness is spread across requests.

The loop, explained

We can now explain the opening of the chapter. I ran greedy decoding for 40 steps from the Chapter 1 prompt and recorded, at each step, the probability of the chosen token, the entropy of the distribution, and the size of the 0.9 nucleus. The first step chose a with probability 0.34 and a nucleus of 959 tokens. By step 13 the chosen token had probability 0.62. By step 23, 0.90 with a nucleus of one. From step 31 onwards every chosen token had probability above 0.91, and half of them above 0.97:

Each repetition of the phrase makes the next repetition more likely. The model has a distribution over what comes next; the loop is its own history fed back as evidence, and the evidence says the pattern is working. Holtzman and colleagues describe this positive feedback in the paper that introduced nucleus sampling, and the measurement above is one concrete trajectory of it. It is a single prompt, a single model and a single greedy path, so it illustrates the mechanism without establishing how common it is.

If the feedback comes from selecting the single most probable token every time, then choosing differently should break it. I tested that, at the cost of introducing a measure of repetition: the fraction of distinct three-token sequences among the 40 generated tokens, where 1.0 means no three-token phrase appeared twice.

Greedy decoding repeated itself heavily: of 38 trigrams, 8 were distinct. Every sampled configuration repeated far less. The measurement is about repetition only. It says nothing about whether the sampled text was any good, and there is reason to doubt it. A sample from the unfiltered \(T = 1\) runs reads: “User-Built models. An inference engine (delegated as an recurrence model) is run upon a nested cadence model (called a factorization) of states of an existing base model.” It does not repeat, and it is not coherent either. Repetition and quality are separate axes, and a repetition metric alone would rank that sentence above the greedy output. Picking among selection rules is a task-dependent decision that needs a task-dependent measure. The Lab should be able to run any of these rules, and the claim “this setting is better” belongs to the evaluation chapters.

What the numbers do not show

• One small model, from 2019. Its distributions are broad and its factual ranking poor. A larger or instruction-tuned model would have lower entropy on most of these prompts and a different loop behaviour.
• Six prompts. They illustrate the range of distributions. They are not a sample of anything, and the single-token prompt is atypical.
• One tokeniser. GPT-2’s BPE does not describe other tokenisers. Different families of models treat spaces, numbers and non-Latin text differently, and the token counts in these tables are specific to GPT-2.
• Embedding neighbours are a property of one input table. They do not establish that models represent meaning, and they are not evidence about sentence-embedding models used for retrieval.
• Repetition is not quality. The trigram measurement says that sampled text repeats less. It does not say it is better.
• No calibration claim. A model assigning 0.99 to its next token is not thereby correct 99 per cent of the time. I have not measured whether its probabilities match outcomes.
• CPU, 64-bit softmax. The naive-softmax failure was shown in 32-bit arithmetic; hardware that computes logits in 16-bit has a narrower range again, and Chapter 12 returns to this.

Inference Lab increment

The Lab now has SamplingConfig and select_token:

@dataclass(frozen=True)
class SamplingConfig:
    greedy: bool = False
    temperature: float = 1.0
    top_k: int | None = None
    top_p: float | None = None
    seed: int = 0

def select_token(logits, config: SamplingConfig) -> int: ...

The configuration is frozen and has a describe() method that returns every input to the choice, so a request record can carry it intact. Together with Chapter 2’s artefact manifest and Chapter 1’s stage timings, a record can now answer: which artefact, which tokeniser, which rule, which seed, and how long each stage took. It still cannot say which hardware or kernel produced the logits, which is the content of Chapter 39.

Two limits are deliberate. The sampler works on Python lists, so it will be replaced by a tensor implementation when the computation moves to a GPU. And it does not yet support the penalties and stop conditions that production samplers add. Chapter 35 returns to constraints on which tokens are legal.

What you should now be able to do

• Tokenise a string with a BPE tokeniser, predict whether a change in spacing or case will change the token count, and explain why the token sequence is the request.
• Say what the embedding table is and what it is not, and find its share of a model’s parameters.
• Explain why softmax is invariant to shifts in the logits and why that makes subtracting the maximum safe.
• Show an example of a numerically wrong but non-crashing softmax and say what in the output would reveal it.
• Explain what temperature, top-k and top-p each change, and why two settings of top-p cannot be compared across different prompts without the distribution.
• Say why a repetition loop under greedy decoding is a property of the distribution and not of the engine, and what measurement would show it.

Key terms

• Token: a discrete unit of text defined by a particular tokeniser.
• Byte-pair encoding: a method for building a vocabulary by repeatedly merging frequent adjacent pairs.
• Embedding table: the learned matrix mapping token ids to vectors.
• Weight tying: using the embedding table also as the output projection.
• Logit: an unnormalised score for a vocabulary entry; only differences matter.
• Softmax: the map from logits to probabilities.
• Subnormal number: a floating-point value smaller than the smallest normal number, with reduced precision.
• Entropy: \(-\sum p_i \ln p_i\), a measure of the spread of a distribution.
• Temperature: a divisor applied to the logits before softmax.
• Top-k, top-p: restrictions of the candidate set by count or by cumulative probability.
• Nucleus: the set of tokens kept by top-p.
• Greedy decoding: selecting the highest-scoring token at every step.

Reader exercises

1. Prove, for the stable softmax, that the largest term in the denominator is exactly 1. Then find the smallest \(\Delta\) between two logits for which the stable form still assigns the smaller one a non-zero probability in 32-bit arithmetic.
2. Compare the next-token distributions for The capital of France is and the same string with a trailing space. How much does the entropy change and does the top token change? What does that imply for prompt templates?
3. Add min_p filtering (discard tokens whose probability is below a fraction of the top token’s). Run it on the six prompts and compare nucleus sizes with top-p.
4. Run the trajectory experiment with T = 0.3, top-p 0.9 and several seeds. At what step does the probability of the chosen token first exceed 0.9? Does the sampler loop?
5. Take the 14,394-token nucleus for The and prepend GPT-2’s end-of-text token as a start marker. Does it narrow? What does that say about the single-token result above?
6. Replace the cosine nearest-neighbour experiment with one over a model you have available. Does weight tying hold? What share of the parameters is the embedding table?

References

• Neural Machine Translation of Rare Words with Subword Units — the paper that introduced byte-pair encoding for neural models.
• The Curious Case of Neural Text Degeneration — introduces nucleus sampling and discusses repetition under likelihood-maximising decoding.
• Transformers generation strategies — official documentation for temperature, top-k, top-p and their ordering.
• PyTorch log_softmax — reference for the numerically stable formulation used as the check.

Evidence in this chapter

Implementation status

• Manuscript: content/books/inference/03-chapter.md
• Runnable examples: examples/ch03/probability.py, examples/ch03/real_distributions.py
• Tests: tests/test_ch03_probability.py, tests/test_ch03_real_distributions.py (the model-dependent commands are not exercised by tests; the numerical helpers, sampler and configuration are)
• Evidence: research/ch03/sources.md, research/ch03/results.json, research/evidence-register.md
• Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It does not claim anything about output quality, calibration or GPU numerics. Chapter 4 opens the box between the embedding table and the logits: what the twelve layers actually compute.