Main Page

Tiny LM

The claim

Tiny LM runs in a web browser with no machine learning library. On five fixed test inputs, its next-token scores match PyTorch to within 1.2e-5.

Why I built it

The models behind chat assistants have billions of parameters and run on someone else’s servers. You can use one every day without seeing how it chooses the next word.

I built every part of Tiny LM from scratch, from cleaning the training text to running the model in your browser. Its page lets anyone watch each step, from cutting text into tokens to choosing the next word.

Tiny LM has 3,737,088 parameters. Its text sounds like its training books and then wanders. It is for learning and understanding how language models work, and no use as a writing tool.

What I built

I trained a four-layer GPT-style model on seven public-domain books from Project Gutenberg, written between 1758 and 1910. Training took 426 seconds on an Apple M4 Pro GPU.

The page downloads the model once and runs it in under 400 lines of TypeScript.

Tiny LM's attention view for the sentence 'The man who saves a little every week will have money when he needs it.' A row of token chips shows what the token 'money' looks at in layer 1, head 1: 21% on 'have', 19% on 'will', and between 3% and 8% on each of the other tokens, itself included. Below it, a triangular heat map has one row for each token and one column for each earlier token, with darker blue for a larger share of attention. Tiny LM's attention view for the sentence 'The man who saves a little every week will have money when he needs it.' A row of token chips shows what the token 'money' looks at in layer 1, head 1: 21% on 'have', 19% on 'will', and between 3% and 8% on each of the other tokens, itself included. Below it, a triangular heat map has one row for each token and one column for each earlier token, with brighter blue for a larger share of attention.
Layer 1, head 1 of the attention view. “money” gives 21% of its attention to “have”, more than to any other word.

AI coding agents (Claude Code) wrote the code under my direction.

Decisions

I trained the model myself. Fine-tuning a small GPT-2 would write far better text, but it would hide the training and make the download much larger.

I wrote the forward pass without a machine learning library. ONNX Runtime Web would run faster and add megabytes of someone else’s code. A full pass over 256 tokens takes about 1 second as a result.

The page runs the weights from the last training step. Keeping the checkpoint with the best validation score would break the training replay, whose final sample has to come from the model on the page.

How I tested it

The parity test feeds the same five inputs to the browser code and to PyTorch, both loading the same weights file. The largest difference in the next-token scores is 1.144e-5. The test fails above 1e-4.

The browser tokeniser matches the Python one on all 38 test cases. On a local build, the first generated token appears 19.9 ms after a click.

What broke

The first training run memorised the books. Validation loss rose after step 2,500 while training loss kept falling. I doubled dropout to 0.2 and cut the run from 6,000 steps to 4,500, and the final validation loss fell from 4.0430 to 3.8146.

An export test could not fail. It reloaded a freshly initialised model, where swapping two tensors changes nothing. It now fills every weight with random values, and it caught all three exporters I broke on purpose.

Text cleaning turned “1/2” into “12”. The cleaner now maps the fraction slash to “/” before it drops other characters.

What next

The model still memorises. Its validation loss is 3.8146, against 3.0932 in training. More public-domain text would help most.