Language Modeling & Next Tokens

The previous lesson traced a Transformer from token IDs all the way to unnormalized vocabulary scores. Now an editor prompt ends with The function returns the. How does text turn into a probability distribution over what comes next, and what mechanics govern the loop after the model picks a continuation?

A token is the atomic unit of text in the model's vocabulary. Tokenizers frequently split long words or prefix tokens with spaces, but our running example treats each visible word and punctuation mark as a distinct token. If the model chooses result, the conditioning context grows to The function returns the result. The updated distribution might next favor a period. This append-and-predict cycle is autoregressive generation.

View larger

Next-token scores build a sequence

An autoregressive language model assigns a categorical probability distribution over the vocabulary for the token that follows a prefix. A GPT-style large language model (LLM) learns these conditional predictions from text.^{[1]Reference 1Improving Language Understanding by Generative Pre-Training.https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf[2]Reference 2Language Models are Unsupervised Multitask Learners.https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf} The basic generation loop consumes one such distribution per selected token. Training and advanced serving can compute scores at multiple positions together; we'll distinguish that computation from selecting a complete answer.

Suppose a model reads this four-token prompt:

"The function returns the"

The model's actual probabilities depend on its tokenizer, parameter weights, and earlier conversation context. In a simplified vocabulary slice, it might rank candidate continuations like this:

Next-token distributions after a prompt

The distribution uses didactic values. The final row combines omitted vocabulary entries so the total probability mass sums to exactly 100%.

Reading the distribution before sampling

The table represents the model's conditional probabilities given the prompt so far. Adding a token can shift the distribution immediately.

These rankings and the explanations in the last column are illustrative. Context changes can shift or preserve the ranking, and models can memorize training phrases as well as condition on context.^{[3]Reference 3Extracting Training Data from Large Language Modelshttps://arxiv.org/abs/2012.07805}

Probability theory explains why this local next-token choice suffices to score entire documents. Under the chain rule of probability, the joint probability of any sequence $W = (w_1, w_2, \dots, w_n)$ factorizes exactly into a product of conditional probabilities:

Here $w_t$ represents the token at position $t$, $n$ is total sequence length, and $w_{<t}$ denotes the history of all preceding tokens. The first term is $P(w_1)$ under the chosen starting context, often a beginning-of-sequence token. When scoring continuations after a user prompt, each factor conditions on that prompt as well. Include an end token to score stopping after these tokens.

This factorization is an exact mathematical equality and requires no independence assumptions. The language modeling objective tasks the network with approximating each individual conditional term $P(w_t \mid w_{<t})$.

Consider guessing continuations for this test runner log:

The test runner says the assertion

Multiple candidates compete for probability mass:

passed
failed
holds

All three candidates are grammatically possible. A language model assigns probabilities to them, and a decoding algorithm makes the final choice. Calibration must be measured, and the test runner's output must be checked against its actual report.

Softmax probabilities from raw logits

Tokenizer vocabularies vary by model and can contain tens or hundreds of thousands of entries, including subwords and special tokens. To trace the mathematics without clutter, use a complete toy vocabulary of five candidates.

Take this prompt:

The compiler reported

The network projects its final hidden state through an unembedding matrix to produce logits, which are unconstrained real-valued scores:

The softmax operator converts these real-valued logits into normalized probabilities:

Subtracting $\max_j z_j$ is a numerical stabilization safeguard. As established in the softmax optimization lesson, subtracting the maximum logit keeps the largest exponent at $\exp(0) = 1$, preventing floating-point overflow while preserving exact probability ratios.

import math

tokens = ["error", "warning", "success", "the", "</s>"]
logits = [3.0, 1.5, -1.0, -2.0, -3.0]

peak = max(logits)
exps = [math.exp(logit - peak) for logit in logits]
total = sum(exps)
probs = [exp / total for exp in exps]

for token, prob in zip(tokens, probs):
    print(f"{token:10s} {prob:.3f}")
print(f"sum        {sum(probs):.3f}")
print("top token ", tokens[max(range(len(probs)), key=lambda i: probs[i])])

error      0.800
warning    0.178
success    0.015
the        0.005
</s>       0.002
sum        1.000
top token  error

error commands 80% of the probability mass, followed by warning at roughly 18%. The probabilities sum cleanly to 1.0. Softmax normalizes the model's internal scores; checking whether a compiler failed requires the test runner's report. The next question is how a model acquires those scores from raw text.

Estimating probabilities from observed n-gram counts

Before neural networks dominated natural language processing, language models estimated sequence probabilities by counting token co-occurrences in text corpora.

A bigram model approximates the full conditional history $w_{<t}$ using only the single immediately preceding token $w_{t-1}$. This simplification is the first-order Markov assumption:

Under Maximum Likelihood Estimation (MLE), the bigram transition probability is the ratio of pair frequency to previous-token frequency:

Here $C(w_{t-1}, w_t)$ counts adjacent pairs, and the denominator counts occurrences of the conditioning token with a following prediction target: $C(w_{t-1}) = \sum_{v \in V} C(w_{t-1}, v)$. Our explicit end token ensures that ordinary final words have a successor. We don't model transitions out of <eos>.

Counting bigram events in code comments

We can verify these calculations directly on a toy corpus of six code comments. <bos> and <eos> act as explicit sequence delimiters, enabling the model to learn when continuations start and stop.

from collections import Counter, defaultdict

messages = [
    "function returns value",
    "function returns value",
    "method returns result",
    "handler returns value",
    "test returns result",
    "test returns error",
]

counts = defaultdict(Counter)
for message in messages:
    tokens = ["<bos>", *message.split(), "<eos>"]
    for previous, current in zip(tokens, tokens[1:]):
        counts[previous][current] += 1

row = counts["returns"]
total = sum(row.values())
for token, count in row.most_common():
    print(f"P({token:6s} | returns) = {count}/{total} = {count / total:.3f}")
print("unseen token probability =", row["exception"] / total)

P(value  | returns) = 3/6 = 0.500
P(result | returns) = 2/6 = 0.333
P(error  | returns) = 1/6 = 0.167
unseen token probability = 0.0

The sparsity collapse and Laplace smoothing

The bigram table shows two limits of this count model:

First, limited history. The bigram model conditions solely on returns. It can't distinguish function returns from test returns, even though both function examples end in value, while the two test examples split between result and error.

Second, sparsity collapse. The pair returns exception never appears in our six training sentences. Its raw count is zero, so $P(\text{exception} \mid \text{returns}) = 0.0$. If an evaluation sequence contains that pair, its probability under the unsmoothed model is zero, so its negative log-likelihood is infinite. One unseen pair can make the model assign zero probability to the whole sequence.

Laplace (add-one) smoothing mitigates zero-count collapse by adding pseudo-counts across the entire vocabulary $V$:

For a separate four-output toy row (value, result, error, exception), the smoothed counts are $[4, 3, 2, 1]$, and the denominator is $6 + 4 = 10$. exception receives probability $0.100$. In a complete language model, add-one smoothing assigns mass to every vocabulary token, including continuations never observed after returns.

# Eight ordinary words, an end token, and the extra candidate exception.
prediction_vocab = sorted({token for rows in counts.values() for token in rows} | {"exception"})
smoothed = {token: (row[token] + 1) / (total + len(prediction_vocab)) for token in prediction_vocab}
print("prediction vocabulary:", len(prediction_vocab))
print(f"P(exception | returns): {smoothed['exception']:.4f}")
print(f"row sum: {sum(smoothed.values()):.3f}")

prediction vocabulary: 10
P(exception | returns): 0.0625
row sum: 1.000

An out-of-vocabulary word still needs an unknown-token or subword policy; add-one smoothing doesn't create new vocabulary entries automatically. Add-one is a teaching method. Discounting and backoff methods such as Kneser-Ney often work better for language modeling.^{[4]Reference 4An Empirical Study of Smoothing Techniques for Language Modelinghttps://dash.harvard.edu/handle/1/25104739}

Widening the context to an n-gram model ($P(w_t \mid w_{t-n+1}, \dots, w_{t-1})$) increases the possible combinations exponentially. With $|V| = 50,000$, there are $|V|^4 = 6.25 \times 10^{18}$ possible four-token tuples. Practical count models store sparse observed n-grams and back off or interpolate. This keeps storage manageable, while sparse evidence and limited sharing across words remain challenges.

Neural language models and distributed representations

An influential 2003 model by Yoshua Bengio and colleagues learned continuous word representations jointly with a next-word predictor.^{[5]Reference 5A Neural Probabilistic Language Modelhttps://www.jmlr.org/papers/v3/bengio03a.html} Shared parameters let evidence from one context influence predictions in others.

Instead of treating tokens as orthogonal atomic symbols where one-hot vectors share zero overlap ($e_i^\top e_j = 0$), a neural language model maps each token ID to a continuous embedding vector in $\mathbb{R}^d$:

1. Look up continuous embeddings for the previous $n - 1$ context tokens in an embedding matrix $E \in \mathbb{R}^{|V| \times d}$.
2. Concatenate these vectors into a single feature vector $x \in \mathbb{R}^{(n-1)d}$.
3. Apply a shared predictor to $x$ to yield vocabulary logits. Bengio's model included a nonlinear tanh hidden layer and an optional direct affine route; our smaller lab below keeps only an affine output map.
4. Pass logits through softmax to form next-token probabilities.

our lab: context IDs -> embedding lookup E -> concatenated x -> affine map -> logits -> softmax

Learned representations can share evidence across related contexts without storing a count for each full prefix. This creates potential for generalization: related words may receive similar representations and unseen phrases may receive useful probabilities. The outcome depends on the training data, model capacity, and optimization.

View larger

Training a tiny neural n-gram with gradient descent

The program below is a complete neural n-gram, small enough to read in one sitting. It reads two context tokens, maps each to a 6-dimensional embedding, concatenates them into 12 features, and projects them to 8 vocabulary logits.

The backward pass applies the exact cross-entropy gradient derived in the softmax lesson: $\frac{\partial \mathcal{L}}{\partial z} = p - y$. Gradients flow through the projection weights back into the embedding rows, updating representations through gradient descent.

import math
import random

random.seed(7)
sentences = [
    "function returns value",
    "function returns value",
    "method returns result",
    "handler returns value",
    "test returns result",
    "test returns error",
]
vocab = sorted({word for sentence in sentences for word in sentence.split()})
to_id = {word: index for index, word in enumerate(vocab)}
n_vocab = len(vocab)
emb_dim = 6

contexts = []
targets = []
for sentence in sentences:
    words = sentence.split()
    for index in range(2, len(words)):
        contexts.append([to_id[words[index - 2]], to_id[words[index - 1]]])
        targets.append(to_id[words[index]])

embeddings = [[random.uniform(-0.3, 0.3) for _ in range(emb_dim)] for _ in range(n_vocab)]
weights = [[random.uniform(-0.2, 0.2) for _ in range(n_vocab)] for _ in range(emb_dim * 2)]
bias = [0.0] * n_vocab

def concat_embeddings(token_ids):
    return embeddings[token_ids[0]] + embeddings[token_ids[1]]

def logits_from(features):
    logits = bias[:]
    for i, value in enumerate(features):
        for j, weight in enumerate(weights[i]):
            logits[j] += weight * value
    return logits

def softmax(logits):
    peak = max(logits)
    exps = [math.exp(logit - peak) for logit in logits]
    total = sum(exps)
    return [exp / total for exp in exps]

def mean_loss():
    total = 0.0
    for token_ids, target in zip(contexts, targets):
        logits = logits_from(concat_embeddings(token_ids))
        peak = max(logits)
        total += (peak - logits[target]) + math.log(sum(math.exp(z - peak) for z in logits))
    return total / len(contexts)

def predict(text):
    token_ids = [to_id[word] for word in text.split()]
    logits = logits_from(concat_embeddings(token_ids))
    return vocab[max(range(n_vocab), key=lambda index: logits[index])]

initial_loss = mean_loss()
learning_rate = 0.4
for _ in range(200):
    grad_embeddings = [[0.0] * emb_dim for _ in range(n_vocab)]
    grad_weights = [[0.0] * n_vocab for _ in range(emb_dim * 2)]
    grad_bias = [0.0] * n_vocab
    scale = 1.0 / len(contexts)
    for token_ids, target in zip(contexts, targets):
        features = concat_embeddings(token_ids)
        dlogits = [prob * scale for prob in softmax(logits_from(features))]
        dlogits[target] -= scale
        for j, grad in enumerate(dlogits):
            grad_bias[j] += grad
        for i, value in enumerate(features):
            for j, grad in enumerate(dlogits):
                grad_weights[i][j] += value * grad
        dfeatures = [
            sum(weights[i][j] * dlogits[j] for j in range(n_vocab))
            for i in range(emb_dim * 2)
        ]
        for offset, token_id in enumerate(token_ids):
            start = offset * emb_dim
            for dim in range(emb_dim):
                grad_embeddings[token_id][dim] += dfeatures[start + dim]
    for i in range(n_vocab):
        bias[i] -= learning_rate * grad_bias[i]
        for dim in range(emb_dim):
            embeddings[i][dim] -= learning_rate * grad_embeddings[i][dim]
    for i in range(emb_dim * 2):
        for j in range(n_vocab):
            weights[i][j] -= learning_rate * grad_weights[i][j]

print(f"loss before training: {initial_loss:.3f}")
print(f"loss after training:  {mean_loss():.3f}")
print("after 'function returns':", predict("function returns"))
print("after 'test returns':    ", predict("test returns"))
assert predict("function returns") == "value"
assert mean_loss() < initial_loss

loss before training: 2.145
loss after training:  0.235
after 'function returns': value
after 'test returns':     result

Notice what happened: the model distinguishes function returns from test returns! It outputs value after function returns and result after test returns.

The loss dropped from 2.145 to 0.235 nats. In our dataset, test returns appears twice with conflicting targets: once with result and once with error. Matching that row's empirical distribution gives probability one half to each. The dataset's cross-entropy infimum is $\frac{2 \ln(2)}{6} \approx 0.231$ nats; finite softmax logits approach the boundary probabilities while assigning nonzero mass to the other outputs. The network's 0.235 is close to that floor.

This fit uses six training targets; unseen-prefix accuracy remains unmeasured. Without a nonlinear hidden layer, each logit is a sum of two position-specific embedding contributions and a bias, limiting its ability to express interactions between context tokens. The original nonlinear predictor had more capacity.

Temperature scaling on output logits

During inference, we don't always take the argmax token. We often sample from the distribution. A temperature parameter $T > 0$ rescales logits before applying softmax:

• $0 < T < 1.0$ (Cooling): Amplifies logit differences and increases the maximum probability. Sampling remains stochastic unless filtering or numerical rounding leaves one choice.
• $T = 1.0$ (Unscaled): Standard softmax probabilities derived directly from learned parameters.
• $T > 1.0$ (Heating): Flattens logit differences. For finite logits and a finite vocabulary, the distribution tends to uniform as $T \to \infty$. Evaluate text quality on the target task; temperature alone changes the sampling distribution.

import math

logits = [3.0, 1.5, -1.0, -2.0, -3.0]
tokens = ["error", "warning", "success", "the", "</s>"]
for temperature in [0.5, 1.0, 2.0]:
    peak = max(logits)
    exps = [math.exp((logit - peak) / temperature) for logit in logits]
    total = sum(exps)
    probabilities = [exp / total for exp in exps]
    ranked = sorted(probabilities, reverse=True)
    if temperature == 1.0:
        print(f"{tokens[0]:<11}{probabilities[0]:.3f}")
        print(f"{tokens[1]:<11}{probabilities[1]:.3f}")
    print(f"T={temperature:.1f} top={ranked[0]:.3f} second={ranked[1]:.3f}")

T=0.5 top=0.952 second=0.047
error      0.800
warning    0.178
T=1.0 top=0.800 second=0.178
T=2.0 top=0.575 second=0.272

At $T = 0.5$, error surges to 95.2% probability. At $T = 2.0$, it drops to 57.5%, leaving significant mass for alternatives. Temperature changes the sampling distribution while leaving the model's learned parameters unchanged.

As $T \to 0^+$, a unique maximum takes all mass; tied maxima share it. The formula is undefined at $T=0$. APIs that accept zero temperature usually use a separate greedy-selection convention.

Causal versus masked language modeling

Causal language modeling is an objective: predict a token from earlier tokens ($w_{<t}$). Bigrams, RNNs, and masked decoder Transformers can all implement it. The objective and the neural architecture are different choices.

In contrast, masked language modeling (BERT) hides tokens and reconstructs them using both left and right context.^{[6]Reference 6BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.https://arxiv.org/abs/1810.04805}

BERT predicts selected original tokens from a corrupted sequence. In its original recipe, 15% of positions are selected; of those, 80% become [MASK], 10% become random tokens, and 10% stay unchanged.^{[6]Reference 6BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.https://arxiv.org/abs/1810.04805} Its masked conditionals aren't directly the left-to-right factors above. Iterative filling or an adapted generation setup is possible, but ordinary BERT isn't trained for the same append-one-token loop.

Teacher forcing and parallel causal training

How does a causal model learn from large corpora? Every document provides supervised labels for free. At each token position $t$, the ground-truth target is the observed token $w_{t+1}$.

The causal attention mask and parallel GEMM

During pre-training, the model receives ground-truth tokens at every position. Feeding the target sequence this way is teacher forcing.

In recurrent neural networks (RNNs), hidden state dependencies forced training to run sequentially step by step. In a decoder Transformer, teacher forcing lets the positions in one layer be computed together:

1. The entire token sequence $x_1, \dots, x_T$ enters the network at once.
2. The causal attention mask sets attention logits $M_{ij} = -\infty$ for all positions $j > i$.
3. Because $\exp(-\infty) = 0$, position $t$ can't attend to future tokens $t+1, \dots, T$.
4. Matrix operations process many positions together within each layer; they don't wait for a sampled token at the preceding position. The model's layers still depend on one another.

One sequence pass isn't constant work or constant latency. For fixed model widths, full self-attention does $\mathcal{O}(T^2)$ score/value arithmetic per layer; projections and feed-forward layers also do work for every token. Parallelism reduces sequential dependencies across positions, and finite hardware determines how much of that work can run together.^{[7]Reference 7Attention Is All You Need.https://arxiv.org/abs/1706.03762}

View larger

Exposure bias and the train-inference gap

Teacher forcing provides a stable, clean training signal, but it creates a fundamental train-inference mismatch known as exposure bias.^{[8]Reference 8Scheduled Sampling for Sequence Prediction with Recurrent Neural Networkshttps://arxiv.org/abs/1506.03099}

Under ordinary teacher-forced training, each prediction conditions on an observed dataset prefix, which needn't be pristine or human-written. In a rollout, the continuation includes the model's selected tokens. A mistake can lead to an unfamiliar prefix and compound, but an unusual sample isn't automatically out of distribution, and the same prefix may have appeared in training. The scheduled-sampling paper studies this possible mismatch in recurrent sequence tasks.^{[8]Reference 8Scheduled Sampling for Sequence Prediction with Recurrent Neural Networkshttps://arxiv.org/abs/1506.03099}

Low teacher-forced loss can coexist with poor free-running generations. Evaluate teacher-forced cross-entropy and generated rollouts for fluency, repetition, drift, and stopping behavior.

Shifting tokens into inputs and target labels

Training a causal decoder requires aligning inputs and targets. Given the sequence [<bos>, function, returns, value, <eos>], input tokens drop the final token, while target labels drop the initial <bos>:

In Hugging Face Transformers 5.6.2, DataCollatorForLanguageModeling(mlm=False) copies input IDs into labels and sets entries matching the padding-token ID to -100.^{[9]Reference 9Transformers Data Collatorshttps://huggingface.co/docs/transformers/main/en/main_classes/data_collator#transformers.DataCollatorForLanguageModeling} The default causal loss shifts labels internally. The versioned source, verified October 10, 2026, pads the label row with a final ignored entry and uses labels[..., 1:]; this is equivalent to pairing all but the last logit with all but the first original label.^{[10]Reference 10Transformers Causal Language-Model Losshttps://github.com/huggingface/transformers/blob/main/src/transformers/loss/loss_utils.py} Don't pre-shift those labels again. The code below instead prepares explicitly shifted arrays for a loss that performs no further shift.

Ignoring a loss target and masking an attention key are separate operations. Also check padding aliases: if pad_token_id == eos_token_id, this collator's token-ID comparison ignores genuine end tokens too. To teach stopping in that setup, construct labels that preserve real end tokens and ignore actual padding positions, for example using an accurate attention mask. Packed documents and instruction-only loss masks need their own explicit boundary and supervision policy.

vocab = {"<bos>": 0, "function": 1, "returns": 2, "value": 3, "<eos>": 4}
tokens = ["<bos>", "function", "returns", "value", "<eos>"]
ids = [vocab[token] for token in tokens]
inputs = ids[:-1]
targets = ids[1:]

print("input token ids: ", inputs)
print("target token ids:", targets)
for position in range(len(inputs)):
    context = " ".join(tokens[: position + 1])
    print(f"{context:25s} -> {tokens[position + 1]}")

input token ids:  [0, 1, 2, 3]
target token ids: [1, 2, 3, 4]
<bos>                     -> function
<bos> function            -> returns
<bos> function returns    -> value
<bos> function returns value -> <eos>

Cross-entropy loss across sequence positions

At each supervised position, cross-entropy loss scores the probability of the observed target token. Averaging across $n$ scored targets gives:

If the network assigns probability 1.0 to the target, the term contributes $-\log(1.0) = 0$ loss. At probability 0.01, the penalty is $-\log(0.01) \approx 4.605$ nats. Across conflicting observations of the same prefix, minimizing average loss favors their frequency distribution, not confidence 1.0 in each incompatible target. Exclude ignored targets from both the loss sum and $n$; require at least one scored target.

import math

targets = [1, 2, 3, 4]  # function, returns, value, <eos>
logits = [
    [0.0, 3.0, 0.5, -1.0, -2.0],
    [-1.0, 0.1, 2.5, 0.0, -2.0],
    [-1.0, 0.0, 0.2, 2.0, -1.5],
    [-2.0, 0.0, -1.0, -0.5, 2.4],
]

true_probs = []
losses = []
for row, target in zip(logits, targets):
    peak = max(row)
    exps = [math.exp(logit - peak) for logit in row]
    total = sum(exps)
    prob = exps[target] / total
    true_probs.append(prob)
    losses.append((peak - row[target]) + math.log(total))

print("true-token probabilities:", [round(prob, 3) for prob in true_probs])
print("per-token losses:        ", [round(loss, 3) for loss in losses])
print(f"mean cross-entropy:       {sum(losses) / len(losses):.3f}")

true-token probabilities: [0.864, 0.824, 0.724, 0.839]
per-token losses:         [0.146, 0.194, 0.323, 0.175]
mean cross-entropy:       0.209

Perplexity as an effective branching factor

Cross-entropy loss in natural logarithms produces units of nats (or bits if using base 2). Exponentiating cross-entropy yields perplexity:^{[11]Reference 11A Mathematical Theory of Communicationhttps://people.math.harvard.edu/~ctm/home/text/others/shannon/entropy/entropy.pdf[12]Reference 12Perplexity — a Measure of the Difficulty of Speech Recognition Tasks.https://doi.org/10.1121/1.2016299}

Perplexity is the reciprocal of the geometric mean of the target token probabilities.

Intuitively, perplexity represents the effective branching factor: the number of equally likely options the model is choosing between at each step.

• A perplexity of 1.0 means the model predicted every target token with 100% confidence ($\mathcal{L} = 0$).
• A uniform predictor over 50,000 tokens has loss $\ln(50,000) \approx 10.82$ and perplexity exactly 50,000. Exponentiating the rounded 10.82 gives only an approximation.
• Our four-position example has perplexity about 1.233. This score summarizes target surprisal; confidence calibration needs a separate evaluation. Perplexity can exceed vocabulary size if a model assigns sufficiently small probabilities to the observed targets.

import math

# Probability the model gave to the token that was actually correct, per position
true_token_probs = [0.5, 0.5, 0.5, 0.125]

per_position_loss = [-math.log(prob) for prob in true_token_probs]
avg_loss = sum(per_position_loss) / len(per_position_loss)
perplexity = math.exp(avg_loss)

print(f"per-position loss: {[round(loss, 4) for loss in per_position_loss]}")
print(f"average loss:      {avg_loss:.4f}")
print(f"perplexity:        {perplexity:.4f}")

product = 1.0
for prob in true_token_probs:
    product *= prob
geo_mean = product ** (1 / len(true_token_probs))
print(f"1 / geometric mean:{1 / geo_mean:.4f}")

per-position loss: [0.6931, 0.6931, 0.6931, 2.0794]
average loss:      1.0397
perplexity:        2.8284
1 / geometric mean:2.8284

Notice how one low-probability target (0.125 at position 4) pulls the average loss from 0.693 to 1.039, elevating perplexity to 2.828. Perplexity penalizes surprise heavily.

Compare token perplexity under the same tokenization, evaluation text, preprocessing, scored targets, and context policy. Vocabulary size does not make tokenization units comparable. Byte tokens and word tokens represent different prediction units, so use per-token perplexity only for like tokenizations. For a common cross-tokenizer unit, compare negative log-likelihood per input byte (bits per byte) under a stated encoding and boundary policy.^{[13]Reference 13Perplexity of fixed-length modelshttps://huggingface.co/docs/transformers/perplexity}

Long-document evaluation also depends on how context is supplied. Disjoint chunks reset context and usually make predictions harder near each boundary; strided windows can preserve preceding context. Score each target once, then pool total negative log-likelihood and target count before exponentiating. Averaging per-batch perplexities generally gives a different result. Ordinary BERT masked-token loss doesn't define this autoregressive sequence perplexity.^{[13]Reference 13Perplexity of fixed-length modelshttps://huggingface.co/docs/transformers/perplexity}

For example, one evaluation batch scores a single target at probability 0.5; another scores three targets at probability 0.125 each. The second batch contributes three times as many scored targets, even if both batches occupy the same padded tensor shape.

import math

batches = [[0.5], [0.125, 0.125, 0.125]]
nll_sums = [sum(-math.log(p) for p in batch) for batch in batches]
counts = [len(batch) for batch in batches]
means = [nll / count for nll, count in zip(nll_sums, counts)]

pooled = math.exp(sum(nll_sums) / sum(counts))
mean_batch_loss_ppl = math.exp(sum(means) / len(means))
mean_batch_ppl = sum(math.exp(mean) for mean in means) / len(means)
print(f"pooled token perplexity:   {pooled:.3f}")
print(f"exp(mean batch loss):      {mean_batch_loss_ppl:.3f}")
print(f"mean batch perplexity:     {mean_batch_ppl:.3f}")
assert math.isclose(pooled, 1 / (0.5 * 0.125 ** 3) ** 0.25)

pooled token perplexity:   5.657
exp(mean batch loss):      4.000
mean batch perplexity:     5.000

Only the pooled result measures surprise per scored token across these four observations. Taking a plain mean of batch losses gives each batch equal weight, so the single easy target counts as much as all three harder targets. Taking a mean of batch perplexities introduces a further difference because exponentiation is nonlinear. Retain the loss sum and count after excluding padding, prompt-only labels, and other ignored targets.

The autoregressive generation loop

In basic autoregressive decoding, the model selects a token before it can score the following token:

1. Input: Feed the prompt tokens into the network.
2. Predict: Compute vocabulary logits and softmax probabilities for the next token.
3. Select: Pick a token via greedy argmax or temperature-scaled sampling.
4. Append: Concatenate the newly selected token to the input sequence.
5. Repeat: Feed the expanded context back into the model to predict the subsequent token.

View larger

A deterministic generation loop in code

This small script demonstrates the sequential mechanics: read context, pick the top token, append it, and check for a stop token.

next_token_table = {
    "The function returns the": [("result", 0.60), ("value", 0.25), ("index", 0.10), ("other", 0.05)],
    "The function returns the result": [(".", 0.80), ("and", 0.20)],
    "The function returns the result.": [("</s>", 1.0)],
}

text = "The function returns the"

for step in range(3):
    choices = next_token_table[text]
    assert abs(sum(probability for _, probability in choices) - 1.0) < 1e-12
    token, _probability = max(choices, key=lambda item: item[1])

    if token == "</s>":
        break

    text = text + ("" if token == "." else " ") + token
    print(step, repr(text))
print("final text:", repr(text))

0 'The function returns the result'
1 'The function returns the result.'
final text: 'The function returns the result.'

This dictionary is a mechanics fixture. A neural language model computes scores from learned parameters for tokenized prefixes within its supported input range. The real tokenizer decodes tokens into text; this script uses a space-and-period shortcut.

Common reasons to stop include:

Greedy selection versus temperature sampling

In greedy decoding, the engine takes the argmax: $\hat{w} = \arg\max_i z_i$. Fixed logits and a fixed tie rule give a deterministic choice. Changing hardware or kernels can perturb close logits, so check repeatability across deployment configurations. Greedy selection can repeat itself and doesn't maximize whole-sequence probability.

For a two-token example, let $P(A)=0.6$, $P(B)=0.4$, and the most likely second token have probability 0.1 after A but 0.9 after B. Greedy starts with A and reaches joint probability $0.6 \times 0.1 = 0.06$; the B continuation reaches $0.4 \times 0.9 = 0.36$. The locally best first choice can lead to a lower-probability sequence. This hypothetical A branch needs at least ten second-token candidates, so its vocabulary contains more than the two first-step choices.

In stochastic sampling, the engine draws a random token weighted by its softmax probability. Unfiltered sampling can pick an unsuitable tail token. Top-$k$ retains the $k$ highest-probability tokens; top-$p$ nucleus sampling retains a highest-probability set whose cumulative mass reaches the threshold $p$, subject to the implementation's tie and minimum-token rules. Renormalize the retained probabilities before sampling. Filtering changes the selected-token distribution; the unfiltered 60% table no longer directly describes selection frequencies.

import math
import random

tokens = ["result", "value", "index"]
logits = [2.4, 1.7, 0.2]

def probabilities(temperature):
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be finite and positive")
    peak = max(logits)
    exps = [math.exp((logit - peak) / temperature) for logit in logits]
    total = sum(exps)
    return [exp / total for exp in exps]

rng = random.Random(4)
greedy = tokens[max(range(len(logits)), key=lambda i: logits[i])]
probs = probabilities(1.4)
samples = [rng.choices(tokens, weights=probs, k=1)[0] for _ in range(8)]

print("greedy:", greedy)
print("sampled at temperature 1.4:", samples)
print("probabilities:", [round(prob, 3) for prob in probs])

greedy: result
sampled at temperature 1.4: ['result', 'result', 'result', 'result', 'result', 'result', 'index', 'value']
probabilities: [0.551, 0.334, 0.115]

Context window limits and memory growth

In the basic loop, selecting a token extends the logical context. When predicting generated token 1,000, the model conditions on the prompt and preceding 999 generated tokens. A KV cache reuses earlier computations, while each new token still requires a model pass for its new position.

Models and serving runtimes specify supported context windows. Full-attention cache requirements usually grow with retained context; sliding-window layers or recurrent models use different retention rules. Before exceeding a deployment's limit, reserve room for output and choose a documented truncation, summary, or rejection policy. Removing old cached entries changes the retained context; cached hidden states may still encode the removed tokens.^{[14]Reference 14KV Cache Strategieshttps://huggingface.co/docs/transformers/main/en/kv_cache} A cache avoids recomputing the prefix and stores the information needed to attend to it during decoding.

Production serving: prefill, decode, and the KV cache

Serving keeps the same autoregressive model but puts it through two different kinds of work: processing the prompt, then generating tokens with earlier outputs as context.

1. The Prefill Phase: Processing the user's prompt tokens all at once.
2. The Decode Phase: Generating new tokens sequentially one by one.

View larger

Count the work before and after caching

In standard multi-head self-attention, token $t$ computes queries $Q_t = X_t W_Q$, keys $K_t = X_t W_K$, and values $V_t = X_t W_V$.

Let $p$ be prompt length and $m$ the number of generated tokens. Without caching, a simple full-prefix implementation processes lengths $p, p+1, \dots, p+m-1$. Projection and feed-forward work sums over those lengths. Full self-attention work sums over their squares. For fixed $p$ and model widths, those totals grow as $\mathcal{O}(m^2)$ and $\mathcal{O}(m^3)$ respectively. Calling all uncached generation quadratic misses the repeated full-attention computation.

The KV cache and memory bandwidth bottlenecks

For a fixed causal Transformer in evaluation mode, unchanged prefix tokens and position assignments produce the same mathematical historical keys and values. Future appends can't change earlier states. This is the reuse invariant behind caching. Changing the model, prefix, positions, or relevant attention policy invalidates it; numerical differences between kernels can still affect results.

The Key-Value (KV) cache stores layer-specific key and value tensors, not raw word embeddings. GPU serving commonly keeps them in High Bandwidth Memory (HBM), while quantized, offloaded, and sliding-window caches use other storage or retention policies.^{[15]Reference 15Efficiently Scaling Transformer Inferencehttps://proceedings.mlsys.org/paper_files/paper/2023/hash/c4be71ab8d24cdfb45e3d06dbfca2780-Abstract-mlsys2023.html[14]Reference 14KV Cache Strategieshttps://huggingface.co/docs/transformers/main/en/kv_cache}

Prefill scores the first generated token. For each subsequent basic decode call:

1. Feed the token just selected from the previous call. Each layer computes its new query, key, and value and runs the other block operations for that position.

2. It appends $k_t$ and $v_t$ to the cached tensors in HBM.

3. It computes attention by multiplying query vector $q_t$ against all accumulated cached keys and values.

4. The final head supplies the scores used to select the following token.

After prefill, projection and feed-forward work per added position is independent of retained sequence length for a fixed model. Full-attention work per position remains linear in the number of visible cached positions. Cache reuse avoids recomputing old projections, while cumulative full-attention work remains quadratic across a growing continuation.

Check the cache against a full pass

Here is the reuse invariant for one attention layer. Fixed input states stand in for positioned layer inputs. We prefill three positions, then feed positions four and five separately. Each new query is at the end of its cache, so all stored keys are allowed. Using nonsquare is_causal=True here would instead align the query to the start of the key row, as the previous Transformer lesson showed.

import torch
import torch.nn.functional as F

torch.manual_seed(3)
hidden = torch.randn(1, 5, 4)
q, k, v = [hidden @ torch.randn(4, 4) for _ in range(3)]
full = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)

cache_k, cache_v = k[:, :3].clone(), v[:, :3].clone()
outputs = [F.scaled_dot_product_attention(q[:, :3], cache_k, cache_v,
                                        is_causal=True, dropout_p=0.0)]
lengths = [cache_k.size(1)]
for position in range(3, 5):
    cache_k = torch.cat([cache_k, k[:, position:position + 1]], dim=1)
    cache_v = torch.cat([cache_v, v[:, position:position + 1]], dim=1)
    allowed = torch.ones(1, cache_k.size(1), dtype=torch.bool)
    outputs.append(F.scaled_dot_product_attention(q[:, position:position + 1],
                   cache_k, cache_v, attn_mask=allowed, dropout_p=0.0))
    lengths.append(cache_k.size(1))

print("cached lengths:", lengths)
print("matches full pass:", torch.allclose(torch.cat(outputs, dim=1), full, atol=1e-6))
assert torch.allclose(torch.cat(outputs, dim=1), full, atol=1e-6)

cached lengths: [3, 4, 5]
matches full pass: True

In a complete decoder, each attention layer needs its own cache, and position IDs must continue from the processed prefix. Cached keys already include that layer's positional treatment, such as RoPE. Prefill and decode still execute normalization, feed-forward blocks, and the vocabulary head; only reusable history is skipped.

How does this change accelerator utilization?

• Prefill often has higher arithmetic intensity: Many prompt positions reuse loaded weight matrices, making large matrix multiplies more likely to be compute-limited.
• Small-batch decode often hits memory bandwidth: One position per sequence does relatively little arithmetic per loaded weight. In a simplified batch-one, two-byte dense weight multiply, one multiply-plus-add gives about 2 FLOPs per 2-byte weight, or 1 FLOP/byte. Attention also reads visible KV data. This estimate isn't the arithmetic intensity of every decode kernel.

Batches reuse weights across requests and can increase arithmetic intensity. Context length, model size, quantization, accelerator hardware, kernel fusion, and distributed communication change the bottleneck. Treat these as workload-dependent tendencies and benchmark the actual system.^{[15]Reference 15Efficiently Scaling Transformer Inferencehttps://proceedings.mlsys.org/paper_files/paper/2023/hash/c4be71ab8d24cdfb45e3d06dbfca2780-Abstract-mlsys2023.html} Speculative decoding can verify several proposed tokens together and accept a block, while preserving the target autoregressive distribution under the exact algorithm.^{[16]Reference 16Fast Inference from Transformers via Speculative Decoding.https://arxiv.org/abs/2211.17192}

Sizing the KV cache in high bandwidth memory

For uniform full-attention layers with equal key/value head dimensions and precision, the raw KV tensor payload for a batch is:

Where:

• $2$: stores both Key and Value matrices
• $L$: number of Transformer layers
• $H_{\text{KV}}$: number of Key/Value attention heads
• $d_h$: stored key/value head dimension (often $d_{\text{model}} / H_Q$, but architecture-dependent)
• $S$: number of cached positions per sequence
• $B$: concurrent batch size
• $P_{\text{bytes}}$: payload bytes per stored value (2 for FP16/BF16, 1 for FP8/INT8)

This assumes each request has the same cached length and each layer keeps every position. For variable lengths, replace $SB$ with their summed cached lengths; for differing layers, sum their separate payloads. Quantization scales, page rounding, allocation reserves, and other model/runtime tensors add memory beyond this formula. A token just selected hasn't necessarily been fed back and cached yet.

batch = 4
layers = 32
kv_heads = 8
head_dim = 128
bytes_per_value = 2

def cache_gib(sequence_length):
    values = batch * layers * sequence_length * kv_heads * head_dim * 2
    return values * bytes_per_value / (1024 ** 3)

for sequence_length in [2_048, 32_768]:
    print(f"{sequence_length:,} tokens: {cache_gib(sequence_length):.2f} GiB")

2,048 tokens: 1.00 GiB
32,768 tokens: 16.00 GiB

At 32,768 cached positions each, our four-request batch has 16 GiB of raw KV payload. At batch size 32, that becomes 128 GiB. Compare it with the chosen model's weight footprint and runtime overhead; layer/head counts alone don't tell us how much memory the weights occupy.

Two architecture choices keep the KV payload practical:

1. Grouped-Query Attention (GQA): Query-head groups share fewer KV heads. Changing 32 KV heads to 8 reduces raw KV payload 4× at equal layer count, head dimension, precision, and cached length. The original paper obtained quality close to multi-head attention in its uptrained T5 experiments; quality at other model and head-ratio settings depends on the model and training recipe. Changing a checkpoint's head count requires a compatible architecture and training/conversion strategy.^{[17]Reference 17GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245}
2. PagedAttention (vLLM): KV blocks can occupy noncontiguous physical storage and be allocated as needed. The original paper measured only 20.4%–38.2% useful token-state occupancy in the compared systems' allocated KV memory, with waste from reserves and fragmentation. This wasn't a universal percentage of all GPU VRAM. Equal block sizes avoid external fragmentation, while partial final blocks can still waste space. Block sizes are implementation/workload choices; the paper evaluated their tradeoffs.^{[18]Reference 18Efficient Memory Management for Large Language Model Serving with PagedAttention.https://arxiv.org/abs/2309.06180}

How next-token predictions can reflect structure

A natural question: if an LLM predicts one token at a time, where do its code and reasoning abilities come from?

Consider this code snippet:

retries = 3
retries -= 1
print(retries)  # output:

What token comes next?

Under a word/token simplification, a short n-gram sees only the end of the comment and misses the assignment and decrement. Its prediction depends on local counts.

A model with sufficient context could predict 2, matching the printed value. Cross-entropy on execution examples can reward state-sensitive predictions, while a pattern shortcut can also produce this answer. Change the initial value and operation, then test unfamiliar combinations to distinguish state tracking from a lucky completion.

Next-token training rewards predictive representations of structure in text. GPT-3 demonstrated task performance from in-context examples; the study measured behavior, while internal mechanisms and reasoning accuracy require separate evidence.^{[19]Reference 19Language Models are Few-Shot Learners.https://arxiv.org/abs/2005.14165} Instruction tuning, preference optimization, and other post-training steps can further change model behavior; the next lesson examines that progression.

A correct completion on this fixture is an initial signal to test against unfamiliar code and inputs.

Failure modes and debugging

Next-token prediction gives you a label at every position. The same objective also creates distinct operational failures.

Treating probabilities as ground-truth facts

High token probability expresses a conditional prediction under the model. It can reflect syntax, supplied evidence, learned facts, and post-training behavior; check factual claims against their sources.

Suppose an internal documentation assistant reads:

The rate-limit page says free-tier requests are

As a hypothetical failure fixture, suppose the model assigns 75% probability to unlimited, while the actual policy is 60 requests per minute. That continuation is wrong. The probability and limit are illustrative inputs; diagnose the error by checking the prompt, supplied policy evidence, model, and generation settings.

Why an LLM isn't a live database

To retrieve live data, connect the model to a database or retrieval tool; next-token generation uses supplied context and learned parameters.

The neural model itself doesn't execute a SQL lookup when it predicts a token. Its weights can encode facts and even reproduce exact training passages, including memorized personal details.^{[3]Reference 3Extracting Training Data from Large Language Modelshttps://arxiv.org/abs/2012.07805} During inference, activations and caches also carry information from the current input. An application can add retrieval or tool access to query an actual database.

For facts that require current, traceable accuracy, provide an authoritative source and check the answer against it. Retrieval and tool calling improve access to evidence; their value depends on source quality and whether the model applies the evidence correctly. A plausible continuation reflects model scoring. Check authoritative evidence for current truth, and inspect the application trace to confirm a live lookup.

Stress-testing prefix distributions and cache budgets

Two small lab edits show how added evidence changes a next-token distribution, and how lower KV storage leaves throughput unanswered.

In count-bigram-events.py, add "function returns exception" to the training list. The new row for returns has 7 total observations: value: 3, result: 2, error: 1, exception: 1. The probability of value drops from $0.500$ to $3/7 \approx 0.429$, and exception takes $1/7 \approx 0.143$.

Now consider what happens in tiny-neural-ngram.py. The prefix function returns now appears twice with target value and once with target exception. A deterministic function can't output two different probabilities of 1.0 for the same input. The cross-entropy objective will drive the network toward assigning roughly 67% probability to value and 33% to exception.

Next, test the arithmetic in kv-cache-memory-budget.py. Halving kv_heads from 8 to 4 reduces this batch's raw payload from 16.00 GiB to 8.00 GiB at 32,768 cached positions. If KV payload alone limits capacity and the usable KV budget is unchanged, you could fit twice as many equal-length requests. Actual concurrency also depends on weight/runtime memory, allocation overhead, scheduling, compute, and latency targets. Changing a number in the calculator doesn't convert a trained checkpoint to a new attention architecture.

Factor a document into next-token probabilities and train that loss on the observed tokens. A KV cache reuses earlier computations. A high probability measures next-token likelihood; verify factual claims against source evidence. The next lesson is how that same loop became an assistant.