A few input rows, a full vocabulary of outputs: from token IDs to sampling

Tokenization ends with a sequence of integers. Those integers do not yet tell us what to write next.

The next two operations both involve the vocabulary, but work in different directions: the input side looks up a few rows by ID; the output side uses a context vector to score the entire vocabulary. Even when the input contains only a and b, the output layer assigns scores to c, d and e. A selection rule then turns those scores into one concrete next-token ID.

This article connects input, output and selection using the same forward computation. First, separate the three responsibilities:

Connect the tokenizer to the model

An actual GPT-2 encoding is:

Hi 你! → [17250, 220, 19526, 254, 0]

There are five tokens. The UTF-8 bytes of 你 are split across two tokens, and the preceding space takes another. The encoding and decoding details are covered in a walk through GPT-2’s tokenizer source.

The overview uses GPT-2’s vocabulary size V=50,257 and hidden width D=768. With input length T=5, the embedding lookup produces 5×768. The last position’s context vector is 1×768, and its output scores are 1×50,257.

encoder.json and vocab.bpe configure the tokenizer. Embeddings, the position table, the Transformer and the output layer are language-model parameters. The tokenizer’s symbol-to-ID table and the model’s ID-to-vector table are different tables, but their ID meanings must match.

The new ID k in the diagram is illustrative; no pretrained GPT-2 prediction was run to obtain it. The next example uses a small model whose numbers can be expanded in full. Its IDs belong to a different vocabulary.

The small model: three positions, five candidates, four features

The vocabulary is a→0, b→1, c→2, d→3, e→4. The input a | b | a becomes [0,1,0]. Batch size is B=1; the text below omits that dimension. The model has one attention head and one Pre-Norm Block, no Dropout, and a position table with eight rows.

All weights are set by hand. The model is untrained. Its purpose is to expose the computation; a probability such as 69% says nothing about learned language ability.

A token ID is an index, not a continuous feature. ID 4 does not carry a meaningful numerical distance from ID 1. The integer selects a vector row; computations operate on that vector. Here T always means sequence length; temperature is written τ. Vocabulary size V is also distinct from attention’s Value vectors.

Input: the same a reads the same E row, but adds a different P row

GPT-2 and this teaching model use learned absolute position vectors. Each position looks up two tables, then adds the vectors elementwise:

X₀[t] = E[id[t]] + P[t]

Both a tokens read E[0], but one adds P[0] and the other P[2]. A position number first looks up a vector; the integers 0, 1 and 2 are not directly added to every feature. Each result still has four features. Addition does not concatenate them into eight dimensions.

token_vectors = token_embedding(ids)                  # (1, 3, 4)
positions = torch.arange(ids.shape[1], device=ids.device)
position_vectors = position_embedding(positions)      # (3, 4)
x0 = token_vectors + position_vectors                 # (1, 3, 4)

Position vectors broadcast across the batch dimension. E and P are parameter tables; the looked-up vectors and x0 are values in this forward pass. This addition describes the implementations used here. It is not a universal claim that every Transformer uses a learned absolute position table.

The Transformer turns each row into a context representation

This example follows the Pre-Norm order:

X₁ = X₀ + Attention(LN₁(X₀))
X₂ = X₁ + MLP(LN₂(X₁))
H  = LN_f(X₂)

LayerNorm normalizes the four features within each position, then applies learned scale and offset. It does not pool statistics across positions. The example uses PyTorch’s default eps=1e-5, with scale initialized to 1 and offset to 0.

The attention branch projects the normalized input to Q, K and Value, computes QKᵀ/√4, masks future scores with −∞, applies row-wise softmax, then takes a weighted sum of Value vectors and projects the result. The actual attention weights are:

             Source 0    Source 1    Source 2
Query 0       1.0000      0.0000      0.0000
Query 1       0.0190      0.9810      0.0000
Query 2       0.4683      0.0101      0.5216

Future positions receive zero weight, and each row sums to one. This softmax distributes weight across input positions. The last row reads from positions 0, 1 and 2. The later output softmax distributes probability across vocabulary candidates instead.

The MLP is Linear 4→16 → ReLU → Linear 16→4. The same MLP processes each position independently; attention handles communication between positions. Both residual additions preserve the 3×4 shape. This teaching structure follows V7’s ReLU choice; GPT-2 uses GELU in the corresponding sublayer.

The last position, t=2, now incorporates the current a | b | a context:

h₂ = [0.9546803832, −1.0673931837, 1.0419750214, −0.9292621017]

These four numbers are not probabilities for four tokens. They are four features for the next layer, which maps D dimensions to V candidate scores.

Output: every candidate has a weight row

lm_head is the language model’s output prediction layer; lm stands for Language Model. It has a different job from an attention head: the former maps context representations to vocabulary scores, while the latter gathers information across input positions.

Here the output layer is nn.Linear(4, 5). PyTorch stores its weight matrix as output dimensions × input dimensions, so W_head has shape 5×4:

z[j] = h₂ · W_head[j] + b[j]
z    = h₂ @ W_head.T + b

       (1×4) @ (4×5) + (1×5) → (1×5)

Applying it to all positions gives Z = H @ W_head.T + b, shape 3×5, with the bias broadcast to each row. Output column j corresponds to vocabulary ID j. Which rows of E were looked up does not restrict the output layer’s five candidate weight rows.

Expand the dot product for b / ID 1:

z[1] = 0.954680×1 + (−1.067393)×0 + 1.041975×1 + (−0.929262)×0 + 0.1
     ≈ 2.096655

The other four candidates receive the same dot-product-plus-bias calculation. A matrix multiplication computes all of them at once.

A logit can be negative; it is not yet a probability. Softmax transforms the entire row into probabilities summing to one:

p[j] = exp(z[j]) / Σₖ exp(z[k])

A dot product is a matching score, not cosine similarity

For nonzero vectors:

h · w = ‖h‖ ‖w‖ cos θ

The dot product depends on both direction and magnitude. Neither vector is first divided by its norm here, and the independent output layer also adds a bias. Calling it a “score” describes its role: produce comparable candidate logits for softmax and the selection rule.

Why GPT-2 writes hEᵀ

GPT-2’s model.py looks up input rows with tf.gather(wte, X) and produces output with tf.matmul(h_flat, wte, transpose_b=True). nanoGPT explicitly ties transformer.wte.weight to lm_head.weight.

Weight tying gives the same parameter table two roles: row lookup on input, and a dot product with every row on output. Counting only these input and output weight matrices gives VD parameters when tied and 2VD when independent. Biases, the position table and the rest of the network are excluded from that count.

Training uses each position; generation uses the last

The teaching text is a | b | a | b. Use its first three entries as input, and shift one position to obtain targets:

Input:   [0, 1, 0]   → a, b, a
Targets: [1, 0, 1]   → b, a, b

Every position asks a next-token question. The causal mask prevents position 0 from looking ahead at position 1’s answer. All three questions can be computed in parallel during training.

At the final position the target is ID 1, so the loss is −ln(p[1]) = −ln(0.690284…) ≈ 0.370652. Average over the three positions:

L = (0.4099546 + 3.3200393 + 0.3706515) / 3
  ≈ 1.36688185

loss = F.cross_entropy(
    logits.reshape(-1, 5),   # Three rows of five raw logits
    targets.reshape(-1),    # The correct IDs: [1,0,1]
)
loss.backward()

cross_entropy accepts raw logits and performs a stable log-softmax and negative log-likelihood internally. No token is sampled first, and no manual softmax is needed. The accompanying experiment checks that the token and position embeddings, Q/K/Value, MLP and lm_head weights all receive nonzero gradients. It computes gradients without performing an optimizer update.

For generation, the full context [0,1,0] is already available. We want the token after it, so we use the final position’s distribution. V7 computes logits for every position and then selects logits[:, -1, :]. nanoGPT’s inference branch applies the output layer only to the last context vector.

How probabilities become an ID

Greedy selection: take the current winner

The largest logit and the largest softmax probability identify the same candidate: b / ID 1 here.

This guarantees only the best current step. Suppose the first step assigns A probability 0.6 and B probability 0.4. After A, the most likely next token has probability 0.5; after B, one token has probability 0.99. The two-step path probabilities are 0.6×0.5=0.30 and 0.4×0.99=0.396. Selecting A first does not produce the highest-probability complete path.

Probability sampling: locate a random number in an interval

For hand calculation, this section explicitly uses rounded probabilities [0.2203,0.6903,0.0292,0.0267,0.0335], which sum to exactly one. Divide [0,1) into intervals with those lengths:

Draw u∈[0,1) uniformly and return the ID of the interval containing it. u=0.50 selects b, u=0.05 selects a, u=0.93 selects c, and u=0.95 selects d. A uniform random number lands in an interval with probability equal to its length. That implements the desired categorical draw.

from bisect import bisect_right
from itertools import accumulate

boundaries = list(accumulate(probabilities))
boundaries[-1] = 1.0
token_id = bisect_right(boundaries, u)

bisect_right finds the first cumulative boundary strictly greater than u. If u equals 0.2203, it enters b’s interval, matching the left-closed, right-open convention. Zero-length intervals cannot be selected.

torch.multinomial(probs, num_samples=1) returns a column index sampled by weight. When columns correspond to vocabulary IDs, that index is the next-token ID. The interval method explains the principle; it is not a claim that PyTorch internally executes this Python loop.

Keeping this distribution fixed, 100,000 draws with Python’s Random(12) produce:

Finite sampling need not match the target probabilities exactly. This is also not a sequence of 100,000 generated tokens: in autoregressive generation, each appended ID changes the context, and the model computes a new distribution for the next step.

★ Temperature: derive what changes from a probability ratio

Before softmax, divide every logit by the same positive number τ:

pᵢ(τ) = exp(zᵢ/τ) / Σⱼ exp(zⱼ/τ), τ > 0

Compare b with a. Their shared denominator cancels:

p_b(τ) / p_a(τ)
= exp(z_b/τ) / exp(z_a/τ)
= exp((z_b − z_a)/τ)

Here the score gap is approximately 2.096655 − 0.954680 = 1.141975:

Lower temperature amplifies the leader’s relative advantage; higher temperature narrows the gap. Labels are rounded, while calculations use full precision.

Positive temperature preserves the ranking, so greedy selection still returns b. It changes the sampling probabilities, not model weights. With a unique maximum, the distribution concentrates on it as τ→0⁺. If the maximum is tied, the limit distributes probability across the tied winners. Code must not divide by zero; use argmax explicitly when greedy selection is wanted.

For stable computation, subtract the maximum logit first: exp((zᵢ−max(z))/τ). The common factor cancels during normalization, preserving the mathematical result while preventing exponential overflow.

top-k: restrict the count, then renormalize

At temperature 1, descending probability order is b, a, e, c, d. With k=2, keep b and a. The other three candidates receive zero probability for this draw.

The retained original mass is approximately 0.6903+0.2203=0.9106. Divide by that mass to obtain a new distribution summing to one:

p′_b = p_b / (p_b + p_a) ≈ 0.758042
p′_a = p_a / (p_b + p_a) ≈ 0.241958
p′_c = p′_d = p′_e = 0

k=2 means two candidates participate in the draw. It neither emits two tokens nor gives each a 50% chance. Renormalization preserves their probability ratio. At k=1, only the winner remains, matching greedy selection in this example.

nanoGPT finds the k-th largest logit, sets smaller logits to −∞, and applies softmax. Since exp(−∞)=0, this automatically renormalizes the retained candidates. Its strict “less than the cutoff” comparison means ties at the boundary can retain more than k entries. This example has no ties. The teaching code below takes exactly k ranked entries, breaking ties by ID.

top-p: include the token that crosses the threshold

Top-p accumulates probabilities in descending order and keeps the shortest prefix whose cumulative mass reaches at least p, then renormalizes it. This follows equations (2) and (3) in §3.1 of the Nucleus Sampling paper, with the cumulative condition ≥ p.

At 0.90, b alone is insufficient. Adding a reaches 91.06%, so keep [1,0]. At 0.92, b+a is insufficient. Add e, keeping [1,0,4]. At 0.95, the first three total only 94.41%. Include c as well, keeping [1,0,4,2]. Sorting changes the display order, not the original token IDs.

Top-k fixes the count; top-p lets the count follow the distribution. With p=0.90, a leader holding 96% needs no companion. Five equally likely candidates at 20% each require all five.

Both filters decide which IDs participate in this step’s draw. They do not delete entries from the tokenizer’s vocabulary. A filtered-out ID can return at the next step as the context changes. That differs from Unigram training’s pruning of the persistent vocabulary.

Source boundary: the old GPT-2 function keeps one fewer candidate here

The inspected GPT-2 src/sample.py, at commit 9b63575, computes descending cumulative probabilities and chooses its cutoff using:

cutoff index = max(count(cumulative_probs <= p) − 1, 0)

At p=0.90 in this example, only the first cumulative value is no greater than 0.90. The count is 1, and the cutoff index is 0. The first candidate’s score therefore becomes the threshold: only b remains. The paper’s rule also keeps a, the candidate that takes cumulative mass across 0.90.

The code and interactive demo in this article use the paper’s rule. A separate Python calculation reproduces the old source’s cutoff-index formula and confirms the difference; its TensorFlow implementation was not executed. The boundary condition determines the behavior, regardless of the function’s top_p_logits name.

Filtering code you can reproduce

This minimal version assumes nonempty, finite logits. Top-k and top-p are applied separately. The complete experiment also includes interval sampling, input checks, repeated draws and the old GPT-2 boundary comparison.

import math

def softmax(logits, temperature=1.0):
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be finite and positive")
    maximum = max(logits)
    weights = [math.exp((z - maximum) / temperature) for z in logits]
    total = sum(weights)
    return [w / total for w in weights]

def ranked_ids(probs):
    return sorted(range(len(probs)), key=lambda i: (-probs[i], i))

def top_p_ids(probs, threshold):
    if not 0 < threshold <= 1:
        raise ValueError("top-p must be in (0, 1]")
    selected, mass = [], 0.0
    for token_id in ranked_ids(probs):
        selected.append(token_id)
        mass += probs[token_id]
        if mass >= threshold:
            break
    return selected

def renormalize(probs, selected):
    allowed = set(selected)
    mass = sum(probs[i] for i in allowed)
    return [p / mass if i in allowed else 0.0
            for i, p in enumerate(probs)]

logits = [0.9546803832054138, 2.0966553688049316,
          -1.067393183708191, -1.1546803712844849,
          -0.929262101650238]
probs = softmax(logits)

kept_k = ranked_ids(probs)[:2]     # [1, 0]
kept_p = top_p_ids(probs, 0.92)   # [1, 0, 4]
after_k = renormalize(probs, kept_k)
after_p = renormalize(probs, kept_p)

assert math.isclose(sum(after_k), 1)
assert math.isclose(sum(after_p), 1)

Output probabilities stay in original ID order, with excluded entries set to zero. A sampled index can therefore be used directly as a token ID. If candidates are compressed into a smaller array instead, its sampled index must be mapped back through the retained-ID list.

Reproduce the complete model

Download the complete five-token model and complete sampling experiment into the same directory. The first includes every fixed weight, the forward pass, gradient checks and recorded numerical values. The second reads its JSON output and uses only the Python standard library.

python gpt_model_interface_demo.py
python gpt_sampling_demo.py

The first experiment requires PyTorch. This verification used Python 3.9.6, PyTorch 2.2.2 and CPU float32. Each script writes a JSON file beside itself. Figures round displayed values; numerical experiments and both demos use unrounded logits. Only the interval calculation and the 100,000 fixed-distribution draws use the explicitly listed four-decimal probabilities.

Checks covered the five GPT-2 IDs and text round trip, embedding-plus-position addition, the causal mask, attention row normalization, output matrix multiplication, cross-entropy and gradients, probability sums and rankings after temperature scaling, top-k and top-p candidate sets, and the 100,000 seeded draws. The source comparison used GPT-2 9b63575 and nanoGPT 3adf61e.

Connect the generation loop

Current ID sequence
  → Model: logits at the last position
  → Greedy; or temperature → filtering → normalization → one draw
  → Append the selected ID
  → Check the end marker or length limit; predict again if continuing

An implementation may combine temperature, top-k and top-p, but their order is part of the algorithm. Each stage reads the distribution produced by the preceding stage. Changing temperature first can change how many candidates top-p needs. The two filters are demonstrated separately here so that each rule remains visible.

Decode through the tokenizer when text is needed for display. A single token may not contain a complete character, so streaming output must handle incomplete byte sequences correctly. Generation stops according to end markers, length limits or other configured rules. That is a different level from BPE encoding’s “no mergeable pair remains” condition.

Input lookup supplies vectors for the current IDs. The model assigns candidate scores under that context. The generation strategy turns those scores into the one ID actually appended at this step.