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 in the right feed order.
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. Hmm....
For some helpful context, in case you didn't check out the puzzle details, a trained 48-block (97 layer) neural network fell apart into 97 shuffled layers. We get the 97 shuffled pieces plus 10k rows of its historical inputs and outputs. Our task 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 all of its labeled outputs.
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 help us hill climb here.
Park's strategy #
The original 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 relative weights (originally L2 ). The hunch here is that later layers tend to have bigger weights vs earlier ones.
- Search for the right order. Try swapping neighbors, but only keep a swap if the error drops, repeat until nothing helps or error drops to 0.
On my machine, steps 1 and 2 take well under a second (this includes script setup time). Step 3 takes about 3-4 seconds (about 95% of the solve time).
My hunch: from my vague memory, physicists or engineers (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...
TL;DR #
The intuition here is: Is it possible to cheaply predict what a potential layer swap's effect on the overall error instead of running the full forward pass again over the entire network? Even if we are able to predict this, it will only work if the prediction is computationally cheaper to perform vs actually re-running the forward pass for the current block combination.
Well it seems that there might be such a method and it's called the adjoint state method. And it is indeed cheaper to compute vs re-running inference on the full neural network.
For a sense of scale: a full forward pass over the neural network is 96 GEMMs (GEneral Matrix Multiply), since we have 97 layers to compute. On the other hand, one adjoint prediction is just 4 GEMMs plus a dot product. The one catch is that for the adjoint computation, we do still perform one full forward and one backward pass, but that gets amortized across all of the predictions.
This has been a solved problem in nuclear engineering for quite some time #
Reactor designers/engineers had a similar problem. For example, what happens if you tweak one material, or move one control rod? You can re-run your model or your experiment of course! But as it turns out in the world of atoms and not bits, re-running the full simulation for every possible tweak is way too expensive.
Back when this method was documented, they were doing this math by hand and if lucky on punch-card machines (perhaps this is what the title of this puzzle is alluding to?!).
The earliest I could find this strategy is from 1945 (review). Basically, solve the linear system once forward and another time backward. The backward solve gives you the importance of every point in the reactor (read: how much the final result cares about a change). You then multiply a potential small local change by its importance and you should know its effect on the whole reactor. Oak Ridge's SCALE tools still work on this basic concept.
Details on the math for this particular problem are here.
What is the practical speedup? #
On my Intel-based box (an i9-13900F, CPU only, w/ SIMD acceleration), its about 20X faster:
| Solve time | Whole script | |
|---|---|---|
| Park's original | 4.0 s | 4.97 s |
| w/ adjoint trick | 0.19 s | 1.28 s |
Fun part: the solve is now so fast that the solution startup time takes longer thant the solve itself. Ha!
Also of note, that as of writing this, this is also the fastest solution that I could find online. If this is no longer the case, let me know!
A Head Scratcher #
The dataset comes with two output columns: pred (what the correctly ordered model predicted before it fell apart) and true (what actually happened and what the model was likely trained originally to predict). To solve the puzzle we only ever need pred, since we are trying to put the model back together. But as a result the true column never gets used! So I’m a little confused why it was included?!
Code #
The solver, benchmark and an academic-style write-up are on GitHub. Park's original solver is here, and his write-up here.