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.

---

Next-token scores, not a whole answer

An autoregressive language model doesn't emit completed answers or prewritten paragraphs in a single computational leap. Instead, it assigns a categorical probability distribution over the vocabulary for the single token that follows the prefix seen so far. A GPT-style large language model (LLM) executes this objective at immense scale.^{[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} Its output at each step is a vector of vocabulary scores, not a fixed sentence.

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

These values are didactic numbers rather than outputs from a specific checkpoint. The final row combines all omitted vocabulary entries so the total probability mass sums to exactly 100%.

Reading the distribution before sampling

The table above doesn't represent static facts stored in a database row. It represents the model's conditional belief given the prompt so far. Add even a single token to the input, and the model's conditional distribution shifts immediately.

The network doesn't memorize one rigid answer for a phrase. Context vectors modulate activations at every layer, yielding a distinct output distribution for every prefix.

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:

$$P(w_1, w_2, \dots, w_n) = \prod_{t=1}^{n} P(w_t \mid w_1, \dots, w_{t-1}) = \prod_{t=1}^{n} P(w_t \mid w_{<t})$$

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 simply $P(w_1)$, conditioned on no prior text. When scoring continuations after a user prompt, each factor conditions on that prompt as well.

This factorization is an exact mathematical equality, not an approximation. It 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:

Multiple candidates compete for probability mass:

All three candidates are grammatically and contextually sound. A language model doesn't require absolute certainty to generate helpful text. It needs calibrated relative probabilities across its vocabulary, leaving the final choice to a decoding algorithm.

Softmax probabilities from raw logits

Real tokenizers span tens of thousands of tokens, typically between 32,000 and 128,000 subwords. To trace the mathematics without clutter, shrink the vocabulary to five candidate tokens.

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:

$$P(w_i) = \frac{\exp(z_i - \max_j z_j)}{\sum_{k=1}^{|V|} \exp(z_k - \max_j z_j)}$$

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.

error commands 80% of the probability mass, followed by warning at roughly 18%. The probabilities sum cleanly to 1.0. Softmax guarantees valid probability math, but it can't tell whether a compiler actually failed. It only normalizes the model's internal scores. 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:

$$P(w_t \mid w_1, \dots, w_{t-1}) \approx P(w_t \mid w_{t-1})$$

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

$$P(w_t \mid w_{t-1}) = \frac{C(w_{t-1}, w_t)}{C(w_{t-1})}$$

Here $C(w_{t-1}, w_t)$ counts how many times the adjacent pair appears, while $C(w_{t-1})$ counts total occurrences of the conditioning token.

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.

The sparsity collapse and Laplace smoothing

The bigram table reveals two severe architectural flaws:

First, Markov amnesia. The bigram model conditions solely on returns. It can't distinguish function returns from test returns, even though our corpus shows that tests frequently return error while functions return value.

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 document contains that valid phrase, taking the negative log-likelihood ($-\log(0)$) yields infinite loss. The entire test corpus probability crashes to zero because of one missing pair!

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

$$P_{\text{Laplace}}(w_t \mid w_{t-1}) = \frac{C(w_{t-1}, w_t) + 1}{C(w_{t-1}) + |V|}$$

If our vocabulary has four potential continuations (value, result, error, exception), the counts become $[4, 3, 2, 1]$, and the denominator becomes $6 + 4 = 10$. exception receives probability $0.100$ instead of zero.

While smoothing prevents infinities, it doesn't solve the core issue. Widening the context to an n-gram model ($P(w_t \mid w_{t-n+1}, \dots, w_{t-1})$) creates an exponential parameter explosion. For a modest vocabulary of $|V| = 50,000$, a 4-gram table demands $|V|^4 = 6.25 \times 10^{18}$ parameters. Most combinations never occur in any training set, and discrete counts can't transfer knowledge between similar words like function and method.

---

Neural language models and distributed representations

Yoshua Bengio and colleagues resolved this combinatorial bottleneck by introducing continuous distributed representations for words.^{[3]Reference 3A Neural Probabilistic Language Modelhttps://www.jmlr.org/papers/v3/bengio03a.html}

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. Project $x$ through shared linear weights $W$ and add a bias vector $b$ to yield vocabulary logits.
4. Pass logits through softmax to form next-token probabilities.

Because words with similar grammatical and semantic roles develop similar embedding orientations, the model generalizes probability mass smoothly. If the training data contains method returns result, the model assigns sensible probability to function returns result without requiring explicit frequency counts for that exact sequence.

Training a tiny neural n-gram with gradient descent

Let's implement a complete neural n-gram model from scratch in Python. 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.

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. An optimal model splits its probability between those two targets, giving a theoretical minimum cross-entropy floor of $\frac{2 \ln(2)}{6} \approx 0.231$. The trained network reaches 0.235, virtually matching the dataset's entropy floor.

Does distinguishing function returns from test returns prove that neural embeddings are strictly necessary?

Answer

No. A discrete trigram table also sees two tokens of context. What neural embeddings provide is parameter sharing: semantically similar words share embedding space. That shared geometry lets neural models generalize to unseen prefixes, which discrete count tables can't do without smoothing.

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:

$$P(w_i) = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$$

• $T < 1.0$ (Cooling): Amplifies logit differences. The top token monopolizes probability mass, making generation focused and deterministic.
• $T = 1.0$ (Unscaled): Standard softmax probabilities derived directly from learned parameters.
• $T > 1.0$ (Heating): Flattens logit differences. The distribution approaches uniform randomness, giving low-probability tokens a higher chance of being sampled.

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 is a sampling control, not a parameter update: it alters decoding entropy without modifying underlying weights.

Causal versus masked language modeling

The architecture that powers modern generative models is causal language modeling (autoregressive generation). It conditions strictly on earlier tokens ($w_{<t}$).

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

Because masked language modeling allows tokens to attend to the right, it can't naturally roll out text token by token. Causal modeling strictly enforces the autoregressive chain rule factorization.

---

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 label is simply 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 rather than its own predictions. This mechanism 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 unlocks massive parallelism:

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. The GPU executes dense General Matrix Multiplies (GEMM) across all $T$ positions simultaneously in an $\mathcal{O}(1)$ parallel forward pass!

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.^{[5]Reference 5Scheduled Sampling for Sequence Prediction with Recurrent Neural Networkshttps://arxiv.org/abs/1506.03099}

During training, position $t$ always conditions on pristine human text from the dataset. During inference, position $t$ conditions on the model's own previously sampled tokens. If the model makes a slight error or samples an unusual token at step 2, it enters an out-of-distribution prefix at step 3. The model was never trained on its own mistaken prefixes, so errors can compound rapidly across long generations.

Because of exposure bias, a low teacher-forced loss doesn't guarantee fluent rollouts. Evaluating a model requires checking both teacher-forced cross-entropy and free-running generations for repetition, drift, and stopping criteria.

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 standard libraries like Hugging Face Transformers, DataCollatorForLanguageModeling with mlm=False takes raw token IDs and masks padding tokens with -100.^{[6]Reference 6Transformers Data Collatorshttps://huggingface.co/docs/transformers/main/en/main_classes/data_collator#transformers.DataCollatorForLanguageModeling} The causal model shifts logits and labels internally: logits[..., :-1, :] pairs with labels[..., 1:].^{[7]Reference 7Transformers Causal Language-Model Losshttps://github.com/huggingface/transformers/blob/main/src/transformers/loss/loss_utils.py}

Cross-entropy loss across sequence positions

At every position $t$, the model emits a probability distribution over the vocabulary. We want the probability assigned to the ground-truth target $x_t$ to be as close to 1.0 as possible. Cross-entropy loss computes the negative log-probability averaged across all $n$ positions:

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

If the network assigns probability 1.0 to the correct token, the term contributes $-\log(1.0) = 0$ loss. If it assigns a tiny probability like 0.01, the penalty is large: $-\log(0.01) \approx 4.605$ nats. Causal pre-training minimizes this average negative log-likelihood across trillions of tokens.

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:^{[8]Reference 8A Mathematical Theory of Communicationhttps://people.math.harvard.edu/~ctm/home/text/others/shannon/entropy/entropy.pdf[9]Reference 9Perplexity — a Measure of the Difficulty of Speech Recognition Tasks.https://doi.org/10.1121/1.2016299}

$$\text{Perplexity} = \exp(\mathcal{L}) = \exp\left(-\frac{1}{n} \sum_{t=1}^n \log P(x_t \mid x_{<t})\right) = \frac{1}{\sqrt[n]{\prod_{t=1}^n P(x_t \mid x_{<t})}}$$

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$).
• If a model guesses uniformly across a vocabulary of 50,000 tokens, its loss is $\ln(50,000) \approx 10.82$, and its perplexity is exactly $\exp(10.82) = 50,000$.
• In our 4-position example above, a mean loss of 0.209 nats yields a perplexity of $\exp(0.209) \approx 1.233$. The model predicts with the certainty of someone picking among roughly 1.2 equally weighted choices per token.

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.

Perplexity is only comparable across models when they share the exact same tokenizer, vocabulary size, and text normalization. A byte-level tokenizer naturally achieves lower perplexity per token than a word-level tokenizer simply because its vocabulary is vastly smaller, not because its predictions are superior.^{[10]Reference 10Perplexity of fixed-length modelshttps://huggingface.co/docs/transformers/perplexity}

A model assigns the true next token a probability of 0.01 at one position. Roughly how much does that position contribute to cross-entropy loss?

Answer

The contribution is -log(0.01) ≈ 4.605 nats. That is a heavy penalty. Confident correct predictions contribute near 0 loss, whereas low-probability surprises drive loss and perplexity up quickly.

---

The autoregressive generation loop

During inference, ground-truth continuations don't exist. The model must generate text one token at a time in a sequential loop:

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.

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.

In a full language model, the transition table isn't a hardcoded dictionary: billions of neural weights compute the distribution dynamically for any arbitrary token prefix.

The loop terminates when one of four conditions is satisfied:

Greedy selection versus temperature sampling

In greedy decoding, the engine always takes the argmax: $\hat{w} = \arg\max_i z_i$. This choice is deterministic, but it often causes repetitive loops in open-ended text.

In stochastic sampling, the engine draws a random token weighted by its softmax probability. Sampling injects diversity, but pure multinomial sampling risks picking ungrammatical tail tokens. Modern serving combines temperature with top-$k$ (filtering to the $k$ most probable tokens) or top-$p$ nucleus sampling (filtering to the smallest set of tokens whose cumulative probability exceeds $p$).

Context window limits and memory growth

Because the model appends every emitted token back into its input, the context length grows at every step. At token 1,000, the model re-evaluates the prompt plus all 999 generated tokens.

Every architecture enforces a maximum context window (such as 8,192 or 128,000 tokens). Exceeding this limit forces the application to truncate historical messages, summarize earlier turns, or reject the request.

---

Production serving: prefill, decode, and the KV cache

Serving an autoregressive Transformer in production introduces a sharp computational divide between two distinct phases:

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

Why naive autoregressive generation is quadratic

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$.

Without caching, generating step $t$ requires re-running linear projections for all $t - 1$ previous tokens. To generate a 2,000-token response from a 1,000-token prompt, the GPU would recompute keys and values for prompt tokens 2,000 times! Over an entire sequence of length $T$, the total operations scale quadratically as $\mathcal{O}(T^2)$.

The KV cache and memory bandwidth bottlenecks

Because causal attention masks prevent future tokens from affecting past states, the Key ($K$) and Value ($V$) vectors of historical tokens never change during generation.

The Key-Value (KV) cache stores the evaluated $K$ and $V$ tensors from all previous positions across all Transformer layers in GPU High Bandwidth Memory (HBM).^{[11]Reference 11Efficiently Scaling Transformer Inference.https://arxiv.org/abs/2211.05102}

At decode step $t$:

1. The GPU only computes $q_t, k_t, v_t$ for the single newly generated token.
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.

This reduces per-step computational complexity from $\mathcal{O}(t)$ to $\mathcal{O}(1)$ projection work and linear $\mathcal{O}(t)$ attention operations.

However, caching introduces a major hardware bottleneck:

• Prefill is compute-bound: Prompt tokens form a large matrix $X \in \mathbb{R}^{S \times d}$. The GPU runs General Matrix Multiplications (GEMM), achieving high arithmetic intensity (FLOPs per byte loaded) and saturating Tensor Cores.
• Decode is memory-bandwidth bound: At each step, the input is a single token vector $x_t \in \mathbb{R}^{1 \times d}$. Computing attention requires executing General Matrix-Vector products (GEMV). The GPU must load gigabytes of weights and cached KV tensors from HBM into SRAM just to perform a handful of operations ($\approx 1\text{--}2$ FLOPs per byte loaded).

In the decode phase, GPU execution units spend most of their time waiting for data to travel across the memory bus.

Sizing the KV cache in high bandwidth memory

The memory footprint of a standard KV cache per request is governed by this exact formula:

$$\text{Memory}_{\text{KV}} = 2 \times L \times H_{\text{KV}} \times d_h \times S \times B \times P_{\text{bytes}}$$

Where:

• $2$: stores both Key and Value matrices
• $L$: number of Transformer layers
• $H_{\text{KV}}$: number of Key/Value attention heads
• $d_h$: head dimension ($d_{\text{model}} / H_Q$)
• $S$: total sequence length (prompt plus generated tokens)
• $B$: concurrent batch size
• $P_{\text{bytes}}$: bytes per value (2 for FP16/BF16, 1 for FP8/INT8)

At 32,768 tokens, our 4-request batch consumes 16 GiB of GPU memory solely for the KV cache. With a concurrent batch size of 32 requests, the cache demands 128 GiB, easily exceeding the memory needed for the model weights themselves!

Two architectural breakthroughs keep modern KV serving manageable:

1. Grouped-Query Attention (GQA): Instead of allocating one KV head per query head ($H_{\text{KV}} = H_Q$), GQA groups query heads to share a smaller number of KV heads (for instance, 32 query heads sharing 8 KV heads). This slashes KV cache memory by 4× with virtually no degradation in model quality.^{[12]Reference 12GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245}
2. PagedAttention (vLLM): Traditional serving engines allocated contiguous memory blocks for the maximum possible context length, wasting 60% to 80% of VRAM on memory fragmentation. PagedAttention partitions the KV cache into fixed-size virtual memory blocks (pages of 16 or 32 tokens), eliminating internal fragmentation and enabling near-optimal memory utilization.^{[13]Reference 13Efficient Memory Management for Large Language Model Serving with PagedAttention.https://arxiv.org/abs/2309.06180}

---

Why next-token objectives learn structured computation

It's common to wonder: if an LLM is only trained to predict the next token, how does it acquire reasoning, code synthesis, and multi-step logic?

Consider this code snippet:

What token comes next?

An n-gram model with a 2-token context only sees retries ==. It might predict 0 because == 0 is common in its training corpus.

A large Transformer reads the full prefix. To drive cross-entropy loss down across millions of similar code files, the network's internal attention heads learn to track variable assignments, decrement operations, and condition checks. Predicting the single token 2 correctly requires internally tracking execution state!

Next-token prediction acts as an information-dense compression objective. To predict the next word across web-scale text, the network must discover circuits that represent syntax, world facts, programmatic states, and conversational structure.^{[14]Reference 14Language Models are Few-Shot Learners.https://arxiv.org/abs/2005.14165}

However, next-token prediction doesn't equal conscious understanding. The model remains a probabilistic engine that mirrors its training distribution.

---

Failure modes and debugging

Next-token prediction supplies a powerful training signal, but it introduces distinct operational failure modes.

Treating probabilities as ground-truth facts

High probability reflects stylistic frequency in the training corpus, not verified truth.

Suppose an internal documentation assistant reads:

The rate-limit page says free-tier requests are

The base model might assign 75% probability to unlimited because marketing copy frequently pairs those words. But if your production API enforces 60 requests per minute, the statistically likely continuation is factually wrong!

The relational database myth

A common misconception treats language models like relational databases that look up rows during generation.

An LLM contains no internal SQL engine or indexed tables. It possesses only numerical weights that parameterize categorical distributions over tokens. When asked for a specific phone number or timestamp, it doesn't query a database; it generates a plausible continuation. If accurate factual recall is required, system designs must provide external sources of truth via retrieval or function calling.

---

Stress-testing prefix distributions and cache budgets

To solidify these concepts, test how prefix evidence and serving dimensions reshape model behavior.

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.

Why can't any neural language model achieve zero loss on a dataset that contains both "function returns value" and "function returns exception"?

Answer

Because identical inputs can't map to conflicting one-hot vectors simultaneously. Softmax outputs must sum to 1.0. When the identical prefix maps to multiple targets, the optimal prediction matches their empirical frequencies, leaving an irreducible entropy floor known as the Bayes error rate.

Next, test the serving economics in kv-cache-memory-budget.py. Halving kv_heads from 8 to 4 (switching to a more aggressive Grouped-Query Attention ratio) reduces cache memory from 16.00 GiB to 8.00 GiB at 32,768 tokens. That 8 GiB savings directly translates to doubling the server's concurrent user capacity.

Language modeling unifies this entire pipeline: from chain-rule factorization and continuous embeddings to teacher-forced training and memory-bandwidth-bound KV cache serving.