An invertible network trains live and bends a Gaussian and its grid like taffy.
A normalizing flow is a generative model made of invertible steps, so it can both draw samples and score any point exactly. This one stacks eight RealNVP-style affine coupling layers: each turns the plane by a fixed angle, keeps one coordinate a and updates the other as b * exp(s(a)) + t(a), where s and t come from a small MLP, so the log-determinant is just s(a). It trains in your browser by maximum likelihood with hand-written backprop through all eight layers and Adam. Because each layer's network sees a single number, s(a) and t(a) are tabulated every few steps, which lets thousands of grid points, a checkerboard painted on the Gaussian and every background pixel be pushed through the flow (or pulled back to measure log density) for a couple of lookups each. The scrub bar chooses how many layers have been applied, and the contours show the density at that stage, so you can see the Gaussian sheared step by step into moons, a pinwheel or a spiral, and see where it struggles, like tearing one blob into two nested rings.
Try it. Paint a new target density with the brush (strokes made within a few seconds of each other build one painting) and the flow retrains on it. Drag the layer bar or press the left and right arrows to apply the layers one at a time, pick Moons, Pinwheel, Rings or Spiral (keys 1 to 4), Space pauses the automatic sweep, G hides the grid and R resets the weights.
Paste this into Claude Code, Codex or any coding agent to get a simple version running, then take it wherever you like.
Build a live 2D normalizing flow demo with JavaScript and the HTML canvas element. Put everything in a single index.html file with no libraries or build step, so I can open it directly in a browser.
Start simple:
- Make a dark canvas that fills the window and stays sharp on high-DPI screens.
- The target data is the classic 'two moons' shape: about 2,000 points on two interleaved half circles with a little Gaussian noise, scaled to roughly [-2, 2].
- Build a flow from 6 to 8 affine coupling layers. Each layer rotates the plane by a fixed angle (use the golden angle times the layer index), keeps the first coordinate a, and maps the second as b' = b * exp(s(a)) + t(a), then rotates back. s and t come from a tiny MLP of a alone (1 input, two hidden layers of 24 tanh units, 2 outputs), with s squashed by 2 * tanh so it stays bounded. Start the last layer's weights at zero so the flow begins as the identity.
- Train by maximum likelihood: run a batch of 64 data points backwards through the layers to z, and minimize 0.5 * |z|^2 plus the sum of the s values. Write the backpropagation through the layers and the Adam optimizer by hand with Float32Arrays, and take as many steps per frame as fit in about 6 ms.
- Draw 1,000 points sampled from a standard Gaussian pushed forwards through the flow, so you watch them become moons as training proceeds.
Once that works, make it beautiful:
- Store each point's position after every layer and add a slider that picks how many layers are applied, interpolating smoothly between them.
- Draw a square grid in Gaussian space pushed through the same layers, so you can see the plane being bent.
- Shade the background by the model's log density (pull each pixel back through the flow) with a few contour lines.
Explain the key ideas in short code comments. When you're done, tell me how to open it and suggest three directions I could take it next, such as painting a custom target density with a brush, a checkerboard texture that shows the flow stretching like taffy, or comparing to targets a flow cannot easily represent, like two nested rings.
