Jane Street published a puzzle: a trained 48-block neural network (97 layers) fell apart into 97 shuffled layers, and you get the pieces plus 10,000 rows of its historical inputs and outputs. Put it back together in the right order and the loss goes to zero.
Hyunwoo Park solved it first. The approach is satisfyingly physical. Pair the layers: each block was split into two, and a simple check on the weights says which two belong together, so all 48 blocks reassemble without touching the data. Then guess an order by sorting blocks by the size of their relative weights — the hunch being that later layers tend to have larger weights. Then search: try swapping neighbours, keep a swap only if the error drops, repeat.
The pairing and the sort take well under a second. The search takes three or four seconds, about 95% of the solve time. Every candidate swap needs a full forward pass over the whole network to see whether it helped.
Cezar Cocu looked at that bottleneck and remembered something from nuclear engineering.
The same question, asked in 1945
Reactor designers had an analogous problem. If you tweak one material or move one control rod, what happens to the final result? You could re-run the simulation for every possible change, but in the physical world that's prohibitively expensive.
The trick, documented as early as 1945, is to solve the system once forward and once backward. The backward solve gives you the importance of every point in the reactor — how much the final answer cares about a change there. Multiply a local change by its importance, and you have an estimate of its effect on the whole system, without re-running anything.
That's the adjoint state method, and Oak Ridge's SCALE tools still work on the same principle.
What it buys on a neural network
Cocu's version asks the same question about a layer swap: can we predict its effect on the total error more cheaply than running the full forward pass again?
A full forward pass over the network is 96 GEMMs — general matrix multiplies, one per layer transition. An adjoint prediction is four GEMMs plus a dot product. The catch is that the adjoint requires one full forward and one backward pass to set up, but that cost is amortized across all the candidate swaps you evaluate afterward.
The result: Park's solver took about 4.0 seconds; with the adjoint trick it takes 0.19 seconds on a CPU-only i9-13900F. The whole script went from 4.97 to 1.28 seconds, which means the solve is now faster than the Python startup around it. Twenty times faster on the part that mattered.
Why this transfers
The two problems look unrelated — a reactor and a shuffled stack of matrices — but they share a structure. In both, you have a system with an output you care about, and you want to know the sensitivity of that output to a local change, for many candidate changes, without re-evaluating the system each time.
That's the general shape of the adjoint method, and it's why it shows up everywhere: weather models, structural engineering, optimal control, and now neural architecture search. The expensive thing is the forward solve. The cheap thing is the backward solve that tells you where the forward solve is sensitive.
The mental move is to stop asking "what is the result if I change this?" and start asking "how much does the result care about this?" Those are duals. The second question is often much cheaper, and it's the one that generalizes across candidates.
The charming detail
The dataset includes two output columns: pred, what the correctly ordered model predicted before it fell apart, and true, what actually happened. The solver only ever needs pred, because the task is reconstruction, not prediction. Cocu notes he's confused about why true was included, and I like that he left the question open rather than inventing an explanation.
It's a small thing, but it's the kind of detail that tells you someone actually did the work rather than narrating it.
The lesson beyond the puzzle
Most expensive searches have this shape: a candidate generator plus a full evaluation per candidate. The naive fix is to make the evaluation faster. The better fix is often to compute the derivative of the evaluation and use it to rank candidates, so most candidates never get evaluated at all.
Park's solver was already clever — the weight-sorting heuristic got it close, and local search finished the job. The adjoint didn't replace the search; it made each step cheap enough that the search stopped being the bottleneck. That's the usual division of labour in optimisation: a good heuristic to get near the answer, a good gradient to converge, and an awareness of which one you're paying for.
The fact that the gradient was invented for reactors in 1945 and applied to a neural network puzzle in 2026 is the part worth remembering. The mathematics doesn't care what the system is made of.
Source: Stealing a Trick from Nuclear Physics to Solve a Puzzle — Cezar Cocu.