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 of the momentum residual block: it 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 Hamiltonian-Step 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 Built from Hamiltonian Steps
A pendulum organizes its motion around energy, while an ordinary residual network has no comparable scalar. This post builds a residual block as a symplectic step of a learned Hamiltonian. The continuous field conserves that Hamiltonian exactly; the discrete network keeps it in the bounded oscillatory band predicted by symplectic integration. The learned pendulum stays within 0.6% where a plain field model drifts 36%, and the classifier survives being run at four times its training depth.