Train a tiny transformer

This is a GPT in miniature, trained while you watch. It reads text one character at a time and learns to predict the next one. Nothing is pre-trained and nothing is downloaded: the model starts from random numbers in your browser, and within a minute or two it goes from noise to spaces, then common words, then line breaks in roughly the right places. It never gets good — the biggest model here has about a hundred thousand parameters, a millionth of a small modern LLM — but every part of a real one is there and you can look at it.

How to use it

Pick a text and a model size. Tiny (one layer, d_model 32, about 15 k parameters) learns word shapes within a minute; Small (two layers, d_model 48, about 60 k) shows real words and line breaks in a couple of minutes; Medium (d_model 64, context 64, an untied output head) needs five minutes or more. Or set the layers, d_model, heads, context, batch, learning rate and seed yourself. Changing the model starts again from fresh random weights; the learning rate can change while it trains.
Press ▶ Train. The loss curve shows the training loss (cross-entropy in nats per character) and, every 50 steps, the loss on the last 10 % of the text, which the model never trains on. Two dashed lines give it scale: ln V is a uniform guess over the V characters; letter frequencies is the best you can do knowing only how common each character is. Anything below that line is the model using context. When the held-out loss stops falling while the training loss keeps going, it has started to memorise.
Samples are generated every 100 steps (change it, or press Sample). Temperature divides the logits before the softmax — 0 always picks the most likely character, 1 samples the model as it is, higher is wilder — and top-k keeps only the k most likely characters. A prompt is fed in first; the model continues it.
Attention shows, for a line of text, how each head of each layer splits each position’s attention over that position and the earlier ones (later ones are masked: that is what makes it a decoder). Hover or tap a cell for the numbers. Untrained heads spread evenly; after a few hundred steps you can usually find one that looks at the previous character and one that looks back to the start of the word.
Character embeddings plots each character’s learned vector projected on its first two principal components (power iteration on the covariance). Vowels tend to drift together, away from consonants, with space and punctuation off on their own.
Parameters and FLOPs counts the weights by group and estimates the arithmetic per training step.
Download weights saves the model as JSON (with its vocabulary); Save here keeps one copy in this browser; Load saved / Open file… bring a model back so you can sample from it or train it further.

The share link carries the text choice and the model settings — not the weights, and never text you pasted.

How it works

The model is the GPT-2 shape, shrunk: a token embedding plus a learned position embedding, then pre-LayerNorm blocks (multi-head causal self-attention, then a GELU MLP four times as wide, each added back to the residual stream), a final layer norm, and an output layer that is, by default, the token embedding matrix transposed (“tied”). The loss is the mean cross-entropy of the next character. The optimiser is AdamW (β₁ 0.9, β₂ 0.99, weight decay 0.01 on the matrices only) with a short warm-up and gradient clipping at norm 1.

There is no autograd library: every operation — matrix multiply, layer norm, GELU, the causal softmax attention, cross-entropy — has its backward pass written out by hand in tiny-transformer-core.js, and the tests check each one, and the whole model, against central finite differences in float64. If you want to see how backpropagation works on a small graph first, the maths of machine learning tool steps through one; the neural network designer builds bigger networks and exports PyTorch.

Training runs in a Web Worker so the page stays responsive; if a worker can’t start it runs on the page in short slices. Everything is plain JavaScript on the CPU over typed arrays — no WebGPU: at this size the time goes into many small matrix multiplies, and moving them to the GPU would add a lot of code for little gain. On a recent laptop the Small model does about ten training steps a second (some 2–3 GFLOP/s).

Characters are the tokens here, so the vocabulary is just the 50–60 distinct characters in the text. Real models use subword tokens — see the tokenizer explorer. For what attention costs at real sizes and how inference caches keys and values, see attention and the KV cache; for optimisers compared on loss surfaces, the classic machine learning tool’s optimiser lab.

Texts

All the built-in texts are in the public domain:

Shakespeare’s sonnets — nineteen of the 154 sonnets (1609), about 12 000 characters.
Nursery rhymes — traditional English rhymes, with Jane Taylor’s “The Star” (1806) and Sarah Josepha Hale’s “Mary’s Lamb” (1830), about 8 000 characters.
Carroll and Lear — “Jabberwocky”, “How Doth the Little Crocodile”, “Twinkle, Twinkle, Little Bat”, “You Are Old, Father William”, “The Walrus and the Carpenter” (Lewis Carroll, 1865 and 1871) and “The Owl and the Pussy-cat” (Edward Lear, 1871), about 6 000 characters.

Your own text takes anything you paste: a few thousand characters is enough to see learning, and more is better. Code, recipes and song lyrics all work.

Browser support

Any current browser; Web Workers are everywhere. Training speed depends on the CPU: a phone is a few times slower than a laptop, so choose Tiny there.

Privacy

Everything runs in your browser. The texts come from this site; text you paste and the weights you save stay in this browser’s storage, and nothing is sent anywhere.