Language Modeling & Next Tokens

A Transformer forward pass produces logits. Softmax turns that vector of scores into a probability distribution over possible next tokens. Language modeling trains that distribution.

A large language model (LLM) does one repeated job: it predicts the next token. You'll build that job from counts into a trainable objective, then see how a causal Transformer uses it during generation.

One token at a time: An LLM doesn't generate a whole paragraph at once. It predicts a single token, appends it to its input, and runs the process again. This is called autoregressive generation.

---

What is a language model?

An autoregressive language model is a probability engine for token sequences: given the tokens so far, it assigns a probability distribution over the vocabulary for what comes next.

Editor autocomplete is a useful first approximation. A language model doesn't treat every possible token as equally likely; it ranks them using the code, text, and instructions already in the context.

The probability of a whole sequence factorizes into a product of next-token predictions, one per position. This is just the chain rule of probability.

$$P(x_1, x_2, \dots, x_n) = \prod_{t=1}^{n} P(x_t \mid x_1, \dots, x_{t-1})$$

Read it left to right: the chance of the first token, times the chance of the second given the first, times the chance of the third given the first two, and so on. Each factor is one next-token distribution. That single building block, "distribution over the next token given everything before it," is the entire job. Generation, training, and perplexity are all built on it.

Suppose a model receives this prompt: "The function returns the"

Its actual probabilities depend on its tokenizer, weights, and preceding context. In a simplified example, it might rank four possible continuations like this:

Probability table

This table is only an excerpt. Tokens omitted from the table receive the remaining probability mass.

Read the distribution before sampling

The table above isn't a list of words the model "knows." It's the model's next-step belief after reading the prompt.

That distinction matters because the same model can behave differently after one extra token.

The model doesn't store one fixed answer for "the function returns the." It computes a distribution from the whole context. Change the context, and you change the probabilities.

Try the mental move:

1. Read this prompt.
2. Predict three likely next tokens.
3. Compare your predictions with the options below.

Possible continuations:

All three are plausible. A model doesn't need certainty to generate text. It needs a probability ranking, then the decoding algorithm picks from that ranking.

Tiny vocabulary example

A production vocabulary contains many tokens, with its exact size set by the model's tokenizer. That's too many for a beginner example, so shrink the world to five tokens.

Use this prompt:

The compiler reported

The model produces logits, which are raw scores before softmax:

Softmax turns those scores into probabilities:

One line in the code deserves explanation before you read it: logits - logits.max().

Why subtract the maximum logit? Because softmax only cares about differences between logits, not their absolute level. If you subtract the same constant from every score, the ranking and final probabilities stay the same. But the numerics get much safer: the largest shifted logit becomes 0, so its exponential is e^0 = 1 instead of some huge number. On tiny toy logits like [3.0, 1.5, -1.0, ...] this barely matters. On real model logits it prevents overflow and keeps the computation stable.

The exact numbers are less important than the shape:

1. One token is far ahead.
2. Another token still has a real chance.
3. Some tokens are valid but unlikely.
4. The probabilities add to 1.

That's next-token prediction in miniature.

From bigram tables to neural n-grams

The smallest context-aware language model is a bigram model. It looks at one previous token and predicts the next token from a table of counts:

If the code-comment corpus often contains returns -> value, the row for returns assigns high probability to value. If it never contains returns -> exception, an unsmoothed table assigns that pair zero probability.

That table already teaches the core language-modeling contract: read context, produce a probability distribution over the next token. Its weakness is context length. A bigram model can't tell the difference between function returns the and test returns the, because it only sees the final token.

Build a count-based language model

Start with six code-comment fragments. Each adjacent token pair becomes one training event: after returns, which token appeared next?

An unseen pair receives probability zero in this unsmoothed table. More data or smoothing helps, but a count table still can't share what it learned about related contexts.

An n-gram model predicts the next token from the previous n - 1 tokens:

For tiny corpora, count tables still work. For real language, the table explodes because most exact contexts are rare. A neural n-gram model addresses that sparsity by learning embeddings and a small MLP (multi-layer perceptron):

1. Look up embeddings for the previous n - 1 tokens.
2. Concatenate them into one vector.
3. Apply a matrix multiply, bias, and a nonlinearity such as the Gaussian Error Linear Unit (GELU).
4. Project to vocabulary logits.
5. Softmax into next-token probabilities.

This is the bridge from count-based autocomplete to modern LLMs. The objective stays next-token prediction. The model family changes: bigram table, neural n-gram MLP, recurrent network, then Transformer. Each step gives the model a better way to use context.

Train a tiny neural n-gram

The next cell trains a small two-token-context, or trigram, model. It isn't a Transformer, but it exposes the same contract: embeddings and shared weights produce logits, and cross-entropy rewards the true next token.

The tiny corpus is deliberately simple, so this cell demonstrates optimization rather than a production-quality model. The important bridge is that a neural model can use one set of learned parameters across contexts instead of storing an isolated probability row for every exact phrase.

Why is a neural n-gram model more useful than a bigram count table?

Answer

A bigram table only sees one previous token and stores a separate row of counts for each observed pair. A neural n-gram model reads several previous tokens, turns them into embeddings, and uses shared weights to generalize across similar contexts. It still predicts the next token, but it can learn patterns such as "cached value" or "default error" instead of memorizing only exact adjacent pairs.

Temperature changes sampling

A temperature parameter rescales logits before softmax: softmax(logits / T), where T > 0. Values below 1 stretch logit gaps and sharpen the distribution; values above 1 compress gaps and give lower-ranked tokens more probability. Adding the same constant to every logit wouldn't change softmax at all. APIs that expose a zero-temperature mode should implement it as greedy decoding rather than divide by zero.

For logits [3.0, 1.5, -1.0, -2.0, -3.0], the two leading probabilities change directly:

Production note: Low temperature makes source-grounded API-doc answers more repeatable, but it doesn't make unsupported facts true. Higher temperature produces more varied drafts for creative tools. Temperature is one lever among several for shaping generation; a later chapter covers the full decoding toolkit.

Causal language modeling

This specific flavor of language modeling, reading from left to right and predicting the future, is called causal language modeling (or autoregressive modeling). GPT-style models use this objective.^{[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}

It's called "causal" because the model is bound by causality: each position can use earlier context, but it can't peek at future tokens. In a Transformer block, a token position may attend to itself and earlier positions; the mask blocks later positions.

GPT-style decoder models are causal language models.

Masked language modeling, used by Bidirectional Encoder Representations from Transformers (BERT), hides a token and predicts it from both left and right context.^{[3]Reference 3BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.https://arxiv.org/abs/1810.04805} That two-sided objective builds input representations; the left-to-right generation loop needs causal modeling.

---

How the model learns the distribution

We keep saying the model "predicts" the next token. Where does that skill come from? From training on text where the next token is already known. Every sentence in the training data is a free set of labeled examples: at each position, the correct answer is the token that came next.

Teacher forcing

During training the model is fed the real previous tokens at every position, not the tokens it would have guessed. This is called teacher forcing. The model never gets to wander off on its own mistakes while learning; each position is scored against ground truth using the true history.

This distinction matters whenever you interpret a training loss or measure generation latency:

Because the true history is known up front, training scores every position of a sequence in a single forward pass. The causal mask is what makes this honest: the logit produced from prefix x_1 ... x_{t-1} can't read its target token x_t.

This side-by-side view helps lock in why training is parallel while generation isn't:

Shift tokens into inputs and labels

For a decoder, one training sequence supplies several labeled predictions. The input is the sequence shifted right; the target is the same sequence shifted left.

Cross-entropy: the training loss

At each position the model outputs a distribution over the vocabulary. We want the probability it assigns to the true next token to be high. The loss for one position is the negative log of that probability, and the loss for a sequence is the average over positions. This is cross-entropy loss, the same objective introduced in the softmax chapter, now applied at every position.

$$\mathcal{L} = -\frac{1}{n} \sum_{t=1}^{n} \log P(x_t \mid x_1, \dots, x_{t-1})$$

If the model gives the true token probability 1.0, that term contributes -log(1) = 0 loss: a perfect prediction. If it gives the true token a tiny probability, -log of a small number is large: a big penalty. Causal pretraining minimizes this average over a large corpus.

The logits at every shifted position can be scored together. This cell calculates stable cross-entropy for the four targets above.

Perplexity: loss in token units

Cross-entropy is measured in nats (because we used natural log), which is hard to feel. Perplexity is the friendlier form: just exponentiate the average loss.

$$\text{Perplexity} = \exp(\mathcal{L})$$

Perplexity is often interpreted as an effective number of equally likely choices per step. A perplexity of 1 means the model assigned probability 1 to every evaluated token. A model that predicts uniformly over a 50,000-token vocabulary has perplexity 50,000, but the converse doesn't follow: a measured perplexity of 50,000 alone doesn't prove every prediction was uniform or that the model learned nothing. Lower is better when the tokenizer, evaluated data, and averaging convention match.^{[4]Reference 4A Mathematical Theory of Communicationhttps://people.math.harvard.edu/~ctm/home/text/others/shannon/entropy/entropy.pdf[5]Reference 5Perplexity — a Measure of the Difficulty of Speech Recognition Tasks.https://doi.org/10.1121/1.2016299}

Work a tiny example by hand. Suppose the model predicts four tokens, assigning the true token these probabilities: 0.5, 0.5, 0.5, 0.125.

A perplexity of about 2.83 means that, averaged across these four steps, the model behaved as if it were choosing between roughly 2.83 equally likely options. The one bad prediction (0.125) drags the average up. That's the whole point of the metric: it tells you, in token units, how surprised the model was by real text.

When you fine-tune a model and watch the loss curve, you're watching average cross-entropy fall. When an eval report uses the same tokenizer, evaluated tokens, and natural-log averaging, perplexity is exp of that same loss. Hold those choices constant before treating a mismatch as a regression.

A model assigns the true next token a probability of 0.01 at one position. Roughly how much does that position contribute to the loss, and is that good or bad?

Answer

The per-position loss is -log(0.01) which is about 4.6 nats, a large penalty. The model was very surprised by the correct token, so that prediction was bad. A confident, correct prediction (probability near 1.0) would contribute close to 0. Training pushes the model to raise the probability it assigns to the tokens that actually occur, lowering this loss.

---

The autoregressive loop

When you ask a chat model to write an API-doc sentence, it doesn't output the whole message in one step. It generates the message one token at a time in a loop.

This is the autoregressive generation loop:

1. Input: The model receives the prompt.
2. Predict: It calculates the probabilities for the very next token.
3. Pick: We take the most likely token or sample from the distribution, optionally adjusting temperature.
4. Append: That chosen token is glued to the end of the prompt.
5. Repeat: The new, longer context determines the distribution for the next token.

This list describes the dependency between generated tokens, not an inefficient serving implementation. Production decode usually reuses earlier attention state through a key-value (KV) cache. It still performs one sequential decode step for each new token, but it doesn't rebuild earlier keys and values. The serving section below makes that tradeoff concrete.

This explains why generation speed is measured in tokens per second (t/s). The model has to run another decoding step for every generated token. A KV cache avoids recomputing earlier keys and values from scratch, but each new token still has to attend to previous context.

A tiny generation loop in code

This example isn't a neural network. It's a small hand-written language model that shows the same loop: look at current text, choose one next token, append it, and repeat.

Real LLMs replace the hand-written table with a Transformer, but the control flow is the same.

The loop ends when one of these things happens:

Picking a token from the distribution

The loop says "sample a token," but how? The model hands you a distribution, and a decoding strategy turns that distribution into one chosen token. Greedy decoding is the simplest strategy: always take the highest-probability token. It's deterministic, but it can fall into loops, repeating "We apologize for the inconvenience. We apologize for the inconvenience..." because "apologize" keeps scoring highest after "We."

The alternative is to sample: draw a token at random in proportion to its probability, so the second- or third-ranked token sometimes wins. Sampling can reduce repetition and add variety, but careless sampling can also hurt quality. The temperature knob from earlier controls how adventurous that draw is. Greedy versus sampling is the core fork; the full toolkit of decoding methods (and how to tune them) gets a dedicated chapter later, so here we only need the idea that one distribution can be turned into text in more than one way.

Common mistake: Setting a high temperature in factual docs generation and wondering why answers drift away from the source, or setting temperature to 0 in a creative app and wondering why every user gets the same paragraph. At each step, the current context produces a distribution; the decoding strategy decides how much you gamble on it.

Use a seeded random generator when testing a sampler, so samples repeat.

The context window limit

Because the output is constantly being appended to the input, the text the model has to read keeps growing. Late in a 1000-token support answer, the model is reading the original prompt plus the 999 tokens it already generated.

This growing sequence must fit into the model's context window, its short-term memory limit. If the context window is 8,000 tokens and the prompt plus generated text reaches 8,001 tokens, the system has to truncate earlier context, reject the request, summarize, or use a longer-context strategy.

What the KV cache changes

Without a cache, the model would recompute attention information for the full prompt at every generation step. That would be wasteful.

The KV cache stores key and value tensors from earlier tokens at each attention layer. New decode queries can attend to these cached tensors instead of rebuilding earlier keys and values at every step.^{[6]Reference 6Efficiently Scaling Transformer Inference.https://arxiv.org/abs/2211.05102}

Serving systems usually split this into two phases:

• Prefill: process the full prompt in parallel and build the initial KV cache.
• Decode: generate one new token at a time, appending each token and reusing the cache.

Serving systems care so much about KV cache memory because the cache is both an accelerator and a memory bill. It speeds generation, but it also consumes GPU memory proportional to context length, number of layers, key/value-head count, head dimension, and batch size.

Production note: If a model seems cheap at 4K context but expensive at 128K context, check KV cache memory before blaming the model weights. Long context changes serving economics.

Estimate cache memory

A simplified cache estimate makes the tradeoff concrete. For each request, layer, token, and key/value head, the server stores one key vector and one value vector. The estimate uses bfloat16 or float16 values at two bytes each.

This is only the KV cache, not model weights, activations, allocator overhead, or batching metadata. It also shows why fewer key/value heads can reduce cache cost.

---

Why next-token training can learn useful structure

Across large mixtures of code, math, dialogue, and documentation, local word association often isn't enough to minimize next-token loss. Reusable features for syntax, arithmetic, and task structure can improve predictions across many examples.

Worked prompt

Consider this prompt: "A job has a 5-second deadline. After 2 seconds elapsed, the remaining budget is"

A simple Markov chain might look at 2 and guess 3 because 2, 3 is common. A larger neural language model can benefit from features that track subtraction-like structure across many related examples, because those features can improve next-token predictions. This is a pressure from the objective, not a guarantee of correct arithmetic.

Next-token prediction explains the training signal, not human-like understanding or guaranteed reasoning. Memorization can occur, and each claimed capability still needs task-specific measurement.^{[1]Reference 1Improving Language Understanding by Generative Pre-Training.https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf[7]Reference 7Language Models are Few-Shot Learners.https://arxiv.org/abs/2005.14165}

Common failure cases

Treating probabilities as facts

The model's top token can be wrong.

Suppose an internal docs assistant sees:

The rate-limit page says free-tier requests are

A model might assign high probability to unlimited because that phrase is common in marketing copy. But if your current API contract says free-tier requests are capped at 60 per minute, the statistically likely continuation is still wrong.

The fix isn't to hope the base model "knows" your API contract. Put the trusted spec into context, retrieve the right document, or use a tool that checks the source of truth.

The "database" myth

An LLM doesn't query a hidden fact table during generation. Learned parameters can reproduce stale, conflicting, or fabricated continuations. Ground current facts with retrieval, tools, or explicit source context.

Mastery check

Key concepts

• Language-model objective: next-token prediction
• Chain-rule factorization of sequence probability
• Bigram and neural n-gram models share the next-token target
• Teacher forcing: training on true history in one parallel pass
• Cross-entropy loss as averaged negative log-likelihood
• Perplexity equals exp(loss): effective choices per token
• Autoregressive generation loop
• Causal vs masked language modeling (GPT vs BERT)
• Context window limits and KV-cache serving tradeoffs
• The database myth: LLMs generate instead of retrieving facts

Evaluation rubric

• Foundational: Defines an autoregressive language model, writes the chain-rule factorization, and traces the append-and-repeat generation loop.
• Foundational: Explains bigram counts, softmax, temperature, and the difference between greedy decoding and sampling.
• Intermediate: Connects neural n-gram and Transformer logits to teacher forcing, cross-entropy, and perplexity.
• Intermediate: Contrasts causal and masked language modeling, then explains why training is parallel while generation is sequential.
• Intermediate: Estimates KV-cache growth and distinguishes likely continuations from grounded facts.

Follow-up questions

If an LLM just predicts the next token, how can it do math or write code?

Answer

To accurately predict the next token in a vast corpus of human text (which includes technical docs, GitHub repositories, and worked reasoning traces), a model often can't rely only on local word association. It has pressure to build useful internal representations of underlying concepts. To predict the answer to '2+2=', an arithmetic-like internal feature is more useful than memorizing every string. By optimizing for next-token prediction over massive, diverse datasets, reasoning-like behavior can emerge as a byproduct of minimizing loss.

Why does generation slow down as the output gets longer?

Answer

Because of autoregressive generation. To generate the 100th token, the model must condition on the previous 99 tokens. To generate the 101st, it must condition on 100 tokens. KV cache prevents rebuilding earlier keys and values, but attention work and memory still grow with context length.

Why don't we pick the highest-probability token every time?

Answer

Picking the top token every time is called greedy decoding, and it can cause repetition loops. A documentation generator might get stuck writing 'This endpoint returns' over and over because 'returns' keeps scoring highest after 'endpoint.' Sampling instead draws a token in proportion to its probability, introducing enough variance to break loops and produce more natural text. Greedy versus sampling is the core fork; a later chapter covers the full decoding toolkit.

How does training (teacher forcing) differ from generation, and why does it matter for speed?

Answer

During training the model is fed the true previous tokens from the dataset at every position, so the targets are fixed and known in advance. The causal mask lets it predict all positions in a single parallel forward pass. During generation no ground truth exists, so the model must feed its own sampled tokens back in one at a time. Training a sequence can be fast and parallel, while generation stays inherently sequential.

Why can a model with low perplexity on API-doc examples still answer your current rate limit incorrectly?

Answer

Perplexity only measures how well the model predicts held-out text from a distribution similar to training or evaluation data. It doesn't guarantee access to today's source of truth. The model may have learned common API-limit language or older company behavior, so a statistically likely continuation can still be wrong unless you ground the answer with retrieval, tools, or fresh context.

Common pitfalls

• Symptom: A creative demo keeps giving nearly identical completions. Cause: Decoding is too deterministic, often from greedy decoding or very low temperature. Fix: Sample from the distribution and raise temperature carefully instead of always taking the top token.
• Symptom: Training loss looks good, but long generations drift after one bad token. Cause: Teacher forcing was confused with real generation, so evaluation never tested the model feeding on its own outputs. Fix: Run autoregressive generation during eval, not loss-only checks.
• Symptom: A model answers a policy question confidently but cites the wrong company rule. Cause: High next-token probability was treated like evidence, even though the model is generating from patterns instead of querying a live source. Fix: Retrieve the current policy or call a tool before trusting the answer.
• Symptom: Generation gets slower or more expensive as outputs get longer. Cause: The key-value cache saves recomputation but still grows with context length, layers, and batch size. Fix: Watch context budgets and cache memory, not model weight size alone.
• Symptom: Someone claims perplexity and cross-entropy disagree. Cause: They forgot perplexity is exp(loss) under natural-log averaging, or compared runs with different token accounting. Fix: Hold tokenizer, evaluated tokens, and averaging convention constant before comparing them.
• Symptom: Bigram tables, neural n-grams, and GPT-style models feel unrelated. Cause: Architecture changes were mistaken for objective changes. Fix: Trace all three back to the same job: predict the next token from context.