microlm

A small language model written from scratch in NumPy, plus a local control surface for taking it apart while it runs.

No PyTorch, no TensorFlow, no autograd. The transformer forward pass and every gradient in the backward pass are written by hand and verified against finite differences. The whole model is 843,520 parameters and trains in about 10 minutes on a laptop CPU.

843,520 params · 4 layers × 4 heads × d128 · context 64 · vocab 328
best validation loss 0.4796  (perplexity 1.62)

Why it looks like this

Most "build a small LLM" projects wrap a framework. This one doesn't, for one reason: if you want to grab a single weight and change it while the model is answering you, it helps enormously that the weights are just arrays in a dict and the math is in a file you can read. The control surface below is only possible because nothing is hidden.

The model is deliberately tiny and the corpus is deliberately narrow. At this size a model cannot learn language, but it can learn the shape of a short conversation — whose turn it is, which reply belongs to which prompt, and when to stop talking. That is the whole target.

Quickstart

Requires Python 3.9+ and NumPy. That is the complete dependency list.

A trained checkpoint is committed, so you can go straight to the dashboard:

pip install numpy
python -m server.app            # open http://127.0.0.1:8000

To rebuild it from nothing:

python data/build_corpus.py     # generate the chit-chat corpus
python -m microlm.train         # train (~10 min), writes checkpoints/model.npz

Verify the hand-written gradients first if you like:

python -m tests.test_grad

1,448 params checked across 28 tensors
All gradients match finite differences (rel err < 1e-05).

Sample output

After training, every reply below is generated, not canned:

user: hello                        bot: Hi, how's it going?
user: how are you                  bot: I'm good, thanks. What about you?
user: what is your name            bot: I'm micro, a very small language model.
user: tell me a joke               bot: My sense of humor didn't fit in the weights.
user: are you chatgpt              bot: Not very. I'm a few hundred thousand parameters.
user: i had a long day             bot: That's rough, take it easy.
user: bye                          bot: Goodbye, take care.

The control surface

python -m server.app serves a local dashboard (standard-library HTTP server, no framework) that holds two copies of the weights: the pristine checkpoint and a live working copy the UI is free to mutate. Every edit hits the live copy, so reverting is always one click away.

Weights — every tensor as a heatmap, click any cell to read its exact value and index, then set / scale / zero / +noise it. Whole-tensor operations include magnitude pruning and reverting a single tensor. Large tensors are block-averaged for display and edits apply to the whole block, which the UI tells you.

Layers & heads — scale any attention branch, MLP branch, or individual attention head between 0 and 2 at inference time. Setting a head to 0 ablates it. Nothing is retrained; the gain multiplies activations inside the forward pass.

Attention — the real attention matrices from the last run, one square per head, rows are queries and columns are keys.

Next token — the model's raw next-token distribution for any text you type, with its entropy. Unfiltered by top-k/top-p, because the point is to see what the model actually believes.

measure loss — re-scores held-out data with the weights as they currently stand, so any edit has a number attached to it.

It does something visible

Ablating all four attention heads in layer 0 and asking the same question twice:

	reply	held-out loss	perplexity
checkpoint	I'm good, thanks. What about you?	0.461	1.59
layer 0 heads zeroed	I like talking to go for me a joke	2.190	8.93

The grammar survives and the meaning does not, which is roughly what you would expect from removing the layer that does the early token mixing.

Architecture

Decoder-only transformer, pre-LayerNorm, GELU MLP, learned positional embeddings, and the output head tied to the token embedding.

layers	4
heads	4 (head dim 32)
model width	128
feed-forward width	512
context	64 tokens
vocabulary	328 word-level tokens
parameters	843,520

Training uses AdamW with decoupled weight decay (applied to matrices only, not to LayerNorm gains or biases), a cosine schedule with 150 warmup steps, and gradient clipping at norm 1.0. Validation loss bottoms out around step 2,200–2,500 and rises after that, so training stops there and the best-validation checkpoint is the one kept.

Layout

microlm/
  ops.py          forward + backward for every primitive
  model.py        transformer assembly, forward, backward, save/load
  tokenizer.py    word-level tokenizer
  corpus.py       jsonl -> token stream, batching
  train.py        AdamW, schedule, checkpointing
  sample.py       temperature / top-k / top-p / repetition penalty
data/
  build_corpus.py generates the chit-chat corpus
server/
  app.py          stdlib HTTP server + JSON API
  static/         the control surface
tests/
  test_grad.py    numerical gradient check

On the gradient check

Central differences carry two competing errors: O(eps²) truncation and O(machine_eps / eps) roundoff. The step that balances them is different for every parameter, so each sampled entry is measured at three step sizes and scored on its best one. A correct gradient matches closely at some step; an incorrect one matches at none. Every tensor lands at or below 2.4e-6 relative error, most around 1e-8.

License

MIT