Deep Equilibrium
4 long-form posts on Deep Equilibrium: machine-learning research by Taha Bouhsine, each built around live, in-browser interactive visualizations.
-
Editing a Deep Equilibrium Network, in JAX/Flax NNX
A runnable companion: build the weight-tied Yat equilibrium operator in Flax NNX, then teach a class by appending rows to the readout (F untouched, exact) or into the shared dynamics (one paste, present at every depth), measure the contraction certificate with power iteration and bisect one gain to restore it, audit the drift of 520 old fixed points, watch a layer-only edit evaporate, and forget by masking. Every number is from a real run.
-
Edit One Operator, Edit Every Depth
One post taught and forgot classes by editing rows of a Yat network, with proofs that nothing else moved. Another melted the stack of layers into a single operator iterated to a fixed point. This is the collision. Every one of those editing proofs rested on a pasted row entering the score once, as one term in one sum, and in an equilibrium network there is no once: whatever you paste is applied at every depth and fed back into its own input, and every fixed point is free to drift. So did melting the stack melt the editability? This post pastes, deletes, and measures: every guarantee that survives is either proved inside the recursion or measured against the real run, fixed point by fixed point.
-
A Network That Is a Fixed Point, in JAX/Flax NNX
A runnable companion: build the Yat deep-equilibrium network in JAX/Flax NNX. One shared operator F(z;x)=tanh(A·φ_W(z)+Ux+z0), solved to its fixed point by damped iteration, trained not by backprop-through-iterations but by implicit differentiation, the adjoint (I−Jᵀ)u=∂L/∂z* run by the same contraction. Plus the weight-tied maze operator that extrapolates from 11×11 to 27×27 by iterating longer. Every number is from a real run: 98.2% on two moons from 1700 shared parameters, ‖J‖ 0.66/0.92, 99.5% on mazes far larger than training.
-
Your Network Is a Stack of Layers. It Could Be a Fixed Point.
A deep network makes you choose its depth before you have seen the problem, and gives every layer its own weights. Share one Yat-kernel operator across all of them instead, and the stack collapses into a single equation: the answer is the fixed point the state settles into. Training makes that operator a contraction, so the settling point is unique and reached from anywhere, the network decides its own depth per input, and the same twenty-four prototypes describe the computation at every step. 98.2% on two moons from 1700 parameters shared across all depth.