jaxverse

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

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 world model learning the consequences of a move.

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.

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 learning loop. The chapters change what we ask that function to learn, and what we do with its answers.

This book makes that idea physical. It begins with a descent down a loss surface, then follows what that learning rule can build. 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 learn two routes from noise to pictures.

Finally, you’ll train a world model from pictures and actions. Give it a destination afterward, use it to rehearse possible movements, and test a prediction by acting. Change the goal and use the same learned dynamics again. Learning how a world behaves and deciding what to do in it become two distinct parts of one experiment.

The experiments are computed here, in your browser, with jax-js. Train, pause, change an input, and watch what follows. 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, and planning through a learned world model. 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. 11 Before the Move World models A little mechanism learns from pictures and actions. Give it a destination afterward, and watch it rehearse a movement before making it. 18′
  13. ∞ Epilogue Where the ideas came from, and where to go next.