David Nichols

Equilibrium Reasoners

Elsewhere I have published playable Marupeke puzzles. On this page you can inspect the intermediate guesses of an iterative neural network solver as it tackles such puzzles. The solver is a 1.3-million parameter MLP-mixer model trained using the framework of the Equilibrium Reasoners (Huang et al., ICML 2026) preprint to solve the puzzles by looping its calculations until they converge to a correct solution.

loop 0loading model…

"Play" will start the solver looping automatically while "Step" will advance the solver a single loop. "Reroll start" will start the solver over from a new random latent. The solver marks both O and X on each square with opacity based on its confidence, so as the solution evolves the guesses the model discards fade away (and possibly fade back in as it reevaluates). The phase portrait is a 2D projection of the evolution of the model's 16,384-dimensional output-side latent onto the first two principal components of the model's solution trajectories.

The paper starts with the following idea. The latent state of the model decodes to a candidate solution to the problem—in my case, a Marupeke puzzle—and we would like successive loops of the model to update its latent state so that the corresponding candidate solution gets closer and closer to the true solution in what the paper calls task-metric space. This suggests a way to think about how to train the model: we want to tune the fixed-point attractor landscape of the latent space so that its fixed-point attractors decode to correct solutions in task-metric space. In this way we are aligning the model's latent space with the task-metric space, which the paper illustrates with a figure adapted below.

Figure 1 from Huang et al. 2026.

In reality, we need to relax our objective: instead of reaching a true fixed point, the model should enter attractor regions in latent space with the property that latent states anywhere in the region decode to the correct solution. In order to converge to those attractors, the paper suggests two scaling axes: depth (more loops) and breadth (aggregating the results of multiple attempts at finding the solution). It turns out that depth scaling alone gets the model to correctly solve most puzzles, so that breadth scaling is only required to rescue the relatively small remainder.

On this page, I implement breadth scaling interactively: one solution attempt plays at a time, and when the solver gets stuck you can click the "Reroll start" button to start the solver on a new random latent state. Because I implement breadth scaling in this way, and because an automatic referee checks the model's attempted solutions against the rules, I deviate from the paper in another way: the solver stops as soon as it reaches the correct solution.

The particular way in which an "equilibrium reasoner" implements looping is to update its latent state z according to the following rule:

zk+1 = zk + λ(f(zk; x) - zk) + βε,

Where x is the puzzle, λ is the step fraction, ε is unit Gaussian noise, and β is the noise injection strength parameter scaling that noise. Huang et al. found that the solver performed best on their tasks (Sudoku and maze pathfinding) with λ = 0.95; my Marupeke solver was trained with λ = 0.95 but runs on this page with λ = 0.25 for better visualization of the model's intermediate guesses. Noise injection helps shake the model out of spurious attractors and find the solution attractor.

Training the model as an equilibrium reasoner involves introducing three special ingredients to the training recipe, which the paper's ablation study shows all matter to the model's performance:

  1. Random initialization. Unlike the models in the papers this one builds on, equilibrium reasoners start iterating from a freshly randomized latent. This makes the model's behavior more stable by training it to converge to the same attractor along multiple paths and enables the model to explore more of the latent space during training.
  2. Noise injection. By injecting the same kind of noise during training that we use during inference, the model gets the same kind of help escaping shallow basins that it gets at inference time. This complements random initialization, since some random starting points might result in the model otherwise getting stuck.
  3. Adaptive computation time. A learned halting mechanism that keeps the model from continuing to loop after solving an easy puzzle helps reduce overfitting and thereby improves model accuracy. Because this benefit is locked in at training time, we can ignore the halting mechanism at inference time and still benefit from having included it.

Training with the above recipe and looping with the above update rule result in a model with a well-tuned attractor landscape. And that means that even though the model is only trained to a depth of 16, it continues to benefit from greater depth scaling at inference time. On a test set band of harder Marupeke puzzles, the model as trained solved 29% of puzzles in a single trajectory at depth 1, 82% at depth 16, and 90.0% at depth 256.