Stealing a Trick from Nuclear Physics to Solve a Puzzle
Jane Street dropped a neural net puzzle.
This is a bit before my time but I am assuming that this is a reference to the terror of dropping punch cards on the floor and being unable to put them back together.
Oops! I am not sure how you'd put the punch card in the correct order if they were not labeled but for a neural network, we might have some hints in the weights themselves. Hm....
For some context, in the puzzle, a trained 48-block network fell apart into 97 shuffled layers. We get the pieces plus 10k rows of its historical inputs and outputs. Our job is to glue it back together in exactly the right order. It is implied that the correct order would result in a loss of 0 across any of its inputs.
Hyunwoo Park solved it first, and spoiler alert: my solution builds on it but makes it significantly faster. Because why not? This is a very verifiable task so LLMs can really assist us here.
TL;DR
The intuition here is: Don't test every swap. Instead run the neural network both forward and backwards once, and using the adjoint method, instead predict (with some error bars) what every swap would do to the error. This only works if the prediction is cheaper to perform vs actually running the forward pass for the current block combination.
For a sense of scale: a full forward pass over the neural network is 96 GEMMs (GEneral Matrix Multiply). On the other hand, one adjoint prediction is just 4 GEMMs plus a dot product. For the adjoin prediction, we do still pay for one forward and one backward pass per round, but that gets ammortized across all of the predictions.
Park's strategy
His solver effectively has three steps:
- Pair. Each block was split into two layers. A simple check on the weights tells which two belong together, so all 48 blocks come back w/o touching the data.
- Guess an order. Sort the blocks by the size of their weights. Later blocks tend to have bigger weights, so this gets close.
- Fix the order. Try swapping neighbors, only keep a swap if the error drops, repeat until nothing helps or error drops to 0.
Steps 1 and 2 take well under a second. Step 3 takes about 3-4 seconds (about 95% of the solve time).
My hunch: from my vague memory, physicists (nuclear ones especially) must have this problem too. They deal in the world of atoms and not bits and can't afford to re-run a whole experiment every time they want to know what a small change might do. So we (w/ Claude ofc) went looking. We found a promising strategy from their toolset, the adjoint method.
Reactor physicists figured this out in 1945
Reactor designers have the same problem. What happens if you tweak one material, or move one control rod? Re-running the full simulation for every possible tweak is way too expensive. In 1945 they were doing this math by hand and on punch-card machines (!), and even today, with thousands of things to tweak, it adds up fast.
Their trick goes back to Wigner in 1945 (review). Basically, 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 result cares about a small change there. You then multiply a small local change by its importance and you should know its effect on the whole reactor, w/o re-solving anything. Oak Ridge's SCALE tools still work this way today.
How much does it actually save?
About 20X faster than Park's original, on my Intel-based box (an i9-13900F, CPU only):
| Solve time | Whole script | |
|---|---|---|
| Park's original | 4.0 s | 4.97 s |
| w/ adjoint trick | 0.19 s | 1.28 s |
Some of that gap comes from other tweaks (like the starting guess), not just the adjoint.
Fun part: the solve is now so fast that importing PyTorch and loading the data take longer than solving. Ha!
A Head Scratcher
10k rows of its inputs and outputs are provided but in order to solve this we only ever need 1,000 of them (the other 9k just double-check the answer), and we never need the true column at all! I'm a little confused why the rest of the dataset was included.
Is there something hiding in there?!
Code
The solver, benchmark and an academic-style write-up are on GitHub. Park's original solver is here, and his write-up here.