Integrators
6 long-form posts on Integrators: machine-learning research by Taha Bouhsine, each built around live, in-browser interactive visualizations.
-
An Error Controller for a Trained Net, in JAX
A runnable companion to the depth-on-demand post: the leapfrog classifier trained with lax.scan at fixed depth, the step-doubling controller that re-renders it to tolerance at inference, the honest work accounting (probes included), and the measured tol^(-1/3) power law. Every figure is rendered from the real Kaggle run.
-
Depth on Demand
The last post made depth a resolution: layers are time steps of a learned flow, and running more of them just renders the same trajectory finer. But every camera knows not to spend equal film on empty sky. This post gives a trained network the integrator's next tool, an error controller that chooses its own step size per input, with no retraining: the same weights, rendered to tolerance. The controller reproduces the reference verdicts at a fraction of the steps, its cost follows the integrator's textbook one-third power law, and the map of where it spends is a genuine surprise: effort tracks the stiffness of the learned flow, not the difficulty of the classification.
-
Reversible Backprop as a custom_vjp in JAX
A runnable companion to the memory post: the momentum block and its exact inverse, the custom_vjp whose backward pass reconstructs the trajectory instead of storing it, XLA's memory_analysis as the measuring instrument, and the (1/mu)^L float cliff reproduced in numpy float32. Every number is from the real Kaggle run.
-
Backprop Without the Memory
Training memory is a tax nobody chose: backprop must hold every activation of the forward pass hostage until the backward pass consumes it, so depth costs memory even when it costs little compute. This post spends the invertibility this series proved two posts ago: a momentum residual block can be run backward, so the backward pass can recompute the past instead of storing it. Measured on the same network, standard backprop's activation memory grows from 13 MB to 674 MB as depth goes 8 to 512; the reversible pass holds flat at 3.2 MB, pays 24% in step time, and returns the same gradient, until a one-line arithmetic of friction and float noise says it cannot.
-
Building the Energy-Conserving Net in JAX/Flax NNX
A runnable companion to the energy-conservation post: the HNN pendulum field as the symplectic gradient of one learned scalar, the plain field model it beats on drift, and the leapfrog classifier whose residual block is a kick-drift-kick step of a learned potential, all as Flax NNX modules with lax.scan doing depth. Every figure is rendered from the real Kaggle run.
-
A Network That Conserves Energy
A pendulum never forgets its energy, and a trained network has no such number to forget. This post builds a residual network whose hidden state carries a conservation law by construction: the block is a symplectic step of a learned energy, so the quantity is held by the architecture, not encouraged by a loss. The learned pendulum keeps its energy to 0.6% where a plain field model leaks 36%, the classifier lands in the pack on accuracy, and the law pays where composition fails: trained at depth 16 and run at four times that, the plain net gives up 31 points on spirals while the leapfrog net holds.