Dictionary of Applied Machine Learning · large language model

large language model — Python demo

Numerical companion to the entry large language model: it recomputes what the entry states and prints one line per check

One block per paragraph of the entry (marked [P...]): each block verifies numerically what the corresponding statement asserts. Self-contained (numpy/matplotlib only), fixed seed. (The entry's scale claim — billions of parameters — is illustrated in miniature: the mechanisms, not the size.)

Run it with python3 llm.py, from any directory — it writes its output files into the current directory. Requires NumPy and Matplotlib only, and uses fixed seeds, so the printed numbers reproduce exactly. Download llm.py

The script, block by block

One cell per block of the script: the code, and what that code printed when it last ran here

setup

"""
llm.py — numerical companion to the glossary entry
'large language model (LLM)'.

One block per paragraph of the entry (marked [P...]): each block verifies
numerically what the corresponding statement asserts. Self-contained
(numpy/matplotlib only), fixed seed. (The entry's scale claim —
billions of parameters — is illustrated in miniature: the mechanisms,
not the size.)

Blocks
------
[P-selfsup]   Self-supervised construction of labeled data points from
              raw text: masking words turns an unannotated corpus into
              (context features, masked-word label) pairs — one per
              position, with zero human labeling effort.
[P-train]     Training via ERM on these pairs: a small next-token model
              (embedding + softmax over the vocabulary) trained by
              gradient descent on the training loss (the negative log probability of the correct next token) drives
              training loss down and beats the uniform-guess baseline.
[P-nexttoken] A trained LLM maps an input token sequence to a
              probability distribution over the next token: outputs are
              nonnegative, sum to one, and the model assigns the
              highest probability to continuations seen in the corpus;
              sampling from the distribution generates text.

Outputs
-------
llm.png : preview figure (checking only).

Data generated by pythondemos/llm.py.
"""

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
report = []


def check(name, ok):
    report.append((name, bool(ok)))
    print(f"  [{'ok' if ok else 'FAIL'}] {name}")


corpus = ("all human beings are born free and equal in dignity and "
          "rights all human beings are endowed with reason and "
          "conscience").split()
vocab = sorted(set(corpus))
V = len(vocab)
tok = {w: i for i, w in enumerate(vocab)}
ids = np.array([tok[w] for w in corpus])

P-selfsup

Self-supervised construction of labeled data points from raw text: masking words turns an unannotated corpus into (context features, masked-word label) pairs — one per position, with zero human labeling effort.

print("[P-selfsup] masked words become labels, contexts become features")
pairs = [(ids[t], ids[t + 1]) for t in range(len(ids) - 1)]
check("one labeled pair per corpus position (no human annotation)",
      len(pairs) == len(corpus) - 1)
check("labels are drawn from the text itself",
      all(0 <= y < V for _, y in pairs))

[P-selfsup] masked words become labels, contexts become features
  [ok] one labeled pair per corpus position (no human annotation)
  [ok] labels are drawn from the text itself

P-train

Training via ERM on these pairs: a small next-token model (embedding + softmax over the vocabulary) trained by gradient descent on the training loss (the negative log probability of the correct next token) drives training loss down and beats the uniform-guess baseline.

print("[P-train] ERM on the constructed pairs")
d_emb = 8
E = 0.1 * rng.normal(size=(V, d_emb))              # embeddings
U = 0.1 * rng.normal(size=(d_emb, V))              # unembedding
def forward(x_ids):
    logits = E[x_ids] @ U
    p = np.exp(logits - logits.max(1, keepdims=True))
    return p / p.sum(1, keepdims=True)
xs = np.array([x for x, _ in pairs])
ys = np.array([y for _, y in pairs])
def xent():
    return -np.mean(np.log(forward(xs)[np.arange(len(ys)), ys] + 1e-12))
loss0 = xent()
for _ in range(800):                               # gradient descent
    P = forward(xs)
    G = P.copy(); G[np.arange(len(ys)), ys] -= 1.0
    gU = E[xs].T @ G / len(ys)
    gE = np.zeros_like(E)
    np.add.at(gE, xs, G @ U.T / len(ys))
    U -= 2.0 * gU; E -= 2.0 * gE
loss1 = xent()
print(f"    training loss: init {loss0:.2f} -> trained {loss1:.2f} "
      f"(uniform baseline {np.log(V):.2f})")
check("training reduces the training loss", loss1 < loss0)
check("the trained model beats the uniform-guess baseline log|V|",
      loss1 < np.log(V) - 0.5)

[P-train] ERM on the constructed pairs
    training loss: init 2.71 -> trained 0.23 (uniform baseline 2.71)
  [ok] training reduces the training loss
  [ok] the trained model beats the uniform-guess baseline log|V|

P-nexttoken

A trained LLM maps an input token sequence to a probability distribution over the next token: outputs are nonnegative, sum to one, and the model assigns the highest probability to continuations seen in the corpus; sampling from the distribution generates text.

print("[P-nexttoken] input sequence -> distribution over the next token")
p_next = forward(np.array([tok["human"]]))[0]
check("the output is a probability distribution (nonneg, sums to 1)",
      np.all(p_next >= 0) and np.isclose(p_next.sum(), 1.0))
check("'human' is followed by 'beings' in the corpus — and gets the "
      "highest next-token probability",
      vocab[int(np.argmax(p_next))] == "beings")
gen = [tok["all"]]
for _ in range(5):                                 # sample a continuation
    gen.append(int(rng.choice(V, p=forward(np.array([gen[-1]]))[0])))
check("sampling from the distributions generates a token sequence",
      len(gen) == 6 and all(0 <= g < V for g in gen))
print("    generated:", " ".join(vocab[g] for g in gen))

# ------------------------------------------------------------ preview
fig, ax = plt.subplots(figsize=(5.6, 3.0))
ax.bar(range(V), p_next)
ax.set_xticks(range(V))
ax.set_xticklabels(vocab, rotation=90, fontsize=6)
ax.set_ylabel("P(next token | 'human')")
ax.set_title("[P-nexttoken] next-token distribution")
fig.tight_layout()
fig.savefig("llm.png", dpi=110)
print(f"\n{sum(ok for _, ok in report)}/{len(report)} checks passed")
assert all(ok for _, ok in report)

[P-nexttoken] input sequence -> distribution over the next token
  [ok] the output is a probability distribution (nonneg, sums to 1)
  [ok] 'human' is followed by 'beings' in the corpus — and gets the highest next-token probability
  [ok] sampling from the distributions generates a token sequence
    generated: all human beings are endowed with

7/7 checks passed