From a prefix to a probability distribution

A model does not receive words as indivisible objects. A tokenizer first maps text to a sequence of integer token identifiers. A token may be a word, part of a word, punctuation or a byte-level fragment, depending on the tokenizer. The model processes those identifiers and produces one score, called a logit, for every token in its vocabulary. A softmax operation converts the logits into non-negative probabilities that sum to one.

For a token sequence $t_1, t_2, \ldots, t_n$, the model estimates

$$P(t_n \mid t_1, t_2, \ldots, t_{n-1}).$$

The notation says that the probability assigned to the next token is conditional on the available prefix. During generation, a decoding rule selects a token from that distribution, appends it to the prefix, and runs the model again. During pretraining, the token that actually followed the prefix supplies the target. No separate annotator has to label each example.

Consider a toy vocabulary after the prefix:

The credit committee approved the

Suppose the model assigns these probabilities:

If the observed next token in the training text is facility, the negative-log-likelihood loss for this position is

$$-\ln(0.46) \approx 0.78.$$

Had the model assigned facility a probability of 0.10, the loss would have been $-\ln(0.10) \approx 2.30$. Training adjusts the parameters so that observed continuations tend to receive more probability in similar contexts. The loss is calculated at many positions in a batch, averaged, and propagated back through the network. A single example contributes little; a large and varied corpus supplies an immense collection of such corrections.

The same distribution supports several generation policies. Greedy decoding would choose facility, the highest-probability token. Sampling could choose proposal or request, preserving variation at the cost of occasional poor choices. Temperature and top-$k$ or nucleus sampling reshape the distribution before selection. These policies change what is emitted, but they do not change the model’s parameters.

There are two important consequences.

First, a generated sentence is not retrieved as a complete record. It is assembled token by token. Memorised passages can influence generation, and training-data reproduction is a genuine risk, but the mechanism at inference time is conditional prediction rather than a database lookup.

Second, probability is not truth. A token can be locally plausible and factually wrong. The model is rewarded during pretraining for matching the continuations found in its data, not for consulting an authoritative source or proving each claim. Fluent wording is therefore weak evidence of factual reliability.

What next-token prediction can teach

To reduce prediction error across diverse text, a model benefits from representing recurring structure. Subject–verb agreement helps it choose a grammatical continuation. Knowledge of document genres helps it distinguish a contract clause from a recipe. Regularities in source code help it close a function or infer a likely variable name. Relations expressed repeatedly in text help with some factual questions. None of these skills needs a dedicated label during pretraining; each can improve the probability assigned to observed tokens.

This account is stronger than saying that the model merely memorises, but weaker than saying that it understands as a person does. A trained network contains distributed numerical representations that support useful generalisation. It can combine patterns in a new prompt, follow examples supplied in context and produce sequences absent from the training corpus. It can also fail on a small change of wording, invent a source, mishandle arithmetic or continue a false premise. Behaviour varies with the data, model, prompt, decoding method and evaluation.

The word understanding hides several possible claims. If it means behavioural competence on a defined task, a model may satisfy the criterion and still be unreliable outside the test distribution. If it implies human experience, intentions or consciousness, next-token performance does not establish it. Current evidence does not justify resolving the philosophical question by assertion. For engineering purposes, we can measure capabilities and failure modes without assigning a human mental state to the system.

Claims about emergent abilities require similar care. Wei and peers used the term for abilities that appear in larger models but are absent from smaller ones, with the transition not predictable by a simple extrapolation.2 Schaeffer and peers later showed that nonlinear or discontinuous metrics can make a smooth change in underlying model outputs look like an abrupt threshold.3 Both observations can matter: larger training runs can yield qualitatively useful behaviour, while the apparent sharpness of a benchmark transition may depend on how performance is measured. We will use emergent only when the measurement and comparison are stated, rather than as a synonym for impressive.

In-context learning illustrates the distinction. A prompt may contain two translation examples followed by a third source sentence. A sufficiently capable model can continue the demonstrated pattern without a gradient update. Its behaviour has adapted to the prompt, but its weights have not changed. Calling that response “learning” is conventional; calling it fine-tuning would be incorrect.

Locating language models within AI

The terminology surrounding language models is easier to use when its boundaries are explicit.

Artificial intelligence is the broad category of systems designed to perform tasks associated with intelligent behaviour. It includes hand-written rules, search, planning, optimisation and learned models. Machine learning is a subset in which behaviour is fitted from data. Deep learning is a subset of machine learning built from multi-layer neural networks. A large language model is a deep neural model trained on language at substantial scale.

Generative AI cuts across data types. It describes systems that generate content such as text, images, audio or code. A text-generating LLM therefore belongs to both deep learning and generative AI. A vision classifier is deep learning but not a language model; an image diffusion model is generative AI but not an LLM.

This taxonomy prevents a common design error: treating “use AI” as a technical requirement. It is not. A requirement should describe the decision, output, evidence and operating constraints. The appropriate method may be a rule, a classifier, a retrieval system, a language model or a combination.

From task-specific models to reusable foundations

Earlier natural language processing (NLP) systems were usually trained for a bounded task and output space. A sentiment classifier produced one of a few labels. A named-entity recogniser marked spans. A translation model mapped one language sequence to another. Specialisation remains valuable: it can make a model smaller, faster to test and easier to constrain.

A foundation model is pretrained on broad data and then reused across tasks. Reuse may require only a prompt, or it may involve retrieval, supervised fine-tuning, preference optimisation or a task-specific output head. The foundation approach reduces the need to train a new language representation from the beginning for every application. It does not remove the need for task design, evaluation or controls.

The choice between specialised and general models is not a contest in which the largest model wins. It is a model-selection problem.

Retrieval is evidence supply, not a truth machine

Retrieval-augmented generation places selected source material in the model’s context before it generates an answer. A typical pipeline indexes documents, retrieves passages for a query and asks the model to answer from those passages. This can make private or current information available without putting it into the model’s weights.

Retrieval introduces its own failure modes. The relevant document may not have been indexed. Chunking may separate a definition from its exception. A search query may use different terminology from the source. The retriever may return a near match with the wrong customer, jurisdiction or document version. The generator may ignore a caveat even when it appears in the context. A citation can point to a passage that does not support the sentence beside it.

For that reason, a governed pipeline needs an evidence gate between retrieval and generation. The gate checks whether enough relevant, authorised and version-correct evidence is present. When the check fails, the proper output is an abstention or a request for more information, not a plausible completion.

Consider an independent, wholly fictional credit-document mini-case for a large UK bank. It combines publicly documented architectural patterns; it does not describe a live system, an internal design or a named organisation. The customer, amounts and documents are invented.

A credit analyst asks:

What is Northbridge Manufacturing Ltd’s interest-cover covenant, and when is it next tested?

The authorised corpus contains a facility agreement, an amendment, the latest credit paper and the bank’s lending policy. Retrieval returns the covenant clause from the original agreement and a later amendment. The evidence gate checks the customer identifier, facility identifier, document status, effective dates and access rights. Because the amendment supersedes the original threshold, the system must use the amended clause while retaining both documents in its provenance record.

The language model may then draft:

The amended minimum interest-cover ratio is 3.0×, tested quarterly. The next test date in the supplied schedule is 30 September 2026. Sources: Amendment 2, clause 4.1; Facility Agreement, schedule 6.

Those values are synthetic. In a real system, a deterministic validator should parse the date, ratio and units; a citation checker should confirm that each claim is supported by the cited span. If the amendment is missing, unsigned or outside the user’s permissions, the system should say that it lacks sufficient authorised evidence and route the query for review. The model’s prose is only one component of the control chain.

This example also separates three forms of information. The model’s parameters contain patterns acquired during pretraining. The prompt and retrieved passages provide temporary context for the current request. The audit record preserves the query, document versions, retrieved spans, validation outcomes and final decision. Conflating those stores creates avoidable security and governance errors.

Pretraining and the adaptation stack

Pretraining starts with a parameterised model and a large token corpus. For causal language modelling, the input at each position contains only earlier tokens, while the target is the next token. Optimisation repeatedly reduces the average prediction loss. The result is a base model: a model fitted to continue text, not automatically a safe or reliable assistant.

Adaptation is broader than a single second stage.

• Prompting and in-context examples change the current input, not the weights.
• Retrieval adds external evidence to the input and requires its own indexing, ranking and provenance controls.
• Supervised fine-tuning updates weights using curated input–output pairs. Instruction tuning is one form; classification fine-tuning is another.
• Parameter-efficient fine-tuning updates a small set of added or selected parameters rather than every base-model weight.
• Preference optimisation uses comparative feedback or a preference objective to shape which responses the model favours.
• Tools and deterministic orchestration connect language generation to search, calculation, databases and approved actions.

The distinctions matter because each intervention changes a different part of the system. Retrieval cannot repair a tokenizer. A prompt cannot update model weights. Fine-tuning cannot ensure that a policy document remains current. Preference optimisation cannot replace access control. Evaluation should identify which layer failed before prescribing a remedy.

It is also inaccurate to describe ChatGPT as simply “GPT-3 plus supervised fine-tuning”. Product implementations and model lineages change, and their full training details are not necessarily public. More generally, a chat assistant may combine a pretrained model with supervised instruction data, preference optimisation, safety policies, system instructions, retrieval, tools and serving-time controls. The model checkpoint and the delivered product are not the same object.

The economics of pretraining and adaptation vary with hardware, model architecture, sequence length, precision, dataset and utilisation. Fixed cost tables become misleading quickly. The durable principle is computational: pretraining distributes a broad representation across a large run; later adaptation reuses that representation for narrower requirements. Whether reuse is cheaper in a particular project must be measured rather than assumed.

The transformer’s original design

Before transformers, recurrent models processed sequence positions in order. Recurrence made the dependency path between distant tokens long and limited parallelism during training. Attention mechanisms allowed a model to form weighted combinations of representations at different positions. The 2017 Transformer made attention the central sequence-processing operation and dispensed with recurrence and convolution in its main architecture.4

The original Transformer was designed for sequence-to-sequence tasks such as machine translation. It contains an encoder stack and a decoder stack:

• Each encoder layer uses unmasked self-attention so that every source position can incorporate information from the complete source sequence, followed by a position-wise feed-forward network.
• Each decoder layer first uses masked self-attention over the target prefix. It then uses encoder–decoder attention, often called cross-attention, to read the encoder’s output. A feed-forward network completes the layer.
• Residual connections and normalisation support optimisation around these sublayers. Positional information is added because attention alone does not encode token order.

Self-attention creates a short path between any pair of positions in the available sequence. It also permits the representations for many positions to be computed in parallel during training. This does not mean that the mechanism solves reference, logic or meaning by itself. An attention weight is an internal mixing coefficient, not a causal explanation of a prediction. Capability depends on the learned projections, feed-forward layers, training data, objective and optimisation as well as the attention pattern.

Encoder-style and decoder-only models

Later model families retained different parts of the original design. Encoder-style models use bidirectional self-attention over an observed sequence and are natural representation builders. They suit tasks such as classification and token labelling, although the architecture does not confine them to those uses.

GPT is conventionally called decoder-only, but the shorthand can mislead. A GPT block is not the original Transformer decoder copied intact. It keeps causally masked self-attention and a feed-forward sublayer, but it omits the encoder and therefore omits encoder–decoder cross-attention. The prompt and generated continuation occupy one sequence. A causal mask prevents each position from reading later positions.

That distinction can be stated precisely:

• the full Transformer decoder attends to both the target prefix and a separate encoded source;
• a decoder-only GPT block attends causally within one sequence and has no separate source encoder;
• an encoder-only block attends across the complete supplied sequence and does not, by itself, define autoregressive generation.

The causal mask is what makes one training sequence yield many prediction examples. At position 5 the model predicts token 6 from tokens 1–5; at position 6 it predicts token 7 from tokens 1–6. During training these positions can be evaluated together because the mask blocks forbidden information. During generation, new tokens still arrive sequentially because the next prefix does not exist until a token has been selected.

The architecture repeats a small set of components: token and positional embeddings, attention, feed-forward networks, normalisation, residual paths and an output projection. Scale changes the width, number of heads, number of layers, training tokens and compute. The recurrence of the block is what makes a from-scratch implementation feasible; each chapter can isolate one mechanism before they are assembled.

Scale, context and compute

The adjective large can refer to several quantities that should not be confused.

Parameter count measures the trainable numerical values in the network. Parameters provide capacity, but a larger parameter count does not guarantee a better model if the data, optimisation or evaluation are poor.

Training-token count measures how much tokenised data is processed during training, including repeated sampling. GPT-3’s largest reported model had 175 billion parameters and was trained for 300 billion tokens.5 These two numbers describe different resources.

Context length is the maximum number of tokens available to a model for one forward pass. GPT-3 used 2,048 tokens.6 A context window is not persistent memory: material outside it is unavailable unless an application retrieves, summarises or stores that material and supplies it again.

Training compute depends on parameter count, token count, sequence length and implementation. Scaling one dimension while starving another can waste compute. Hoffmann and peers trained Chinchilla, a 70-billion-parameter model, on 1.4 trillion tokens and found that model size and training tokens should rise together under the compute budgets and model family they studied.7 Their result corrected a period of parameter-heavy scaling. It should not be reduced to a universal ratio that ignores data quality, architecture or the intended inference budget.

Scale can improve average loss and many downstream results. It also increases the importance of data governance, evaluation and operational constraints. The correct comparison for a product is not parameter count alone; it is whether a candidate system meets the required quality, latency, cost, privacy and control thresholds on representative data.

The build ahead

Our target is a compact GPT-style model whose internals remain visible. We are not reproducing the resources or performance of a frontier system. The value of the exercise lies in connecting each observed behaviour to a tensor operation, a parameter update or a serving decision.

The sequence is deliberate.

1. Represent text. We will tokenise raw text, construct input–target windows and map token identifiers and positions to vectors. This establishes the shapes consumed by the network.
2. Build attention. We will derive query, key and value projections, apply scaling and causal masks, then extend one attention head into multi-head attention.
3. Assemble GPT. We will add normalisation, feed-forward layers, residual connections and the vocabulary projection, checking dimensions and parameter counts as we go.
4. Pretrain and generate. We will calculate cross-entropy loss, update weights, monitor held-out loss and compare decoding policies. We will also load compatible pretrained weights to test the fidelity of our implementation.
5. Adapt the model. We will use labelled examples for classification and instruction–response examples for supervised instruction tuning, making the changed objective explicit in each case.
6. Govern the application. We will connect model behaviour to provenance, validation, abstention, access control, monitoring and human review through independent, fictional credit-document examples.

The next chapter starts at the boundary between language and mathematics. Text must become token identifiers, and token identifiers must become vectors. Once that bridge is explicit, “predict the next token” becomes a sequence of operations we can implement and inspect.

Build check

Before moving on, verify that you can do the following without relying on the chapter’s wording.

1. Given a next-token probability of 0.25 for the observed token, calculate its negative-log-likelihood: $-\ln(0.25) \approx 1.39$.
2. Explain why greedy decoding and sampling can produce different text from the same model without changing its weights.
3. Choose an architecture for each case: a fixed regulatory threshold, a high-volume three-label classifier, a cited answer from current policy, and an authorised update to a database. A defensible answer is, respectively: deterministic code; a compact classifier; retrieval plus an evidence-gated generator; and a typed tool call with validation and authorisation.
4. Distinguish the full Transformer decoder from a GPT block. The former includes cross-attention to an encoder; the latter has causal self-attention within one sequence and no encoder cross-attention.
5. State the three GPT-3 quantities kept separate in this chapter: 175 billion parameters, 300 billion training tokens and a 2,048-token context window.
6. Trace the credit-document query through retrieval, evidence checks, generation, deterministic validation and review. Identify the point at which the system must abstain if the operative amendment is unavailable.

Notes

What an embedding represents

An embedding is a vector whose coordinates are learned for a task. In a language model, each vocabulary item has one row in a token-embedding matrix. Training adjusts those rows so that the rest of the network can predict future tokens more accurately.

This definition is more useful than saying that an embedding is a numerical form of meaning. Meaning is neither assigned to individual coordinates nor guaranteed by geometric distance. A training objective can nevertheless produce useful neighbourhoods: tokens used in similar contexts often acquire vectors that are close under a measure such as cosine similarity. A two-dimensional projection may reveal groups, but it discards most of the geometry and depends on the projection method.

There is also a distinction between token embeddings and contextual representations. The embedding lookup is static: a particular token ID always selects the same row. After that row has passed through transformer blocks, its hidden state depends on the surrounding tokens. The token bank therefore starts from one vocabulary row whether it occurs beside river or mortgage; attention creates the context-sensitive representations later.

Before any lookup can happen, the tokenizer must decide what counts as a vocabulary item.

Choosing the units

A word-level vocabulary makes common text short, but it has to assign a symbol to every spelling, inflection, name and formatting variant it intends to support. A character vocabulary covers new words with a small symbol set, at the cost of longer sequences. Subword tokenisation occupies the space between those choices.

The example segmentation in the table is illustrative. A real tokenizer may split the word differently because its vocabulary was learned from a particular corpus. It may also include a preceding space in a token. GPT-2’s tokenizer, for example, can distinguish payment at the start of a string from payment after another word.

A tokenizer does three related jobs:

• pre-tokenisation divides text into regions on which merge rules operate;
• a vocabulary and merge table maps byte sequences to integer IDs; and
• special-token rules reserve IDs for boundaries or control markers.

The model sees only the resulting IDs. Token ID 12052 does not resemble its spelling, and adjacent IDs have no implied similarity. An ID is an index into two learned structures: the input embedding table and, later, the output vocabulary projection.

Byte-level BPE

Byte Pair Encoding was introduced as a compression method and adapted to subword modelling by Sennrich, Haddow and Birch. GPT-2 uses a byte-level form. Its path is more precise than the common description “start with characters and merge the most frequent pair”.

For GPT-2-style tokenisation:

1. a regular expression pre-tokenises the Unicode string;
2. each resulting piece is represented as UTF-8 bytes;
3. a reversible byte-to-Unicode mapping gives each of the 256 byte values a convenient internal symbol;
4. ranked BPE merges combine frequent adjacent byte sequences within those pieces.

Common byte sequences acquire compact tokens. Less frequent text remains split into smaller units. Because the base alphabet covers every byte value, valid Unicode text is representable even when a word or script was rare in the tokenizer’s training corpus.

The byte detail matters. One token can contain only part of a multi-byte UTF-8 character. Decoding that token in isolation may therefore produce invalid text. When inspecting boundaries, use the token’s raw bytes; decode the complete token sequence to recover the original valid string.

GPT-2’s vocabulary contains 50,257 entries, including its <|endoftext|> special token. The following code uses that fixed encoding for the synthetic clause.

import tiktoken

encoding = tiktoken.get_encoding("gpt2")
clause = (
    "If Northstar Retail Ltd misses an interest payment, the Agent "
    "must notify the Lenders within two Business Days."
)

token_ids = encoding.encode(clause)
token_bytes = [
    encoding.decode_single_token_bytes(token_id)
    for token_id in token_ids
]

print("vocabulary:", encoding.n_vocab)
print("token count:", len(token_ids))
print("token IDs:", token_ids)
print("token bytes:", token_bytes)
assert encoding.decode(token_ids) == clause

vocabulary: 50257
token count: 22
token IDs: [1532, 2258, 7364, 26702, 12052, 18297, 281, 1393, 6074,
11, 262, 15906, 1276, 19361, 262, 406, 7338, 1626, 734, 7320, 12579,
13]
token bytes: [b'If', b' North', b'star', b' Retail', b' Ltd',
b' misses', b' an', b' interest', b' payment', b',', b' the',
b' Agent', b' must', b' notify', b' the', b' L', b'enders',
b' within', b' two', b' Business', b' Days', b'.']

Northstar becomes North and star; capitalised Lenders becomes L and enders. Neither split is a corrupted word. The merge table simply has no higher-ranked token for that exact byte sequence in that context.

user_text = "The contract literally prints <|endoftext|> here."

# Treat the characters as ordinary user text.
user_ids = encoding.encode(user_text, disallowed_special=())

# Insert a trusted boundary as an ID constructed by the application.
training_ids = user_ids + [encoding.eot_token]

From a token stream to prediction examples

A causal language model learns to predict the next token at every visible position. For a contiguous sequence

$$[t_0,t_1,t_2,t_3,t_4],$$

one training example of context length four is:

$$x=[t_0,t_1,t_2,t_3], \qquad y=[t_1,t_2,t_3,t_4].$$

The input and target have equal length. y[j] is the token immediately after x[j] in the original stream. The causal attention mask introduced in the next chapter ensures that the representation at position $j$ cannot inspect x[j + 1] while making that prediction.

Long streams yield many examples. A window starts at index $s$, takes context_length input tokens, then takes the same number of target tokens from $s+1$. The stride is the distance from one start index to the next.

import torch
from torch.utils.data import Dataset


class NextTokenDataset(Dataset):
    def __init__(self, token_ids, context_length, stride):
        if context_length < 1:
            raise ValueError("context_length must be positive")
        if not 1 <= stride <= context_length:
            raise ValueError(
                "stride must be between 1 and context_length"
            )

        self.context_length = context_length
        self.ids = torch.as_tensor(token_ids, dtype=torch.long)
        if self.ids.numel() <= context_length:
            raise ValueError(
                "need at least context_length + 1 tokens"
            )

        stop = int(self.ids.numel()) - context_length
        self.starts = range(0, stop, stride)

    def __len__(self):
        return len(self.starts)

    def __getitem__(self, index):
        start = self.starts[index]
        end = start + self.context_length
        inputs = self.ids[start:end]
        targets = self.ids[start + 1:end + 1]
        return inputs, targets


dataset = NextTokenDataset(
    token_ids,
    context_length=8,
    stride=8,
)
inputs, targets = dataset[0]

print(inputs.shape, len(dataset))
print(inputs.tolist())
print(targets.tolist())
assert torch.equal(inputs[1:], targets[:-1])

torch.Size([8]) 2
[1532, 2258, 7364, 26702, 12052, 18297, 281, 1393]
[2258, 7364, 26702, 12052, 18297, 281, 1393, 6074]

The class stores one token tensor and indexes views into it. It also constrains the stride to at most the context length so that the starts do not leave untrained gaps in the target stream. With 22 tokens, context length 8 and stride 8, it produces two complete examples. The short remainder is omitted. That choice is explicit; it should be measured rather than hidden.

from torch.utils.data import DataLoader


loader = DataLoader(
    dataset,
    batch_size=2,
    shuffle=False,
    drop_last=False,
)

batch_inputs, batch_targets = next(iter(loader))
print("inputs:", batch_inputs.shape)
print("targets:", batch_targets.shape)
assert torch.equal(
    batch_inputs[:, 1:],
    batch_targets[:, :-1],
)

inputs: torch.Size([2, 8])
targets: torch.Size([2, 8])

Embedding lookup

Let the vocabulary size be $V$ and model width be $D$. The trainable token embedding matrix is

$$E \in \mathbb{R}^{V \times D}.$$

For token ID $i$, the lookup returns row $E_i$. An input tensor of IDs with shape $[B,T]$ therefore becomes a floating-point tensor with shape $[B,T,D]$.

The lookup is equivalent to multiplying a one-hot row vector by $E$. If $V=5$ and the ID is 2, the one-hot vector

$$[0,0,1,0,0]$$

selects the third row. nn.Embedding avoids constructing the mostly zero one-hot vector.

Here is a hand-worked example with three vocabulary items selected from a toy five-token vocabulary. The vectors are assigned for inspection; in the model they begin from an initialisation scheme and are learned.

$$\begin{aligned} E_2 &= [0.2,-0.4,0.6] && \text{(` covenant')}\\ E_4 &= [-0.1,0.7,0.2] && \text{(` breach')}\\ E_1 &= [0.5,0.0,-0.3] && \text{(`.`)} \end{aligned}$$

For IDs $[2,4,1]$, lookup produces a $3\times3$ matrix:

$$\begin{bmatrix} 0.2 & -0.4 & 0.6\\ -0.1 & 0.7 & 0.2\\ 0.5 & 0.0 & -0.3 \end{bmatrix}.$$

The rows preserve sequence order. Similarity between ID values never enters the calculation: ID 4 is no closer to ID 2 than it is to ID 1.

For the GPT-2 Small configuration, $V=50{,}257$ and $D=768$, so the token embedding table has

$$50{,}257 \times 768 = 38{,}597{,}376$$

trainable parameters. An ID batch with $B=8$ and $T=1{,}024$ has shape $[8,1024]$; after lookup its shape is $[8,1024,768]$. Repeated IDs select the same row, and their gradient contributions accumulate when the loss is backpropagated.

Adding order

Content-only self-attention without position information is permutation-equivariant: permuting input rows produces the corresponding permutation of output rows. It cannot infer the original order from token vectors alone. A causal mask introduces a direction because each position sees a different prefix. GPT-2 nevertheless supplies an explicit signal for position and distance.

GPT-2 uses a learned absolute-position matrix

$$P \in \mathbb{R}^{C \times D},$$

where $C$ is the configured context length. For a sequence of length $T$, positions $0$ through $T-1$ select $T$ rows. The model input is

$$X_{b,t}=E_{\text{ID}_{b,t}}+P_t.$$

Token and position vectors have the same width $D$, so element-wise addition retains shape $[B,T,D]$.

Continue the hand-worked example with:

$$P_0=[0.1,0.0,0.0],\quad P_1=[0.0,0.1,0.0],\quad P_2=[0.0,0.0,0.1].$$

The first combined row is

$$[0.2,-0.4,0.6]+[0.1,0.0,0.0] =[0.3,-0.4,0.6].$$

All three rows can be reproduced with trainable PyTorch embeddings.

import torch
from torch import nn


token_weights = torch.tensor([
    [ 0.1,  0.0, -0.1],
    [ 0.5,  0.0, -0.3],
    [ 0.2, -0.4,  0.6],
    [-0.2,  0.1,  0.4],
    [-0.1,  0.7,  0.2],
])
position_weights = torch.tensor([
    [0.1, 0.0, 0.0],
    [0.0, 0.1, 0.0],
    [0.0, 0.0, 0.1],
])

token_embedding = nn.Embedding.from_pretrained(
    token_weights,
    freeze=False,
)
position_embedding = nn.Embedding.from_pretrained(
    position_weights,
    freeze=False,
)

input_ids = torch.tensor([[2, 4, 1]])  # [B=1, T=3]
positions = torch.arange(input_ids.shape[1])

token_vectors = token_embedding(input_ids)       # [1, 3, 3]
position_vectors = position_embedding(positions) # [3, 3]
model_input = token_vectors + position_vectors   # [1, 3, 3]

print(model_input)
assert model_input.shape == (1, 3, 3)

tensor([[[ 0.3000, -0.4000,  0.6000],
         [-0.1000,  0.8000,  0.2000],
         [ 0.5000,  0.0000, -0.2000]]], grad_fn=<AddBackward0>)

Both embedding modules have requires_grad=True. Later layers supply the loss; backpropagation adjusts the token and position rows that participated in the batch.

The complete input contract

The pipeline can now be stated as a set of testable tensor contracts.

Tokenisation is a discrete preprocessing operation in this build. The vocabulary and merges are trained or selected before model training and then fixed. The embedding tables are ordinary trainable parameters. This boundary is useful when debugging: if decoded IDs do not reproduce the source, inspect the text and tokenizer; if IDs are correct but vectors or gradients are wrong, inspect the model path.

The synthetic clause also shows why provenance should survive preprocessing. A model consumes IDs, but an audit or human reviewer needs to recover the source document, page, clause span and tokenizer version. Keep those references beside the tokenised example rather than attempting to reconstruct them from a generated answer.

Build check

Run these checks before connecting the input pipeline to attention.

1. Round trip representative text. Verify decode(encode(text)) == text for plain English, curly punctuation, pound values, combining marks, intended scripts and extracted document samples. Inspect full token sequences rather than assuming each token decodes alone.

2. Freeze and fingerprint the tokenizer. Record the encoding name, vocabulary and merge digest, normalisation policy, software version and special-token map beside the checkpoint. Fail model loading when the fingerprint differs.

3. Test fixed token IDs. Keep a small regression set that includes the synthetic clause. An unexplained change in IDs is a release failure even when round-trip decoding still succeeds.

4. Measure the target distribution. Report token counts, truncation and tail loss by language and document type. Universal byte coverage is not a substitute for this measurement.

5. Keep trusted markers out of user control. Tokenise literal user text as ordinary text and let the application insert document or role markers. Treat special-token handling as one protocol check among several security controls.

6. Split before windowing. Prevent overlapping clauses, documents or customer-related material from crossing evaluation partitions. Confirm that a window never straddles an unintended boundary.

7. Reconcile token accounting. For each run, count source, predicted, padded, masked and dropped tokens. Record stride, context length, drop_last and any packing policy.

8. Assert shifts and shapes. For every unpadded contiguous example, check inputs[1:] == targets[:-1]. After batching, IDs should have shape $[B,T]$, embeddings $[B,T,D]$, and all IDs must lie in $[0,V)$.

9. Test masks separately. A causal mask, padding mask and loss mask solve different problems. Construct a padded example and verify that padded labels contribute zero loss and padded keys receive no attention.

10. Check position bounds. For learned absolute embeddings, require $T \le C$ before lookup. Test the first and last valid positions and make truncation an explicit caller decision.

11. Inspect numerical health. Confirm embedding outputs are finite, on the intended device and in the intended dtype. Backpropagate a tiny loss and verify that used rows receive gradients.

12. Decode a batch during training. Periodically decode input and target rows with their provenance. A readable sample often exposes duplicated boundaries, extraction artefacts or an off-by-one shift before they consume a training run.

At this point the model has an ordered matrix of vectors and a target token for each position. Chapter 3 builds the causal attention operation that lets each vector gather information from the permitted part of that matrix.

Notes and primary sources

1. Rico Sennrich, Barry Haddow and Alexandra Birch, “Neural Machine Translation of Rare Words with Subword Units”, Proceedings of ACL, 2016. This paper adapts BPE to an open-vocabulary subword translation setting.

2. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei and Ilya Sutskever, “Language Models are Unsupervised Multitask Learners”, OpenAI, 2019. Sections 2.2–2.3 describe GPT-2’s byte-level BPE design, 50,257-entry vocabulary and 1,024-token context configuration.

3. OpenAI, tiktoken 0.13.0 repository and README, with the encoding API in core.py, accessed 28 July 2026. These are the implementation sources for encode, decode, decode_single_token_bytes, and special-token handling used in the tokenizer example above.

4. Ashish Vaswani and peers, “Attention Is All You Need”, NeurIPS, 2017. Section 3.5 explains why a transformer without recurrence or convolution needs position information and presents sinusoidal encoding.

5. Jianlin Su and peers, “RoFormer: Enhanced Transformer with Rotary Position Embedding”, 2021. The paper defines RoPE by applying position-dependent rotations to queries and keys.

6. Shouyuan Chen and peers, “Extending Context Window of Large Language Models via Positional Interpolation”,

   2023. The need for interpolation and further fine-tuning illustrates why a RoPE implementation alone does not establish reliable extrapolation.

From one-vector compression to direct access

Early neural machine-translation systems used a recurrent encoder and decoder. The encoder processed the source sentence from left to right and passed a fixed-size summary to the decoder. That design placed a demanding compression step between the input and output. Details from a long source sentence had to survive inside one vector before the decoder could use them.

Bahdanau and peers changed the interface. Their decoder could inspect all encoder states at each output step and construct a new weighted summary. When generating a translated date, it could give more weight to the source date; when generating a subject, it could favour the relevant noun phrase. The weights were learned with the rest of the model.

The transformer removed recurrence from the main sequence path. Its self-attention mechanism applies the same broad idea within one sequence: positions exchange information directly through learned weighted sums.

We will build the decoder form used by a GPT-style model. It differs from the original translation setting in one important respect: a token may use only itself and earlier tokens. The causal mask enforces that boundary.

The tensor contract

Let:

• $B$ be batch size;
• $T$ be sequence length;
• $D$ be model width;
• $H$ be the number of attention heads; and
• $D_h = D/H$ be the width of one head.

The hidden-state tensor entering an attention layer has shape $[B,T,D]$. A multi-head layer returns the same shape so it can sit inside a residual block.

For one head, the scaled dot-product calculation is:

$$\operatorname{Attention}(Q,K,V) = \operatorname{softmax}\left( \frac{QK^\mathsf{T}}{\sqrt{D_h}} + M \right)V$$

where $M$ is a mask. The shapes for one batch item are:

With all heads present, the score tensor has shape $[B,H,T,T]$. This four-dimensional object is the source of both attention’s flexibility and its cost.

A non-trainable warm-up

Before introducing learned projections, consider six input vectors arranged as a matrix $X$ with shape $[6,3]$. A basic similarity-weighted mix is:

scores = X @ X.T
weights = torch.softmax(scores, dim=-1)
context = weights @ X

The entry scores[i, j] is the dot product between positions i and j. Softmax is applied across the last dimension, so every row becomes a distribution over source positions. context[i] is then the weighted sum of all rows of X for query position i.

Suppose one row of scores is:

[0.8, 1.9, 1.1, 0.2]

Subtracting the maximum before exponentiation gives the same softmax with better numerical stability:

$$\begin{aligned} s' &= [-1.1,\;0.0,\;-0.8,\;-1.7] \\ \exp(s') &\approx [0.333,\;1.000,\;0.449,\;0.183] \\ \operatorname{softmax}(s) &\approx [0.169,\;0.509,\;0.229,\;0.093]. \end{aligned}$$

The second source vector supplies just over half of this context vector. The calculation does not copy that vector; it blends all four according to the weights.

This warm-up has three limitations.

First, a single vector plays query, key and value at once. The model cannot learn separate match and payload representations. Second, the score matrix comes directly from the current hidden space; the layer has no attention parameters to train. Third, no information boundary prevents a position from using future tokens.

X @ X.T is symmetric, but the row-wise softmax result need not be. Each row has its own normalising denominator. Once separate query and key projections are used, even the raw score matrix is generally asymmetric.

Queries, keys and values

A trainable attention layer projects each hidden state three ways:

$$Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V.$$

W_Q, W_K and W_V are learned parameter matrices. They share an input but not a job.

• A query represents the features a position uses to seek relevant context.
• A key represents the features against which other positions match.
• A value carries the information that a successful match can contribute.

The familiar database analogy is helpful up to a point: a lookup key determines which record to select, while the returned record may contain different fields. Attention is softer than a database lookup because it normally mixes many values rather than selecting exactly one.

Why divide the scores by $\sqrt{D_h}$? Under a simple assumption that query and key coordinates are independent with mean zero and variance one, their dot product has variance proportional to $D_h$. Large unscaled scores make softmax very sharp. Most derivatives then become small, which can hamper learning. Dividing by $\sqrt{D_h}$ keeps the score scale more stable as the head width changes.

For GPT-2 Small, $D=768$, $H=12$, and $D_h=64$. The divisor is therefore $\sqrt{64}=8$. The head dimension is 64, not 768; 768 is the combined model width across all twelve heads.

The causal boundary

During training, all positions in a sequence are processed in parallel. Without a mask, position 3 could inspect the correct token at position 4 while learning to predict it. Loss would look excellent, but generation would fail because position 4 does not exist yet at inference time.

A causal mask blocks score $(i,j)$ whenever $j>i$. In matrix form it is upper triangular:

query 0:  ✓  ×  ×  ×
query 1:  ✓  ✓  ×  ×
query 2:  ✓  ✓  ✓  ×
query 3:  ✓  ✓  ✓  ✓

Blocked logits are set to negative infinity before softmax. Their exponential is zero, so they receive no weight.

Removing this mask does not turn GPT into BERT. It creates a bidirectional attention pattern, but BERT also uses an encoder stack, a different training corruption objective and its own input conventions. Architecture, objective and data interface all matter.

Causal masking is not the only mask a deployed layer may need. Batches often contain padding. A key-padding mask prevents real query positions from using padding keys. The two masks express different rules:

• the causal mask blocks information from the future; and
• the padding mask blocks positions that are not data.

When variable-length sequences are padded on the right, their combined mask can be broadcast to $[B,1,T,T]$. Query padding also needs careful downstream handling; otherwise unused padded rows can still produce values even though real tokens do not attend to padded keys.

A complete multi-head implementation

One head learns one set of query, key and value projections. Multiple heads perform several attention calculations in parallel, concatenate their outputs, and apply a final output projection. The heads are not assigned jobs such as “syntax” or “dates” by the programmer. Training may produce specialised, overlapping or redundant behaviours.

The implementation below supports causal masking and an optional key-padding mask. It keeps line lengths modest for narrow EPUB screens.

import torch
from torch import nn


class MultiHeadCausalAttention(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        context_length: int,
        dropout: float = 0.0,
        qkv_bias: bool = False,
    ) -> None:
        super().__init__()
        if d_model % num_heads != 0:
            raise ValueError("d_model must be divisible by num_heads")

        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.d_model = d_model

        self.q_proj = nn.Linear(
            d_model, d_model, bias=qkv_bias
        )
        self.k_proj = nn.Linear(
            d_model, d_model, bias=qkv_bias
        )
        self.v_proj = nn.Linear(
            d_model, d_model, bias=qkv_bias
        )
        self.out_proj = nn.Linear(d_model, d_model)
        self.attn_dropout = nn.Dropout(dropout)

        causal = torch.triu(
            torch.ones(
                context_length,
                context_length,
                dtype=torch.bool,
            ),
            diagonal=1,
        )
        self.register_buffer(
            "causal_mask",
            causal,
            persistent=False,
        )

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        batch, tokens, _ = x.shape
        x = x.view(
            batch,
            tokens,
            self.num_heads,
            self.head_dim,
        )
        return x.transpose(1, 2)

    def forward(
        self,
        x: torch.Tensor,
        key_is_padding: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch, tokens, width = x.shape
        if width != self.d_model:
            raise ValueError("unexpected model width")
        if tokens > self.causal_mask.shape[0]:
            raise ValueError("sequence exceeds context length")

        q = self._split_heads(self.q_proj(x))
        k = self._split_heads(self.k_proj(x))
        v = self._split_heads(self.v_proj(x))

        scores = q @ k.transpose(-2, -1)
        scores = scores / (self.head_dim**0.5)

        causal = self.causal_mask[:tokens, :tokens]
        scores = scores.masked_fill(
            causal[None, None, :, :],
            -torch.inf,
        )

        if key_is_padding is not None:
            expected = (batch, tokens)
            if key_is_padding.shape != expected:
                raise ValueError(
                    f"padding mask must have shape {expected}"
                )
            scores = scores.masked_fill(
                key_is_padding[:, None, None, :],
                -torch.inf,
            )

        weights = torch.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)
        context = weights @ v

        context = context.transpose(1, 2).contiguous()
        context = context.view(batch, tokens, self.d_model)
        return self.out_proj(context)

The three projections initially produce $[B,T,D]$. _split_heads reshapes that to $[B,H,T,D_h]$. The score multiplication therefore yields $[B,H,T,T]$. After the value multiplication, transposition and reshape restore $[B,T,D]$.

contiguous() matters because transpose changes the tensor’s strides. A subsequent view requires a compatible contiguous memory layout.

The causal mask is registered as a buffer so it follows the module across devices without being optimised. Marking it non-persistent avoids storing a reconstructible triangular matrix in every checkpoint. Either choice is valid if the loading contract is explicit.

When dropout is active, it is applied to attention probabilities during training. The retained entries are rescaled in the usual dropout manner; a realised row therefore need not sum exactly to one. Evaluation mode disables dropout.

Tests that catch real mistakes

Successful execution is not enough. Attention bugs can produce plausible shapes while violating causality.

def test_attention_contract() -> None:
    torch.manual_seed(7)
    layer = MultiHeadCausalAttention(
        d_model=12,
        num_heads=3,
        context_length=8,
        dropout=0.0,
    )
    layer.eval()

    x = torch.randn(2, 5, 12)
    y = layer(x)
    assert y.shape == x.shape
    assert torch.isfinite(y).all()


def test_future_token_cannot_change_the_past() -> None:
    torch.manual_seed(8)
    layer = MultiHeadCausalAttention(
        d_model=12,
        num_heads=3,
        context_length=8,
        dropout=0.0,
    )
    layer.eval()

    original = torch.randn(1, 5, 12)
    changed = original.clone()
    changed[:, 4, :] += 100.0

    y_original = layer(original)
    y_changed = layer(changed)

    torch.testing.assert_close(
        y_original[:, :4, :],
        y_changed[:, :4, :],
        atol=1e-6,
        rtol=1e-6,
    )

The second test deliberately changes the final input. Outputs at positions zero through three must remain unchanged. The final output may change because it is allowed to use the final input.

A fuller test suite should also check:

• a padding key cannot affect real-token outputs;
• gradients reach all four projection layers;
• invalid head counts and excessive sequence lengths fail clearly;
• dropout is reproducible under a fixed seed in training mode;
• serial single-head and vectorised multi-head implementations agree when parameters are copied carefully; and
• output values match a trusted framework implementation for a fixed fixture.

Position remains a separate problem

Self-attention does not encode token order from content vectors alone. Without position information and without a directional mask, permuting the input permutes the outputs in the same way. The operation is permutation-equivariant.

The causal mask adds a directional visibility structure: earlier and later positions no longer have identical neighbourhoods. A decoder still needs a position representation to distinguish locations effectively. GPT-2 adds learned absolute position embeddings before the first transformer block. Many later decoders rotate queries and keys using rotary position embeddings. Neither choice is contained in the attention formula itself.

Long-context behaviour also requires more than substituting a rotary encoding. The training distribution, positional scaling method, numerical precision, attention implementation and evaluation length all affect extrapolation. A model trained only on short sequences should not be presumed reliable at a much longer length because its code accepts a larger tensor.

The quadratic cost

The score matrix contains $T^2$ entries per head. Doubling sequence length roughly quadruples the score calculation and the size of a materialised attention matrix:

Those counts describe the logical score pairs. An implementation need not store the whole matrix at once.

FlashAttention reorganises exact attention into tiles so intermediate scores need not make repeated round trips to high-bandwidth memory. It reduces memory traffic and the amount of stored intermediate state; it does not change dense attention’s quadratic pairwise work into a linear algorithm.

Autoregressive inference has a different optimisation. Once a prefix has been processed, the keys and values for its tokens can be cached. The next decoding step computes a query for the new token and attends to the cached keys and values. This avoids recomputing the entire prefix, but the cache grows linearly with sequence length, layers and key–value heads. Each new token still compares against the visible prefix.

Multi-query attention shares one set of key–value heads across many query heads. Grouped-query attention uses an intermediate number of key–value groups. Both reduce cache size and memory bandwidth. They do not remove the need for multiple query heads, nor do they make the context free.

When evaluating a long-context design, separate at least four quantities:

1. training-time arithmetic;
2. training-time activation memory;
3. prefill latency for an existing prompt; and
4. per-token decode latency and key–value cache memory.

One phrase such as “supports 128K” hides all four.

Parameter and cache accounting

The projection matrices also deserve a concrete count. Ignoring bias terms, a standard multi-head layer with equal input and output width has:

$$3D^2 \quad\text{parameters for } W_Q,W_K,W_V$$

and:

$$D^2 \quad\text{parameters for the output projection}.$$

At $D=768$, that is $4\times768^2=2{,}359{,}296$ attention parameters per block. Splitting the width into twelve heads does not multiply that count: the three large projection matrices already produce all heads at once.

The key–value cache has a different shape. For a conventional layer, the cached keys and values for one sequence each contain $[H,T,D_h]$ elements. Together they require $2TD$ elements per layer. For GPT-2 Small at $T=1{,}024$, that is:

$$2 \times 1{,}024 \times 768 = 1{,}572{,}864$$

elements per layer, before batch size and numeric precision are considered. Across twelve layers in 16-bit storage, the idealised payload is about 36 MiB for one sequence. Real serving memory includes allocation overhead, temporary buffers and framework-specific layout.

This calculation explains why grouped-query attention can matter even when it does not alter the number of query heads. If twelve query heads share, for example, three key–value groups, the cached key–value width falls by a factor of four. Quality and throughput still need measurement on the intended model and hardware; the arithmetic only identifies the resource being reduced.

Failure clinic

Most first implementations fail in ordinary ways. The resulting tensors may still look plausible, which is why each mistake should have a targeted test.

weights = torch.softmax(scores, dim=-1)

What attention weights can and cannot tell us

An attention map is a record of one internal mixing operation. It can help an engineer inspect masks, identify dead or highly concentrated heads, and compare behaviour across fixtures. It is not, by itself, a faithful causal explanation of a model’s output.

Several factors break the simple story that “the highest weight is the most important word”:

• values and the output projection transform what a weight carries;
• information may have been mixed by earlier layers;
• residual connections provide paths around the attention sublayer;
• different heads may compensate for one another; and
• alternative attention patterns can sometimes produce similar outputs.

The literature includes both demonstrations of these limitations and arguments for narrower, carefully specified uses of attention as explanation. The practical conclusion is modest: state what was measured. “Head 4 assigned weight 0.31 to token 17” is an observation. “Token 17 caused the credit decision” is a much stronger claim requiring intervention-based or otherwise validated attribution evidence.

For consequential systems, explanations should be anchored in application evidence: cited source passages, deterministic calculations, feature provenance, rule traces, model version, input snapshot and human review. An attention heatmap may accompany engineering diagnostics; it should not stand in for that audit trail.

Worked case: a covenant clause

Consider this synthetic clause from a fictional facility agreement:

The borrower shall ensure that consolidated net debt to EBITDA does not exceed 3.50× on each quarter-end test date.

A decoder representation near 3.50× can use earlier tokens for the metric, operator and subject. Multiple heads may learn useful mixing patterns. The application still should not ask an unconstrained language-model output to become the system of record.

A safer boundary is:

source page and coordinates
        ↓
candidate span extraction
        ↓
typed parse:
  metric = "net_debt_to_ebitda"
  operator = "<="
  threshold = Decimal("3.50")
  test_frequency = "quarter_end"
        ↓
deterministic schema and arithmetic checks
        ↓
human confirmation for ambiguous or high-impact cases

The model can propose spans and normalised fields. Deterministic code should check that the operator agrees with the source language, parse the number with decimal arithmetic, preserve units, retain page provenance and reject values outside the permitted schema. If two clauses conflict, the system should surface both rather than average them into a fluent answer.

This division of labour follows directly from the mechanism. Attention is good at contextual representation. It does not guarantee exact extraction, calibrated confidence or policy authority.

Build check

Before placing the layer inside a transformer block, verify that you can answer and test each of these questions.

1. Which dimension does softmax normalise, and why?
2. Why is the scale factor based on head width rather than total model width?
3. What shapes do q, k, v, scores and the output have?
4. What distinct failures do causal and padding masks prevent?
5. Can changing a future input alter an earlier output in evaluation mode?
6. Why is out_proj present after concatenating the heads?
7. Which costs remain quadratic, and which memory traffic does a tiled implementation avoid?
8. Why is a key–value cache useful, and what resource does it consume?
9. What exactly can an attention map support as evidence?
10. Which deterministic controls remain necessary in the covenant example?

The component we carry into Chapter 4 has a compact contract: $[B,T,D]$ enters, $[B,T,D]$ leaves, no output position uses a future input, and all contextual mixing remains differentiable. The next task is to place it inside a stable residual block.

Notes and further reading

1. Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio, “Neural Machine Translation by Jointly Learning to Align and Translate” (2014), arXiv:1409.0473.
2. Ashish Vaswani and peers, “Attention Is All You Need” (2017), arXiv:1706.03762.
3. Tri Dao and peers, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (2022), arXiv:2205.14135.
4. Joshua Ainslie and peers, “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints” (2023), arXiv:2305.13245.
5. Sarthak Jain and Byron C. Wallace, “Attention is not Explanation” (2019), ACL Anthology N19-1357.
6. Sarah Wiegreffe and Yuval Pinter, “Attention is not not Explanation” (2019), ACL Anthology D19-1002.

Layer normalisation stabilises each token vector

For one token vector $x \in \mathbb{R}^{D}$, LayerNorm computes the mean and population variance across its $D$ features:

$$\mu = \frac{1}{D}\sum_{i=1}^{D}x_i, \qquad \sigma^2 = \frac{1}{D}\sum_{i=1}^{D}(x_i-\mu)^2.$$

It then normalises the vector and applies a learned scale $\gamma$ and shift $\beta$:

$$\operatorname{LayerNorm}(x) = \gamma \odot \frac{x-\mu}{\sqrt{\sigma^2+\epsilon}} +\beta.$$

The small constant $\epsilon$ prevents division by zero. At initialisation, $\gamma$ is one and $\beta$ is zero, so a non-constant vector leaves the operation with approximately zero mean and unit variance. During training, the two learned vectors allow each feature to acquire a useful scale and offset.

Statistics are computed independently for every token in every example. LayerNorm therefore does not depend on other items in the batch, and its calculation is the same in training and evaluation. This differs from BatchNorm, which estimates statistics across examples.

The placement of normalisation changes the residual block. The original Transformer used post-normalisation: a sublayer ran first, its residual was added, and the sum was normalised. GPT-2 moved normalisation before each sublayer and added a final LayerNorm after the block stack. This pre-normalisation arrangement often gives more manageable gradients at initialisation and can simplify optimisation. “Often” matters. Pre-LN is not strictly superior for every depth, objective, initialisation or training recipe; post-LN models can train well with suitable warm-up, scaling and other stabilisation.

LayerNorm operates on hidden features. It should not be described as normalising the human-scale values mentioned in text. The token sequence “£24 million” is represented by learned vectors; LayerNorm does not inspect the currency amount and rescale it as an accounting number.

GELU supplies a smooth nonlinearity

Two linear maps in succession collapse to one linear map. A feed-forward network needs a nonlinearity between them if it is to represent nonlinear functions. The Gaussian Error Linear Unit is

$$\operatorname{GELU}(x) = x\Phi(x),$$

where $\Phi(x)$ is the cumulative distribution function of a standard normal variable. The code in this chapter evaluates that definition through the error function. A frequently seen expression using tanh and the coefficient 0.044715 is an approximation, not the definition.

GELU differs from ReLU most visibly around zero. ReLU clips every negative value to zero; GELU has a small negative region and approaches zero smoothly as the input moves left. Its derivative also changes smoothly. These properties can affect optimisation, but they do not prove that GELU wins in every architecture or training regime.

The feed-forward network expands, transforms and contracts

Attention mixes information across sequence positions. The feed-forward network applies the same nonlinear transformation to each position independently:

$$\operatorname{FFN}(x) = W_2\,\operatorname{GELU}(W_1x+b_1)+b_2.$$

In the reference model, $W_1$ expands $D$ features to $4D$, and $W_2$ contracts them to $D$. For $D=768$, the intermediate width is 3,072. This four-times ratio is a design choice inherited by many early Transformers; it is not a mathematical requirement. Gated networks often choose a different intermediate width to keep their parameter budget comparable.

The expansion gives the activation function more channels on which to operate. The contraction is equally important: it restores $[B,T,D]$, allowing the result to be added to the residual stream and passed to the next block.

With biases, one $D \rightarrow 4D \rightarrow D$ feed-forward network contains

$$D(4D)+4D+(4D)D+D = 8D^2+5D$$

parameters. At $D=768$, that is 4,722,432 parameters. This exceeds the parameter count of the attention projections in the same block, although parameter count alone does not measure runtime: attention also forms matrices whose size grows with sequence length.

Residual paths provide an additive route

A residual sublayer has the form

$$y = x + f(x).$$

Its Jacobian with respect to $x$ is

$$\frac{\partial y}{\partial x} = I + J_f(x).$$

The identity term supplies a direct additive path for activations and gradients. It does not guarantee that every gradient component has magnitude at least one. The branch Jacobian can rotate, amplify or partly cancel the identity contribution; an eigenvalue of $J_f$ near $-1$ can make the corresponding eigenvalue of $I+J_f$ near zero. Across many blocks, the total Jacobian is still a product of block Jacobians.

The practical benefit is more modest and more defensible: the optimiser does not have to transmit all information through every learned transformation. A block can begin near an identity mapping, and shorter gradient routes are available. Initialisation, normalisation, depth and the learned branch all remain relevant.

In a pre-normalisation decoder block, the two updates are

$$x' = x + \operatorname{Dropout} \left(\operatorname{Attention}(\operatorname{LN}_1(x))\right)$$

and

$$y = x' + \operatorname{Dropout} \left(\operatorname{FFN}(\operatorname{LN}_2(x'))\right).$$

Both learned branches return $[B,T,D]$, so both additions are well-defined. Dropout is active during training and disabled in evaluation.

One transformer block

The block now has a precise division of labour. LayerNorm controls the scale presented to each learned sublayer. Causal attention mixes information from permitted earlier positions. The feed-forward network transforms each position separately. Residual additions carry the stream through both operations.

The word “identical” can be misleading when blocks are stacked. They share an architecture, but each block normally owns different learned parameters. Reusing one block object repeatedly would tie all block weights and create a different model.

The causal mask is also part of the block contract. At position $t$, attention may read positions $0$ through $t$, never positions $t+1$ onwards. Padded batches need an additional key-padding mask. The compact implementation accepts that optional mask; ordinary packed language-model batches can omit it because every supplied position is real.

A complete runnable reference model

This is the sole model implementation used in the chapter. It includes exact GELU, custom LayerNorm, causal multi-head attention, the feed-forward network, pre-normalisation blocks, learned absolute positions, tied token embeddings and a greedy generation interface.

import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass(frozen=True)
class GPTConfig:
    vocab_size: int = 50_257
    context_length: int = 1_024
    d_model: int = 768
    n_heads: int = 12
    n_layers: int = 12
    dropout: float = 0.1
    qkv_bias: bool = True
    ff_multiplier: int = 4
    tie_embeddings: bool = True

    def __post_init__(self) -> None:
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if min(
            self.vocab_size,
            self.context_length,
            self.d_model,
            self.n_heads,
            self.n_layers,
            self.ff_multiplier,
        ) <= 0:
            raise ValueError("configuration values must be positive")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")


class LayerNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))
        self.bias = nn.Parameter(torch.zeros(d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=-1, keepdim=True)
        variance = x.var(dim=-1, keepdim=True, unbiased=False)
        normalised = (x - mean) * torch.rsqrt(variance + self.eps)
        return self.weight * normalised + self.bias


class GELU(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0)))


class FeedForward(nn.Module):
    def __init__(self, cfg: GPTConfig) -> None:
        super().__init__()
        hidden_size = cfg.ff_multiplier * cfg.d_model
        self.in_proj = nn.Linear(cfg.d_model, hidden_size)
        self.activation = GELU()
        self.out_proj = nn.Linear(hidden_size, cfg.d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.out_proj(self.activation(self.in_proj(x)))


class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: GPTConfig) -> None:
        super().__init__()
        self.d_model = cfg.d_model
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.d_model // cfg.n_heads

        self.qkv = nn.Linear(
            cfg.d_model,
            3 * cfg.d_model,
            bias=cfg.qkv_bias,
        )
        self.out_proj = nn.Linear(cfg.d_model, cfg.d_model)
        self.attention_dropout = nn.Dropout(cfg.dropout)

        mask = torch.triu(
            torch.ones(
                cfg.context_length,
                cfg.context_length,
                dtype=torch.bool,
            ),
            diagonal=1,
        )
        self.register_buffer(
            "causal_mask",
            mask.view(1, 1, cfg.context_length, cfg.context_length),
            persistent=False,
        )

    def forward(
        self,
        x: torch.Tensor,
        key_is_padding: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch_size, n_tokens, _ = x.shape
        query, key, value = self.qkv(x).chunk(3, dim=-1)

        def split_heads(tensor: torch.Tensor) -> torch.Tensor:
            return tensor.view(
                batch_size,
                n_tokens,
                self.n_heads,
                self.head_dim,
            ).transpose(1, 2)

        query = split_heads(query)
        key = split_heads(key)
        value = split_heads(value)

        scores = query @ key.transpose(-2, -1)
        scores = scores * (self.head_dim ** -0.5)
        mask = self.causal_mask[:, :, :n_tokens, :n_tokens]
        scores = scores.masked_fill(mask, float("-inf"))

        if key_is_padding is not None:
            expected = (batch_size, n_tokens)
            if key_is_padding.shape != expected:
                raise ValueError(
                    f"padding mask must have shape {expected}"
                )
            scores = scores.masked_fill(
                key_is_padding[:, None, None, :],
                float("-inf"),
            )

        weights = F.softmax(
            scores,
            dim=-1,
            dtype=torch.float32,
        ).to(query.dtype)
        weights = self.attention_dropout(weights)
        context = weights @ value

        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, n_tokens, self.d_model)
        return self.out_proj(context)


class TransformerBlock(nn.Module):
    def __init__(self, cfg: GPTConfig) -> None:
        super().__init__()
        self.norm_1 = LayerNorm(cfg.d_model)
        self.attention = CausalSelfAttention(cfg)
        self.norm_2 = LayerNorm(cfg.d_model)
        self.feed_forward = FeedForward(cfg)
        self.residual_dropout = nn.Dropout(cfg.dropout)

    def forward(
        self,
        x: torch.Tensor,
        key_is_padding: torch.Tensor | None = None,
    ) -> torch.Tensor:
        attention_out = self.attention(
            self.norm_1(x),
            key_is_padding=key_is_padding,
        )
        x = x + self.residual_dropout(attention_out)
        feed_forward_out = self.feed_forward(self.norm_2(x))
        return x + self.residual_dropout(feed_forward_out)


class GPTModel(nn.Module):
    def __init__(self, cfg: GPTConfig) -> None:
        super().__init__()
        self.cfg = cfg
        self.token_embedding = nn.Embedding(
            cfg.vocab_size,
            cfg.d_model,
        )
        self.position_embedding = nn.Embedding(
            cfg.context_length,
            cfg.d_model,
        )
        self.embedding_dropout = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList(
            TransformerBlock(cfg) for _ in range(cfg.n_layers)
        )
        self.final_norm = LayerNorm(cfg.d_model)
        self.lm_head = nn.Linear(
            cfg.d_model,
            cfg.vocab_size,
            bias=False,
        )

        self.apply(self._initialise)
        residual_std = 0.02 / math.sqrt(2 * cfg.n_layers)
        for block in self.blocks:
            nn.init.normal_(
                block.attention.out_proj.weight,
                mean=0.0,
                std=residual_std,
            )
            nn.init.normal_(
                block.feed_forward.out_proj.weight,
                mean=0.0,
                std=residual_std,
            )

        if cfg.tie_embeddings:
            self.lm_head.weight = self.token_embedding.weight

    @staticmethod
    def _initialise(module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward_features(
        self,
        input_ids: torch.Tensor,
        key_is_padding: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, tokens]")

        batch_size, n_tokens = input_ids.shape
        if n_tokens > self.cfg.context_length:
            raise ValueError("sequence exceeds configured context length")
        if key_is_padding is not None:
            expected = (batch_size, n_tokens)
            if key_is_padding.shape != expected:
                raise ValueError(
                    f"padding mask must have shape {expected}"
                )

        positions = torch.arange(n_tokens, device=input_ids.device)
        token_vectors = self.token_embedding(input_ids)
        position_vectors = self.position_embedding(positions)
        x = self.embedding_dropout(token_vectors + position_vectors)

        for block in self.blocks:
            x = block(
                x,
                key_is_padding=key_is_padding,
            )

        return self.final_norm(x)

    def forward(
        self,
        input_ids: torch.Tensor,
        key_is_padding: torch.Tensor | None = None,
    ) -> torch.Tensor:
        hidden = self.forward_features(
            input_ids,
            key_is_padding=key_is_padding,
        )
        return self.lm_head(hidden)


@torch.inference_mode()
def generate_greedy(
    model: GPTModel,
    input_ids: torch.Tensor,
    max_new_tokens: int,
) -> torch.Tensor:
    was_training = model.training
    model.eval()
    try:
        for _ in range(max_new_tokens):
            window = input_ids[:, -model.cfg.context_length :]
            logits = model(window)
            next_token = logits[:, -1, :].argmax(
                dim=-1,
                keepdim=True,
            )
            input_ids = torch.cat((input_ids, next_token), dim=1)
        return input_ids
    finally:
        model.train(was_training)


def count_parameters(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())


def expected_parameter_count(cfg: GPTConfig) -> int:
    d_model = cfg.d_model
    hidden_size = cfg.ff_multiplier * d_model

    qkv = 3 * d_model * d_model
    if cfg.qkv_bias:
        qkv += 3 * d_model
    attention = qkv + d_model * d_model + d_model

    feed_forward = (
        d_model * hidden_size
        + hidden_size
        + hidden_size * d_model
        + d_model
    )
    two_layer_norms = 4 * d_model
    all_blocks = cfg.n_layers * (
        attention + feed_forward + two_layer_norms
    )

    embeddings = (
        cfg.vocab_size * d_model
        + cfg.context_length * d_model
    )
    final_layer_norm = 2 * d_model
    untied_head = (
        0 if cfg.tie_embeddings else cfg.vocab_size * d_model
    )
    return embeddings + all_blocks + final_layer_norm + untied_head

The three projections for queries, keys and values share one Linear module for efficient implementation, then split along the last dimension. Each tensor is reshaped from $[B,T,D]$ to $[B,H,T,D/H]$. The attention result reverses that reshape before the output projection.

Softmax is evaluated in float32 and cast back to the query dtype. This small implementation choice reduces avoidable numerical trouble in lower precision. More specialised kernels can perform the same operation with different memory and execution strategies.

The initialisation scales the two residual-output projections by $1/\sqrt{2L}$, where $L$ is the number of blocks. There are two residual branches per block. This follows the depth-aware idea described with GPT-2; it is part of the training recipe rather than a change to the forward-pass algebra.

The forward pass, shape by shape

For a batch of two six-token sequences, the default configuration produces:

The logit vector at position $t$ predicts the token at $t+1$ during next-token training. All positions can be processed in parallel because the causal mask prevents information from leaking backwards from future targets.

A synthetic banking fragment such as “Leverage remains within the agreed limit” changes none of this machinery. After tokenisation it is a sequence of identifiers. The architecture can learn statistical relations among those tokens; it does not calculate a leverage ratio or establish that a limit is genuinely satisfied. Those claims need sourced inputs and deterministic calculation outside the base model.

Parameter accounting and weight tying

The parameter ledger should match the code, including biases and tying.

For $D=768$ with QKV biases:

• QKV and attention output projections: 2,362,368 parameters;
• $768 \rightarrow 3{,}072 \rightarrow 768$ feed-forward network: 4,722,432;
• two LayerNorm modules: 3,072; and
• one complete block: 7,087,872.

Twelve blocks therefore contain 85,054,464 parameters. The token table contains 38,597,376, the position table contains 786,432, and final LayerNorm contains 1,536.

The language-model head has shape $[V,D]$, the same shape as the token-embedding table. When tie_embeddings=True, both modules refer to the same Parameter. The head computes dot products against the shared rows; it is not an inverse of the embedding lookup. PyTorch’s parameter iterator counts the shared object once.

The default model in this chapter contains exactly 124,439,808 unique parameters. Untying the head adds $50{,}257 \times 768 = 38{,}597{,}376$, producing 163,037,184. Size labels such as “124M” are rounded, and published reports do not always use identical counting conventions. The executable count is the authority for this implementation.

Autoregressive generation

Generation repeatedly runs the model on a visible prefix, selects a token from the logits at the final position and appends that token. The reference function uses greedy selection, so argmax can operate directly on logits; computing softmax would not change the winner.

The reference generator crops to the configured context window and recomputes the window at every step. That keeps the implementation legible but wastes work. A production decoder commonly returns keys and values from each layer and reuses them on the next step. Cache support changes the attention interface and position handling, so it should be added and tested explicitly rather than implied by a comment.

Greedy decoding is deterministic only under a fixed execution environment and deterministic kernels. Sampling, temperature, top-$k$ and top-$p$ are later decoding choices; they do not change the model architecture. Randomly initialised weights will still yield valid token identifiers, but the decoded sequence has no reason to be coherent. Training supplies learned distributions.

Reference checks

The checks below instantiate a tiny version of the same model. They verify shape, normalisation, causal isolation, a backward pass, weight sharing, generation length and parameter arithmetic. The default 124M-shaped count is tested analytically without allocating the full model.

torch.manual_seed(7)

tiny_cfg = GPTConfig(
    vocab_size=101,
    context_length=16,
    d_model=32,
    n_heads=4,
    n_layers=2,
    dropout=0.0,
)
tiny_model = GPTModel(tiny_cfg)

# Shape and parameter-accounting checks
input_ids = torch.randint(0, tiny_cfg.vocab_size, (2, 7))
logits = tiny_model(input_ids)
assert logits.shape == (2, 7, tiny_cfg.vocab_size)
assert count_parameters(tiny_model) == expected_parameter_count(tiny_cfg)

# The embedding and output head must be one shared Parameter
assert tiny_model.lm_head.weight is tiny_model.token_embedding.weight

# LayerNorm uses population variance over the final axis
sample = torch.randn(3, 5, tiny_cfg.d_model)
normalised = LayerNorm(tiny_cfg.d_model)(sample)
assert torch.allclose(
    normalised.mean(dim=-1),
    torch.zeros(3, 5),
    atol=1e-5,
)
assert torch.allclose(
    normalised.var(dim=-1, unbiased=False),
    torch.ones(3, 5),
    atol=1e-4,
)

# Later tokens must not change logits at earlier positions
tiny_model.eval()
prefix_a = torch.tensor([[1, 2, 3, 4]])
prefix_b = torch.tensor([[1, 2, 9, 10]])
logits_a = tiny_model(prefix_a)
logits_b = tiny_model(prefix_b)
assert torch.allclose(
    logits_a[:, :2],
    logits_b[:, :2],
    atol=1e-6,
    rtol=0.0,
)

# A masked interior key cannot change later real-token logits
padded_a = torch.tensor([[1, 2, 3, 4]])
padded_b = torch.tensor([[1, 99, 3, 4]])
key_is_padding = torch.tensor([[False, True, False, False]])
masked_a = tiny_model(
    padded_a,
    key_is_padding=key_is_padding,
)
masked_b = tiny_model(
    padded_b,
    key_is_padding=key_is_padding,
)
assert torch.allclose(
    masked_a[:, 2:],
    masked_b[:, 2:],
    atol=1e-6,
    rtol=0.0,
)
unmasked_a = tiny_model(padded_a)
unmasked_b = tiny_model(padded_b)
assert not torch.allclose(
    unmasked_a[:, 2:],
    unmasked_b[:, 2:],
    atol=1e-6,
    rtol=0.0,
)

# A next-token loss must backpropagate into the tied embedding
tiny_model.train()
logits = tiny_model(input_ids)
loss = F.cross_entropy(
    logits[:, :-1].reshape(-1, tiny_cfg.vocab_size),
    input_ids[:, 1:].reshape(-1),
)
tiny_model.zero_grad(set_to_none=True)
loss.backward()
assert tiny_model.token_embedding.weight.grad is not None
assert torch.isfinite(loss)

# Greedy generation appends exactly the requested number of tokens
generated = generate_greedy(
    tiny_model,
    input_ids[:1, :3],
    max_new_tokens=4,
)
assert generated.shape == (1, 7)

# The chapter's full configuration has a reproducible unique count
assert expected_parameter_count(GPTConfig()) == 124_439_808

print("All Chapter 4 reference checks passed.")

Expected console output:

All Chapter 4 reference checks passed.

The causal test is especially valuable. Two inputs share positions zero and one but differ afterwards. If their first two output positions change, the mask is wrong or information has leaked through another path. The test does not establish language quality; it establishes one architectural invariant.

Further tests belong beside training code: loss reduction on a tiny corpus, checkpoint save-and-load equivalence, mixed-precision checks, device transfer, seeded reproducibility, tokenizer compatibility and comparison with a trusted checkpoint where exact compatibility is intended.

Modern variations, scoped carefully

The reference model uses early GPT-style choices because they expose the block clearly. Later architectures often replace individual parts:

RMSNorm changes the normalisation equation and usually removes the learned shift. RoPE commonly replaces the absolute position table and changes how queries and keys are formed. SwiGLU changes both the activation and the feed-forward parameter ledger. GQA changes attention tensor shapes and the KV cache formula. Each is a local architectural variation, but none is a search-and-replace operation in a trained checkpoint.

The residual skeleton remains recognisable: normalise, transform, add; normalise, transform, add. That continuity makes the model built here a useful reading tool for newer architectures, provided the differences are examined rather than waved away as implementation detail.

Build check

Before training, confirm the following:

• d_model is divisible by n_heads;
• every block accepts and returns $[B,T,D]$;
• the causal mask prevents later tokens from affecting earlier logits;
• LayerNorm uses variance with unbiased=False;
• train and evaluation modes switch dropout correctly;
• the position range cannot exceed the configured table;
• the tokenizer vocabulary equals vocab_size;
• the language-model head is deliberately tied or deliberately untied;
• the unique parameter count matches the analytic ledger;
• logits have shape $[B,T,V]$;
• the shifted next-token loss uses position $t$ to predict token $t+1$; and
• the model can save, reload and reproduce the same evaluation logits.

These checks catch architecture errors before a long training run turns them into an expensive mystery.

The assembled machine

The decoder now has a complete forward path. Token and position embeddings create the residual stream. Each block uses causal attention to mix earlier context and a feed-forward network to transform each position. LayerNorm conditions both branches, residual additions preserve a direct route, and the final projection produces vocabulary logits.

The untrained architecture contains no corpus-derived patterns. Its random parameters produce arbitrary logits, which the generation loop exposes one token at a time. Chapter 5 supplies the objective, batches, gradients, optimiser and checkpoints that turn this tested structure into a trained language model.

Primary-source notes

1. Vaswani et al., “Attention Is All You Need”, 2017, introduced the Transformer’s attention and position-wise feed-forward structure with residual connections and layer normalisation.
2. Ba, Kiros and Hinton, “Layer Normalization”, 2016, defined normalisation from within-example layer statistics with learned gain and bias.
3. Hendrycks and Gimpel, “Gaussian Error Linear Units”, 2016, defined GELU as $x\Phi(x)$.
4. He et al., “Deep Residual Learning for Image Recognition”, 2015, presented residual learning as a way to ease optimisation of deeper networks.
5. Radford et al., “Language Models are Unsupervised Multitask Learners”, 2019, documented GPT-2’s pre-normalisation placement, final normalisation, 50,257-token vocabulary, 1,024-token context and depth-aware residual initialisation.
6. Press and Wolf, “Using the Output Embedding to Improve Language Models”, 2016, analysed sharing the input embedding with the output embedding.
7. Xiong et al., “On Layer Normalization in the Transformer Architecture”, 2020, compared the initialisation and optimisation behaviour of pre-LN and post-LN Transformers.
8. Zhang and Sennrich, “Root Mean Square Layer Normalization”, 2019; Su et al., “RoFormer”, 2021; Shazeer, “GLU Variants Improve Transformer”, 2020; and Ainslie et al., “GQA”, 2023, are the primary references for the four modern variations summarised above.

One token stream, two shifted views

Causal language modelling supplies its own targets. Given token identifiers

$$[t_0, t_1, t_2, t_3, t_4],$$

the input is

$$[t_0, t_1, t_2, t_3],$$

and the target is the same stream shifted one place to the left:

$$[t_1, t_2, t_3, t_4].$$

At input position 0, the model predicts $t_1$. At position 1, it predicts $t_2$, using $t_0$ and $t_1$. The causal mask prevents position 1 from reading $t_2$ even though all positions are processed in one tensor. A batch with shape [batch, sequence] therefore yields one supervised classification target per non-padding position.

The dataset and loader pattern from Chapter 2 already returns this pair. The loss function can therefore remain small:

import torch
import torch.nn.functional as F


def causal_lm_loss(
    model,
    input_ids,
    target_ids,
    reduction: str = "mean",
):
    """Next-token cross-entropy for a fixed-length batch."""
    logits = model(input_ids)  # [batch, sequence, vocabulary]
    if logits.shape[:2] != target_ids.shape:
        raise ValueError(
            f"Logits {logits.shape[:2]} and targets "
            f"{target_ids.shape} do not align"
        )

    return F.cross_entropy(
        logits.reshape(-1, logits.size(-1)),
        target_ids.reshape(-1),
        reduction=reduction,
    )

reshape merges the batch and sequence dimensions. The vocabulary dimension remains intact, so each row is a classification over every token. If a dataset uses padding, pass an ignore_index and count only unmasked target positions. The fixed-length windows used here contain no padding.

For a small experiment, reducing the model’s context length from GPT-2’s 1,024 tokens to 256 reduces memory and computation. It also changes the model configuration. Before importing GPT-2 weights later in the chapter, the positional-embedding shape and context limit must be restored to 1,024.2

Cross-entropy and perplexity

For one target token $y$, let the model assign probability $p_y$. Its negative log-likelihood is

$$\ell = -\ln(p_y).$$

If $p_y=0.5$, the loss is about 0.69 nats. If $p_y=0.01$, it is about 4.61. The logarithm penalises a model that gives very little probability to the observed token. Averaging this quantity across target positions gives the cross-entropy reported during training.

Uniform predictions provide a useful implementation check. With a vocabulary of 50,257 tokens, a perfectly uniform distribution has loss $\ln(50{,}257)\approx10.82$. Randomly initialised logits need not be exactly uniform, so an observed initial loss may differ. A large difference deserves investigation: the inputs and targets may be misaligned, a mask may leak future tokens, or a numerical value may be non-finite.

Perplexity is the exponential of mean token-level negative log-likelihood:

$$\operatorname{PPL}=\exp(\overline{\ell}).$$

A loss of 3.0 corresponds to perplexity of about 20.1. This transformation does not turn the metric into a count of facts known or an estimate of factual confidence. It is best read as a restatement of average token prediction loss under a particular evaluation pipeline.

That qualification matters. Perplexity changes with the corpus, tokenizer, normalisation, treatment of whitespace, context length and handling of out-of-vocabulary text. A tokenizer that represents the same sentence with more tokens changes the denominator. Results are directly comparable only when the prediction units and data preparation are the same. Perplexity on a held-out short story cannot be compared casually with perplexity on code or financial contracts, and lower perplexity does not guarantee better citation support, reasoning or instruction following.

Use cross-entropy and perplexity to answer a limited question: how well does this model predict these held-out tokens under this tokenisation and context? Other claims need other tests.

Designing the train and validation sets

The split must happen before overlapping windows are created. If a document is windowed first and neighbouring windows are then assigned at random, repeated phrases can appear in both sets. Validation loss will measure recognition of near-duplicates rather than generalisation.

For a collection of documents, split by document, source or time according to the intended deployment. A time-based split is appropriate when the model will encounter later documents. A source-based split tests transfer to a new source. For one continuous story, a contiguous final segment is a defensible teaching split, although conclusions from such a small sample must remain modest.

The data record for a repeatable run should include:

• immutable identifiers or hashes for source documents;
• the tokenizer name, vocabulary and version;
• filtering, normalisation and deduplication rules;
• split membership and the random seed used to assign it;
• context length, stride and treatment of incomplete windows;
• counts of documents, tokens and prediction targets after filtering.

Overlapping windows are legitimate inside the training split. They expose more positions with a full left context, although they also cause some tokens to contribute to several updates. The risk arises when overlapping or duplicated content crosses the split boundary. Record the stride and report tokens processed separately from distinct source tokens so that repetition is visible.

Document packing needs an explicit rule as well. Concatenating unrelated documents without a separator teaches the model an artificial transition from the end of one document to the start of another. Insert a designated end-of-text token, or use attention and loss masks that prevent cross-document prediction. If sensitive and public corpora are mixed, preserve provenance at window level; a later deletion or audit must be able to identify which examples contained a given source.

Keep a test set untouched until model and hyperparameter choices are finished. Repeatedly selecting a checkpoint from “test” results turns that set into validation data.

The following cell makes the teaching run concrete. Download the pinned the-verdict.txt snapshot named in the note, save it beside the program and verify its digest before use. The split occurs on raw text before either partition is tokenised or windowed.

from pathlib import Path
import hashlib

import tiktoken
from torch.utils.data import DataLoader


corpus_path = Path("the-verdict.txt")
expected_sha256 = (
    "b41e41a68f0398a3154ae69e2e4c0e2694e17fe0d"
    "66730536837f1b01935b31f"
)
corpus_bytes = corpus_path.read_bytes()
actual_sha256 = hashlib.sha256(corpus_bytes).hexdigest()
if actual_sha256 != expected_sha256:
    raise ValueError("corpus digest does not match the pinned snapshot")

corpus_text = corpus_bytes.decode("utf-8")
split_at = int(0.90 * len(corpus_text))
train_text = corpus_text[:split_at]
validation_text = corpus_text[split_at:]

encoding = tiktoken.get_encoding("gpt2")
train_ids = encoding.encode(
    train_text,
    disallowed_special=(),
)
validation_ids = encoding.encode(
    validation_text,
    disallowed_special=(),
)

context_length = 256
stride = 256
train_dataset = NextTokenDataset(
    train_ids,
    context_length=context_length,
    stride=stride,
)
validation_dataset = NextTokenDataset(
    validation_ids,
    context_length=context_length,
    stride=stride,
)

train_generator = torch.Generator().manual_seed(123)
train_loader = DataLoader(
    train_dataset,
    batch_size=2,
    shuffle=True,
    drop_last=False,
    generator=train_generator,
)
train_eval_loader = DataLoader(
    train_dataset,
    batch_size=2,
    shuffle=False,
    drop_last=False,
)
val_loader = DataLoader(
    validation_dataset,
    batch_size=2,
    shuffle=False,
    drop_last=False,
)

cfg = GPTConfig(context_length=context_length)
model = GPTModel(cfg)

Validation should run with dropout disabled and gradient tracking off. The following function sums token losses before dividing, rather than averaging batch means. That detail gives every target token equal weight when the last batch is smaller.

@torch.inference_mode()
def mean_nll(model, data_loader, device):
    was_training = model.training
    model.eval()
    total_nll = 0.0
    total_tokens = 0

    try:
        for input_ids, target_ids in data_loader:
            input_ids = input_ids.to(device)
            target_ids = target_ids.to(device)
            logits = model(input_ids)

            total_nll += F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                target_ids.reshape(-1),
                reduction="sum",
            ).item()
            total_tokens += target_ids.numel()
    finally:
        model.train(was_training)

    if total_tokens == 0:
        raise ValueError("Cannot evaluate an empty dataloader")
    return total_nll / total_tokens

Record both train and validation loss against optimiser updates and tokens processed. Epoch numbers alone are ambiguous when datasets or accumulation settings change.

An optimisation step

One update has a strict order:

1. run the forward pass;
2. calculate mean next-token loss;
3. backpropagate gradients;
4. optionally unscale and clip those gradients;
5. update parameters and the learning-rate schedule;
6. clear gradients for the next effective batch.

AdamW is a practical optimiser for this build. It maintains moving estimates of the first and second gradient moments and applies weight decay separately from the loss gradient, as proposed by Loshchilov and Hutter.3 Its settings are hyperparameters, not universal constants. In a larger run, biases and normalisation parameters are often placed in a parameter group without weight decay.

Three refinements make the update safer to reuse across run sizes.

Warmup followed by decay

Early gradients are produced by uncalibrated activations and optimiser moment estimates. A warmup raises the learning rate gradually. After warmup, cosine decay reduces it towards a floor as the update budget is consumed. This schedule is one sound default, not a theorem about every dataset. Cosine-shaped annealing was explored in the Stochastic Gradient Descent with Warm Restarts (SGDR) experiments; our schedule uses one decay without warm restarts.5

The scheduler below expresses the rate as a multiplier of the optimiser’s peak learning rate:

import math


def lr_multiplier(step, warmup_steps, total_steps, min_ratio=0.1):
    if (
        step < 0
        or total_steps <= 0
        or not 0 <= warmup_steps <= total_steps
        or not 0.0 <= min_ratio <= 1.0
    ):
        raise ValueError("Invalid schedule")
    if step < warmup_steps:
        return (step + 1) / max(1, warmup_steps)

    decay_steps = max(1, total_steps - warmup_steps)
    progress = min(
        1.0,
        (step - warmup_steps + 1) / decay_steps,
    )
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_ratio + (1.0 - min_ratio) * cosine

Here step is a zero-based update index. If the run executes total_steps updates, the final index is total_steps - 1, where the multiplier reaches min_ratio. The warmup’s final update reaches the peak; the first decay update then begins moving below it.

Mixed precision without silent numerical errors

Automatic mixed precision (AMP) lets selected operations use float16 or bfloat16 while numerically sensitive work stays in float32. It can reduce memory use and improve throughput on supported hardware. The result depends on the device and kernels; casting the entire model to half precision by hand is not an equivalent procedure.

float16 has limited exponent range. Small gradients can underflow, so a gradient scaler multiplies the loss before backpropagation and reverses that scale before the update. bfloat16 has a wider exponent range and commonly does not need loss scaling, but it has fewer mantissa bits. Neither format is guaranteed to work for every model.6

When clipping scaled gradients, unscale them once, after all micro-batches have accumulated and immediately before clipping. Updating the scaler inside an accumulation window would mix gradients with different scales. The following single-device loop implements that order. It enables automatic mixed precision only on CUDA and otherwise runs in float32.7

def train_one_epoch(
    model,
    train_loader,
    optimizer,
    scheduler,
    scaler,
    device,
    accumulation_steps,
    max_grad_norm,
    global_step,
    tokens_seen,
    use_amp,
    amp_dtype,
):
    if accumulation_steps < 1 or len(train_loader) == 0:
        raise ValueError("Training requires batches and accumulation >= 1")
    if max_grad_norm <= 0:
        raise ValueError("max_grad_norm must be positive")

    model.train()
    optimizer.zero_grad(set_to_none=True)
    batch_count = len(train_loader)
    window_tokens = 0

    for batch_index, (input_ids, target_ids) in enumerate(train_loader):
        input_ids = input_ids.to(device)
        target_ids = target_ids.to(device)

        window_end = min(
            (
                batch_index // accumulation_steps + 1
            ) * accumulation_steps,
            batch_count,
        )

        with torch.autocast(
            device_type="cuda",
            dtype=amp_dtype,
            enabled=use_amp,
        ):
            loss = causal_lm_loss(
                model,
                input_ids,
                target_ids,
                reduction="sum",
            )

        scaler.scale(loss).backward()
        token_count = target_ids.numel()
        window_tokens += token_count
        tokens_seen += token_count

        if batch_index + 1 != window_end:
            continue

        scaler.unscale_(optimizer)
        for parameter in model.parameters():
            if parameter.grad is not None:
                parameter.grad.div_(window_tokens)
        torch.nn.utils.clip_grad_norm_(
            model.parameters(),
            max_norm=max_grad_norm,
        )

        old_scale = scaler.get_scale()
        scaler.step(optimizer)
        scaler.update()
        step_was_skipped = (
            scaler.is_enabled()
            and scaler.get_scale() < old_scale
        )

        if not step_was_skipped:
            scheduler.step()
            global_step += 1

        optimizer.zero_grad(set_to_none=True)
        window_tokens = 0

    return global_step, tokens_seen

Here is the corresponding setup:

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)
model.to(device)

accumulation_steps = 4
epochs = 5
updates_per_epoch = math.ceil(
    len(train_loader) / accumulation_steps
)
total_updates = epochs * updates_per_epoch
warmup_updates = max(1, int(0.1 * total_updates))

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=4e-4,
    betas=(0.9, 0.95),
    weight_decay=0.1,
)
scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=lambda step: lr_multiplier(
        step, warmup_updates, total_updates
    ),
)

use_amp = device.type == "cuda"
amp_dtype = (
    torch.bfloat16
    if use_amp and torch.cuda.is_bf16_supported()
    else torch.float16
)
use_scaler = use_amp and amp_dtype == torch.float16
scaler = torch.amp.GradScaler("cuda", enabled=use_scaler)

The numerical safeguards still need monitoring. Log non-finite losses, gradient norm before clipping, the scaler value, learning rate and skipped updates. Compare validation loss from a mixed-precision run with a float32 baseline before assuming equivalence.

A staged run protocol

A long run should be the last step in the experiment, not the first. A short sequence of tests catches most implementation errors while they remain easy to diagnose.

Check the data contract. Decode several input and target windows. Confirm that every target is shifted by one token, document separators appear where expected, and no validation source occurs in training. Count targets after masking rather than inferring the number from file size.

Check the random baseline. With dropout disabled, calculate loss on fixed training and validation batches before any update. Compare the result with the uniform-loss reference and inspect the logits for finite values. The aim is not to force exact uniformity; it is to detect a gross shape, masking or target bug.

Overfit one batch deliberately. Reuse one small batch for enough updates to drive its loss down. If the optimiser cannot fit a single batch, increasing the dataset or runtime will not help. This diagnostic may reveal detached tensors, parameters excluded from the optimiser, an incorrect causal mask or a learning rate that produces non-finite values. Discard the resulting weights.

Test resume equivalence. Run a few updates, save, load into newly constructed objects and continue. Compare the next batch, learning rate, scaler and loss with an uninterrupted run on the same software and device. Exact equivalence may require a stateful sampler and deterministic kernels, but a large difference still exposes missing state.

Run a bounded pilot. Use the intended train and validation paths for a small fraction of the update budget. Inspect throughput, memory, clipping rate, skipped AMP updates and fixed-prompt samples. Only then schedule the full educational run.

Every run needs a unique identifier and an immutable configuration record. Log the source revision, model configuration, optimiser and schedule, tokenizer, dataset fingerprint, seed, device type and software versions beside the metrics. A loss curve without that context cannot be reproduced or compared fairly.

Reading the loss curves

Training loss should fall as the optimiser fits the sampled windows. Validation loss often falls at first and then flattens. A widening gap can indicate overfitting, but it can also reflect a domain mismatch, duplicate removal that changed one split more than the other, or noisy validation data.

For this deliberately tiny single-story run, memorisation is expected: the model has far more capacity than the corpus can support. The run demonstrates mechanics, not a general-purpose language model. Stop according to a stated rule, such as the lowest smoothed validation loss within a fixed update budget. Keep the selected checkpoint and report the selection rule; do not inspect the test set after every update or use it to select a checkpoint.

Text samples complement loss curves. Generate from a fixed set of prompts with both greedy decoding and a fixed sampling seed. Samples can expose repetition, broken Unicode or data leakage that average loss hides. They remain diagnostic examples rather than a statistically adequate evaluation.

Compute-efficient scaling also requires balancing model capacity and training tokens. The Chinchilla experiments found that, under their model family and compute budgets, increasing only parameter count left performance on the table.8 That result motivates recording both parameters and tokens processed; it does not supply a universal token-to-parameter ratio for every corpus.

Checkpointing a training trajectory

Saving only model.state_dict() is enough to distribute weights for inference. It is not enough to resume training faithfully. AdamW has moment estimates; the scheduler has a position; a float16 gradient scaler has a scale; random-number generators (RNGs) determine dropout, sampling and data order.

A resumable checkpoint should contain:

• model, optimiser, scheduler and scaler state;
• completed optimiser step, epoch, batch cursor and tokens processed;
• CPU, CUDA, Python and NumPy random-number-generator state when those generators are used;
• the model configuration and a fingerprint of the dataset manifest;
• the state of the generator or sampler that determines training order; and
• sampler cursor state for an exact mid-epoch resume, especially with distributed shuffling or worker prefetching.

The implementation below writes to a temporary path and replaces the previous checkpoint only after serialisation succeeds.

import os
import random
from pathlib import Path

import numpy as np


def capture_rng_state():
    return {
        "torch_cpu": torch.get_rng_state(),
        "torch_cuda": (
            torch.cuda.get_rng_state_all()
            if torch.cuda.is_available()
            else None
        ),
        "python": random.getstate(),
        "numpy": np.random.get_state(),
    }


def save_checkpoint(
    path,
    *,
    model,
    optimizer,
    scheduler,
    scaler,
    global_step,
    tokens_seen,
    data_state,
    model_config,
):
    path = Path(path)
    temporary = path.with_suffix(path.suffix + ".tmp")
    state = {
        "schema_version": 1,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "scaler": scaler.state_dict(),
        "global_step": global_step,
        "tokens_seen": tokens_seen,
        "data_state": data_state,
        "model_config": model_config,
        "rng": capture_rng_state(),
    }
    torch.save(state, temporary)
    os.replace(temporary, path)


def load_checkpoint(
    path, *, model, optimizer, scheduler, scaler, device
):
    # weights_only=False is safe only for a checkpoint you trust.
    state = torch.load(
        path,
        map_location=device,
        weights_only=False,
    )
    if state.get("schema_version") != 1:
        raise ValueError("Unsupported checkpoint schema")

    model.load_state_dict(state["model"])
    optimizer.load_state_dict(state["optimizer"])
    scheduler.load_state_dict(state["scheduler"])
    scaler.load_state_dict(state["scaler"])

    rng = state["rng"]
    torch.set_rng_state(rng["torch_cpu"].cpu())
    random.setstate(rng["python"])
    np.random.set_state(rng["numpy"])
    if rng["torch_cuda"] is not None:
        if len(rng["torch_cuda"]) != torch.cuda.device_count():
            raise RuntimeError("CUDA topology differs from checkpoint")
        torch.cuda.set_rng_state_all(
            [value.cpu() for value in rng["torch_cuda"]]
        )

    return {
        "global_step": state["global_step"],
        "tokens_seen": state["tokens_seen"],
        "data_state": state["data_state"],
        "model_config": state["model_config"],
    }

data_state should contain the next epoch or a batch cursor together with the training-order generator or sampler state. A stateless dataloader cannot resume at a mid-epoch cursor by itself; either restore a stateful sampler or recreate the same shuffle and skip the already processed batches. Multi-worker prefetching makes the requirement stricter.

After saving, test the checkpoint in a fresh process: load it, calculate loss on a fixed validation batch and perform one controlled update. Keep more than one recent checkpoint and a separate milestone copy. Reproducibility is bounded by software, hardware and kernel choices; identical seeds do not guarantee bitwise identity across different platforms or framework releases.9

At epoch boundaries, a compact driver can evaluate and save both the latest trajectory and the checkpoint selected by validation loss:

def fit(
    *,
    model,
    train_loader,
    train_eval_loader,
    train_generator,
    val_loader,
    optimizer,
    scheduler,
    scaler,
    device,
    epochs,
    accumulation_steps,
    max_grad_norm,
    use_amp,
    amp_dtype,
    checkpoint_dir,
    model_config,
    dataset_fingerprint,
    resume_state=None,
):
    checkpoint_dir = Path(checkpoint_dir)
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    history = []
    resume_state = resume_state or {}
    restored_data = resume_state.get("data_state", {})
    global_step = resume_state.get("global_step", 0)
    tokens_seen = resume_state.get("tokens_seen", 0)
    start_epoch = restored_data.get("next_epoch", 0)
    best_validation = restored_data.get(
        "best_validation", float("inf")
    )
    generator_state = restored_data.get(
        "train_generator_state"
    )
    if generator_state is not None:
        train_generator.set_state(generator_state.cpu())

    for epoch in range(start_epoch, epochs):
        global_step, tokens_seen = train_one_epoch(
            model,
            train_loader,
            optimizer,
            scheduler,
            scaler,
            device,
            accumulation_steps,
            max_grad_norm,
            global_step,
            tokens_seen,
            use_amp,
            amp_dtype,
        )

        train_nll = mean_nll(
            model, train_eval_loader, device
        )
        validation_nll = mean_nll(model, val_loader, device)
        history.append(
            {
                "epoch": epoch + 1,
                "step": global_step,
                "tokens_seen": tokens_seen,
                "train_nll": train_nll,
                "validation_nll": validation_nll,
            }
        )

        improved = validation_nll < best_validation
        if improved:
            best_validation = validation_nll

        data_state = {
            "next_epoch": epoch + 1,
            "batch_in_epoch": 0,
            "dataset_fingerprint": dataset_fingerprint,
            "best_validation": best_validation,
            "train_generator_state": (
                train_generator.get_state().cpu()
            ),
        }
        save_args = {
            "model": model,
            "optimizer": optimizer,
            "scheduler": scheduler,
            "scaler": scaler,
            "global_step": global_step,
            "tokens_seen": tokens_seen,
            "data_state": data_state,
            "model_config": model_config,
        }
        save_checkpoint(
            checkpoint_dir / "latest.pt",
            **save_args,
        )
        if improved:
            save_checkpoint(
                checkpoint_dir / "best-validation.pt",
                **save_args,
            )

    return history

To resume, first construct the same model, optimiser, scheduler, scaler, dataloaders and train_generator; call load_checkpoint; verify its model configuration and dataset fingerprint; then pass the returned dictionary as resume_state. The driver restores the generator before it creates the next training iterator. The separate, non-shuffled train_eval_loader prevents metric collection from consuming that generator. Do not silently resume on a different token order or schedule budget. Pass dataclasses.asdict(cfg) as model_config so the next chapter can reconstruct GPTConfig explicitly.

For a large corpus, full train-set evaluation at every epoch may be too expensive. Use a fixed, representative monitoring subset between less frequent full evaluations, and label the two measurements distinctly. The selected checkpoint should still be chosen by a declared validation procedure rather than by whichever generated sample looks best.

Temperature and top-k decoding

Training produces a conditional probability distribution. Decoding decides how to select from it.

With greedy decoding, choose the largest logit. The output is deterministic for a deterministic forward pass. With temperature sampling, divide logits by a positive temperature before softmax:

$$P_i(T)=\frac{\exp(z_i/T)} {\sum_j \exp(z_j/T)}.$$

Temperatures below 1 sharpen the distribution; temperatures above 1 flatten it. Temperature preserves the logit ranking. It adds neither knowledge nor factual reliability.

Top-$k$ filtering retains only the $k$ largest logits and sets the rest to negative infinity before softmax. It removes the low-probability tail but can also exclude a valid rare token. Both parameters belong in evaluation records.

This batch-size-one generator makes the two controls explicit:

@torch.inference_mode()
def generate(
    model,
    input_ids,
    *,
    max_new_tokens,
    context_size,
    temperature=0.0,
    top_k=None,
    eot_id=None,
    generator=None,
):
    if input_ids.size(0) != 1:
        raise ValueError("This generator expects batch size one")
    if temperature < 0:
        raise ValueError("temperature must be non-negative")
    if top_k is not None and top_k < 1:
        raise ValueError("top_k must be positive")

    was_training = model.training
    model.eval()
    output = input_ids

    try:
        for _ in range(max_new_tokens):
            context = output[:, -context_size:]
            logits = model(context)[:, -1, :]

            if top_k is not None:
                k = min(top_k, logits.size(-1))
                threshold = torch.topk(
                    logits, k
                ).values[:, -1, None]
                logits = logits.masked_fill(
                    logits < threshold,
                    float("-inf"),
                )

            if temperature == 0:
                next_id = logits.argmax(dim=-1, keepdim=True)
            else:
                probabilities = torch.softmax(
                    logits / temperature,
                    dim=-1,
                )
                next_id = torch.multinomial(
                    probabilities,
                    num_samples=1,
                    generator=generator,
                )

            output = torch.cat((output, next_id), dim=1)
            if eot_id is not None and next_id.item() == eot_id:
                break
    finally:
        model.train(was_training)

    return output

For a repeatable sample, create a device-compatible torch.Generator with a fixed seed. For an application, choose decoding settings from task evaluation. A lower temperature may reduce variation, but it cannot ground an answer in documents that were never supplied.

Evaluating more than token loss

A pretraining report should contain at least four views:

1. token-weighted train and validation loss on versioned splits;
2. loss by relevant domain, document length and other known slices;
3. fixed-prompt samples under recorded decoding settings;
4. tests for the intended downstream properties, including abstention and evidence use where those properties matter.

The fourth view prevents perplexity from becoming a proxy for everything. Consider a fictional test harness for the credit-document assistant introduced earlier. It represents the control demands of a large UK bank; it does not describe a live system, internal experiment, client data or employer policy. Its expected result uses the book’s canonical covenant fields: metric, operator, threshold and test_frequency.

The generator must return a constrained record:

SYNTHETIC_CASES = [
    {
        "id": "operative-amendment",
        "question": "State the current interest-cover covenant.",
        "evidence": [
            {
                "source": "AMEND-02#4.1",
                "text": (
                    "From 30 June 2026, interest cover "
                    "must be at least 3.0x, tested quarterly."
                ),
            }
        ],
        "expected_fields": {
            "metric": "interest_cover",
            "operator": ">=",
            "threshold": 3.0,
            "test_frequency": "quarterly",
        },
        "expected_citations": ["AMEND-02#4.1"],
        "should_abstain": False,
    },
    {
        "id": "missing-operative-document",
        "question": "State the current interest-cover covenant.",
        "evidence": [],
        "expected_fields": {},
        "expected_citations": [],
        "should_abstain": True,
    },
]


def grade_case(case, result):
    abstention_ok = (
        result.get("abstain") == case["should_abstain"]
    )
    if case["should_abstain"]:
        return abstention_ok and not result.get("fields")

    fields_ok = all(
        result.get("fields", {}).get(name) == value
        for name, value in case["expected_fields"].items()
    )
    citations_ok = set(case["expected_citations"]).issubset(
        result.get("citations", [])
    )
    return abstention_ok and fields_ok and citations_ok