Stealing a Trick from Nuclear Physics to Solve a Puzzle

A blogger built on Hyunwoo Park's solution to Jane Street's "dropped a neural net" puzzle by applying the adjoint state method from nuclear engineering, cutting the layer-reordering search from roughly 3-4 seconds to a fraction of that. The puzzle asks solvers to restore a trained 48-block, 97-layer neural network that was split into 97 shuffled layers, using the shuffled pieces plus 10,000 rows of historical inputs and outputs. The key insight is that predicting a layer swap's effect on overall error costs 4 GEMMs plus a dot product, versus 96 GEMMs for a full forward pass, with one forward and one backward pass amortized across all predictions.

Stealing a Trick from Nuclear Physics to Solve a Puzzle Jane Street dropped a neural net puzzle https://huggingface.co/spaces/jane-street/droppedaneuralnet . This is a bit before my time but I am assuming that this is a reference to the terror of dropping punch cards https://en.wikipedia.org/wiki/Computer programming in the punched card era Identification and sequence 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 https://hynwprk.github.io/blog/2026/i-dropped-a-neural-net/ 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: 1. 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. 2. Guess an order. Sort the blocks by the size of their relative weights originally L2 https://github.com/hynwprk/droppedaneuralnet/pull/1 . The hunch here is that later layers tend to have bigger weights vs earlier ones. 3. 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 https://en.wikipedia.org/wiki/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 https://www.osti.gov/pages/servlets/purl/1438213 . 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 https://scale-manual.ornl.gov/6.3.3/tsunami-1d.html tools still work on this basic concept. Details on the math for this particular problem are here https://github.com/cezarc1/droppedaneuralnet-puzzle/blob/main/academic puzzle note.pdf . 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 https://github.com/cezarc1/droppedaneuralnet-puzzle/blob/main/academic puzzle note.pdf are on GitHub https://github.com/cezarc1/droppedaneuralnet-puzzle . Park's original solver is here https://github.com/hynwprk/droppedaneuralnet , and his write-up here https://hynwprk.github.io/blog/2026/i-dropped-a-neural-net/ .