jaxverse
Contents

jaxverse

A little universe of learning machines, designed and built by Neo Mohsenvand.

GitHub LinkedIn

An interactive book, computed as you read

A little universe of learning machines

Every model in this book is real, and every one of them trains here, in your browser — from a single neuron bending a line to a language model learning chess.

Plate I The loss landscape, live

gradient descent with momentum Adam

Three optimizers descending the same surface, right now, on your machine. Click anywhere on the map to drop them somewhere new and watch the routes diverge — the whole of the Prologue is in this picture.
Begin the descent View the source

Machine learning has a reputation for mystery it doesn’t deserve. Under every headline model — the ones that talk, draw, and play — sits one modest idea: build a smooth, adjustable function, measure how wrong it is, and nudge every knob a little downhill. Do this a few million times. That’s the whole trick. Everything else is architecture and appetite.

This book makes that idea physical. It is arranged as a descent: a prologue, then chapter by chapter, each one lower on the loss surface than the last. You’ll fit curves, bend space until tangled spirals come apart, teach a network to read handwriting, watch a map of meaning assemble itself without a single label, train a small language model three different ways — the same pipeline, in miniature, that produced the models everyone is talking about — and end by pulling pictures out of pure noise.

Nothing here is a recording. The numbers move because a network is training on your GPU as you watch, courtesy of jax-js. Pause anywhere. Break things. The models are small and do not mind.

A second book, for the builders

Beneath every plate that trains something sits a quiet line like the one above — and the one above works: open it for the machinery behind the map at the top of this page. Every block holds the real jax-js code from this repository, annotated, with the stagecraft on its own tab. Followed chapter by chapter, the blocks add up to a short course in jax-js itself: pytrees, valueAndGrad, jit, devices, workers, a whole transformer. Each chapter also ends with a downloadable lab — a standalone npm project that runs the same model with nothing but npm install.

Contents — the descent

  1. 0 The Descent Optimization Learning is falling downhill on a landscape of error. Meet the loss surface, the gradient, and the little step rule the rest of this book repeats. 10′
  2. 1 The Approximator Neural networks A neuron is a bump of influence; a layer is a sum of bumps. Watch a tiny network sculpt itself into any curve you draw. 12′
  3. 2 Bending Space Representation Why does deep learning work at all? Because a network is a smooth deformation of space — watch it untangle spirals until a straight line can tell them apart. 13′
  4. 3 Telling Things Apart Supervised learning Ten thousand handwritten digits, a stack of layers, and a rule for being less wrong. Train a classifier and then look inside it. 14′
  5. 4 The Hidden Map Representation learning Squeeze every digit through two numbers and back. No labels, no answers — yet a map of meaning appears on its own. 12′
  6. 5 The Next Token Self-supervised learning Predict what comes next: the simplest game in the world, and the one behind every large language model. Train one, live, on three hundred thousand words of children’s stories. 15′
  7. 6 Learning from Reward Reinforcement learning No examples, no answers — only consequences. Watch a policy discover a path through a world that only ever says “more” or “less”. 12′
  8. 7 Teaching Taste Preference learning Some things can be checked. Most can’t. Fit a judge to nothing but which-of-these-two, watch an optimizer take it apart — then meet the one idea that stops it. 16′
  9. 8 Rook Everything at once One small language model learns chess three ways: pretrained on games, fine-tuned on better ones, then sharpened by verifiable reward. The modern pipeline, in miniature. 17′
  10. 9 Out of the Static Generative modelling Ruining a picture is easy, and you can do it a grain at a time. Learn to take one grain back and you have a machine that draws things nobody has ever drawn. 15′
  11. 10 The Straight Path Flow matching Aim the model at the shortest route from noise to picture, then learn to steer it. Ask for a boot, or walk the label from sneaker to boot, and watch it draw a garment it was never shown. 16′
  12. ∞ Epilogue Where the ideas came from, and where to go next.