Two models, trained from scratch, shipped to a browser tab.

Why this exists

The project started as an exercise in building a language model from first principles: attention, then a small GPT, then a model that has to reason rather than recite. Family relationships are a good test because the answer is checkable, chains can be made longer to stress the model, and a wrong answer is easy to spot.

Architecture

The relationship model reads the question and story with names replaced by entity ids (@0 to @31), writes a chain of family facts, then the answer. The name model reads a text twice and labels each word of the second copy as a name or not.

Dataset design

The relationship model learned from CLUTRR, a benchmark of short family stories that need several reasoning steps. Evaluation uses itsgen_train234_test2to10 configuration: validation stories need 2 to 4 steps, test stories 2 to 10. Training used 64,554 stories from the other splits, after removing any story that matched a held-out story up to the names.

Three kinds of data were added on top, because the original benchmark always has an answer:

• Unanswerable questions: a friend or stranger, two families that never meet, or a person who is not in the story. Target: “unknown”.
• Distractors: unrelated sentences mixed into answerable stories.
• Wrong chains: paths with one invented, swapped or near-miss step, labelled with a <nofact> token so the model learns to check each step against the story.

Training

• PyTorch on an 8 GB Apple M2 laptop (Metal backend).
• AdamW (β 0.9, 0.95; weight decay 0.1), learning rate 5×10^-4 with 300 warm-up steps, batch 32, dropout 0.1.
• Loss on the chain and the answer, plus a small next-word loss on the story text (weight 0.1).
• 30% of questions were sub-questions taken from the dataset’s proofs, so short hops are practised often.
• The final version adds hard-negative mining: every 200 steps, the model’s own beam search runs on 512 unanswerable questions, and every convincing wrong chain it finds becomes a training example for the step checker.

Evaluation

Accuracy is measured on held-out stories the model never saw, with a reference search setting (beam 8, four person-id orderings voting) and the lighter setting this site uses (beam 4, one ordering). On 950 labelled questions the light setting matched the reference on average at about one sixth of the compute.

v8 kept v4’s accuracy while learning to say “unknown”: it now refuses 92% of questions about a person who is not in the story. An end-to-end check on plain text with invented names (name model, rules and relationship model together) gave 90% correct on validation stories, with only 2% confidently wrong.

Error analysis

• Long chains. Accuracy falls from 99% at two steps to roughly a third beyond six. A 13M-parameter model has little room to hold long chains of facts.
• Friends and strangers. The hardest unanswerable case: friendship sentences look like family sentences. The model alone catches 46% (reference setting); a rule that checks whether two people are ever mentioned together brings two-family stories to 96.5%.
• Distractors. Unrelated sentences cost about 11 points on validation stories (78.2% vs 89.8%).

Confident answers are much more reliable than tentative ones; the limits page shows accuracy by chain length.

The name model

A causal model cannot look ahead, but deciding whether “May” is a name needs the words after it. The fix is to feed the text twice and read the labels from the second copy, where every word can attend to the whole first copy. Trained on 64,554 stories with a pool of 630 names, plus places, companies, languages, months and holidays, and with short snippets and lower-case text so it does not depend on position or capitals:

• F1 0.998 on names it never saw in training.
• Real places or companies tagged as people: 7.3% (6.7% in short texts, down from 52% before short-text training).

Running in the browser

A 460-line C engine implements the forward pass (layer norm, attention with a shared key/value cache, GELU, matrix multiplies) and is compiled to WebAssembly with Emscripten into six builds: relaxed SIMD, SIMD and plain scalar, each with and without threads. Each build is 28 to 46 KB. At load, the page tests which builds the browser accepts and tunes matrix-multiply tiles and the thread count for the CPU.

Weights stay at their original 32-bit precision. To shrink the download losslessly, the four bytes of every float are split into separate planes before gzip, which compresses better: the relationship model goes from 52.0 MB to 44.2 MB. The browser streams the parts straight into WebAssembly memory as they arrive. Results match PyTorch to within 1.5×10^-5 and pick the same answer and chain on all 60 reference questions.

Why WebAssembly, not WebGPU

A GPU sounds like the obvious choice for a neural network, so it was measured rather than assumed. The benchmark (tools/bench-webgpu.html) runs this model’s real workload both ways on an Apple M2: the WebAssembly engine as shipped, and a WebGPU compute kernel doing the same matrix multiplications with the same weight shapes, reading results back after every step as the search requires.

Even counting only its matrix work, the GPU is three to nine times slower per step. Beam search is a chain of small decisions: each step processes 1 to 12 rows, then the CPU must read the results to choose the next step. At that size the cost of sending work to the GPU and reading it back outweighs the GPU’s raw speed, and uploading the weights and compiling shaders adds almost a second at start-up. WebGPU would pay off for much larger batches or a much larger model; for this one, a tuned CPU engine that adapts to the device (relaxed SIMD, SIMD or plain WebAssembly, with or without threads, tiles and thread count tuned on load) is faster everywhere and works on far more devices.

Explaining answers step by step

The map turns the model’s chain of facts into a worked solution: a family tree drawn by generation, built in reasoning order, with each step combined with the last using fixed kinship rules (“father’s wife is mother”). Those rules are code, not model output, and are unit-tested. Checked against the model on 299 held-out validation questions, they produce a step-by-step deduction for 277 and agree with the model’s answer on 265. When they disagree, the page shows the warning rather than the comfortable story. The layout is tested on 735 chain shapes at phone and desktop widths for overlapping labels.

Worker architecture

The page never runs the model itself. A module Web Worker loads the engine and both models once, then answers messages: find the people in a text, or answer a question. Every request carries an id so a late answer for an edited story is ignored. Threaded builds start their own helper workers, which needs a cross-origin isolated page (Cross-Origin-Opener-Policy and Cross-Origin-Embedder-Policy headers); without them the single-threaded build runs instead.

Caching and offline use

• Model files live under a content-hashed path, so a new model never mixes with old runtime code and files can be cached forever.
• Weights are split into parts of at most 16 MiB (61.4 MB in total), downloaded three at a time, checked against SHA-256, and stored in the browser’s Cache Storage.
• A damaged file is detected (checksum, gzip check or the CRC-32 of the decoded weights), discarded, and downloaded again. Only this site’s model files are touched.
• A service worker keeps the pages and engine available offline. Old model versions are deleted when a new one loads.

Latency and memory

WebAssembly memory is about 105 MB single-threaded and 201 MB with threads. Older devices fall back to the slower builds automatically; the answer is identical, only the time changes.

Load the model on the workbench and come back to see the numbers measured on this device.

Challenges and lessons

• Teaching “I don’t know” directly failed. Training the model to answer “unknown” at the answer position did not learn. Checking each step instead worked.
• Models find shortcuts. One version learned that “extra people” meant a made-up chain. Adding wrong chains to clean stories removed the shortcut.
• Search fools a verifier. Beam search is good at finding chains that slip past the step checker. Mining those chains as training data fixed it.
• Small machines need care. Mining 512 questions at once used 7.7 GB on an 8 GB laptop; batches of 64 kept training stable.
• Position matters. The name model failed on short texts it had never seen at that length; training on short snippets cut that error from 52% to 6.7%.

Source code: github.com/ArAnirudh2901/LLM-From-Scratch