Build an LLM from scratch

Next-token prediction

Frame LMs as next-token predictors

55 min2/8 in module

Learning objectives

  • Frame LMs as next-token predictors
  • Implement a bigram language model
  • Sample text and inspect learned table statistics

The objective that scales to GPT

Every large language model — GPT-4, Claude, LLaMA — is trained on one deceptively simple objective: predict the next token given all previous tokens. There is no separate module for "reasoning" or "being helpful" in pretraining. There is next-token prediction at massive scale, followed by post-training that steers behavior.

That simplicity is why this module starts here, in miniature. If you understand next-token prediction on a bigram table, you understand the skeleton of a billion-parameter transformer. The architecture grows; the training target does not.

Formally, a language model assigns a probability distribution over the vocabulary at each position:

P(token_t | token_1, token_2, ..., token_{t-1})

In a bigram model, you make a strong simplifying assumption: only the previous token matters.

P(token_t | token_{t-1})

That is Markov order 1. Trigram models use two previous tokens; n-gram models use n−1. Transformers relax the fixed window entirely — but the prediction target remains the same.

Callout — generation is repeated prediction: To generate text, sample one token, append it to the context, predict the next, repeat. ChatGPT is this loop with a very good P and fancy decoding.

From counts to probabilities

Training a bigram model is counting. Scan the tokenized corpus; for each adjacent pair (a, b), increment count[a][b]. Convert counts to probabilities:

P(b | a) = count(a, b) / count(a)

Add smoothing so unseen pairs do not get zero probability. Laplace (add-one) smoothing:

P(b | a) = (count(a, b) + 1) / (count(a) + V)

where V is vocabulary size. Without smoothing, any pair absent from training has P = 0 — generation gets stuck forever.

Worked example

Corpus (word-level for clarity): "the cat sat on the mat the cat"

Bigram counts from "the":

Next word Count
cat 2
mat 1

Raw: P(cat | the) = 2/3, P(mat | the) = 1/3.

At generation time, if current word is "the", sample "cat" with 67% probability and "mat" with 33%.

Perplexity: how surprised is the model?

Perplexity is the standard metric for language models. It asks: on average, how many tokens wide is the model's effective choice at each step?

Perplexity = exp(average cross-entropy loss)

Lower is better. A perplexity of 100 means the model is as uncertain as choosing uniformly among 100 tokens. A perplexity of 3 means it is usually picking among ~3 plausible continuations.

For a bigram model on a tiny corpus, perplexity will look good on training data and terrible on held-out text — the table is memorizing pairs, not learning grammar. That gap foreshadows overfitting and the need for neural models with generalization.

Sampling text from a bigram LM

Generation loop:

def generate(start_token, length, bigram_probs, sample_fn):
    tokens = [start_token]
    for _ in range(length - 1):
        prev = tokens[-1]
        probs = bigram_probs[prev]  # dict: next_token -> probability
        next_t = sample_fn(probs)
        tokens.append(next_t)
    return tokens

Sampling strategies (preview of lesson 3.6):

  • Greedy — always pick argmax. Repetitive, boring.
  • Multinomial — sample from P. Stochastic, more natural.
  • Temperature — sharpen or flatten P before sampling.

Bigram output looks like word salad with local coherence ("the cat the cat the mat") but no long-range structure. That failure mode motivates neural LMs: you need more context than one token.

Bigrams vs neural LMs: what you gain

Property Bigram Neural LM (coming next)
Context 1 token Hundreds to millions
Storage Sparse count table Dense weight matrices
Generalization None for unseen n-grams Similar contexts share statistics
Training Count + normalize Gradient descent on cross-entropy
Perplexity on tiny data Can look OK Can overfit too, but captures patterns

The bigram model teaches the data pipeline: tokenize → build training pairs → estimate P → sample. You will reuse that pipeline when you swap the count table for an MLP and later a transformer.

Inspecting learned statistics

After training, inspect your table — do not only generate text.

Questions to answer:

  • What are the top continuations for "the", ".", and your corpus-specific tokens?
  • Which pairs have high count but low probability (because the predecessor is common)?
  • Where does smoothing dominate (rare predecessors)?

A heatmap of the top 50 tokens × top 50 successors reveals structure: punctuation patterns, common bigrams ("of the", "in the"), and corpus artifacts.

Callout — the table is the model: For a bigram LM, there are no hidden layers. Every behavior is readable from counts. Use this phase to build intuition before weights become opaque.

Engineering problem (staff framing)

Causal LMs are trained as next-token predictors. Everything else (chat, tools, RAG) is scaffolding on that objective.

Diagram — Causal LM objective

flowchart LR
  C[Context tokens] --> M[Model]
  M --> Softmax[Softmax over vocab]
  Softmax --> Sample[Sample / argmax]
  Sample --> Append[Append token]
  Append --> C

Precise definitions & mental model

  • Causal mask — position i attends ≤ i.
  • Cross-entropy — −log p(true next token).
  • Teacher forcing — train on gold prefixes.
  • Perplexity — exp(mean CE); useful but not aligned with product quality alone.

Tradeoffs — when to use what

Teacher forcing (stable train) vs free-running (exposes exposure bias).

Failure modes (interview + on-call)

Off-by-one shift in labels; leaking future tokens in mask; optimizing PPL only.

Production & OSS practices

Log CE on held-out domain corpora as a canary when swapping models.

Deep dive (FAANG / OSS bar)

Loss on the shifted sequence

Labels are tokens x_1…x_T predicting x_2…x_{T+1} (or equivalent packing). An off-by-one in the shift is the classic silent bug: loss decreases while generations are garbage.

Packing and EOS

Production pretraining packs documents with EOS separators. Toy models that omit EOS learn to ramble across sample boundaries. Even in mini-GPT, insert a separator token between poems/files.

Exposure bias (awareness level)

Teacher forcing trains on gold prefixes; at inference the model consumes its own errors. Mitigations (scheduled sampling, RL, better decoding) exist; for this course, know the name and why beam search is not a free lunch for open-ended chat.

Micro-project: Bigram LM

In m3/bigram/ (or your course-portfolio monorepo):

  1. Tokenize a corpus using your tokenizer from lesson 3.1 (or a simple whitespace split for a first pass).
  2. Build bigram counts and convert to smoothed probabilities. Save as bigram.json or a sparse matrix.
  3. Implement generate(start_token_id, max_tokens) with multinomial sampling.
  4. Generate at least 200 tokens from three different start tokens; save samples in samples.txt.
  5. Compute perplexity on a held-out chunk of the same corpus.
  6. Write STATS.md: top-5 successors for 5 chosen tokens, perplexity number, and one paragraph on why output does or does not look like your corpus.

Optional: compare Laplace vs no smoothing on an held-out pair that never appeared in training — show zero-probability failure.

Checklist

  • Bigram counts and smoothed probabilities saved
  • Generation produces varied (non-greedy) samples
  • Perplexity computed on held-out text
  • STATS.md with table inspection and qualitative assessment
Project checklist0/3 done

ShipAI delivery model is: