I Dropped a Neural Net
Organizations: Carnegie Mellon University
Abstract
A recent Dwarkesh Patel podcast with John Collison and Elon Musk featured an interesting puzzle from Jane Street: they trained a neural net, shuffled all 96 layers, and asked to put them back in order. Given unlabelled layers of a Residual Network and its training dataset, we recover the exact ordering of the layers. The problem decomposes into pairing each block's input and output projections ( possibilities) and ordering the reassembled blocks ( possibilities), for a combined search space of , which is more than the atoms in the observable universe. We show that stability conditions during training like dynamic isometry leave the product for correctly paired layers with a negative diagonal structure, allowing us to use diagonal dominance ratio as a signal for pairing. For ordering, we seed-initialize with a rough proxy such as delta-norm or then hill-climb to zero mean squared error.
Figures & tables
| Round | Swaps | MSE | Cum. swaps |
|---|---|---|---|
| 0 (BT init) | – | 0.002989 | 0 |
| 1 | 21 | 0.001025 | 21 |
| 2 | 7 | 0.000553 | 28 |
| 3 | 5 | 0.000200 | 33 |
| 4 | 3 | 0.000052 | 36 |
| 5 | 1 | 0.000000 | 37 |
| Round | Swaps | MSE | Cum. swaps |
|---|---|---|---|
| 0 (seed) | – | 0.036794 | 0 |
| 1 | 23 | 0.021690 | 23 |
| 2 | 20 | 0.012919 | 43 |
| 3 | 21 | 0.008559 | 64 |
| 4 | 15 | 0.005873 | 79 |
| 5 | 11 | 0.004097 | 90 |
| Round | Swaps | MSE | Cum. swaps |
|---|---|---|---|
| 0 (seed) | – | 0.075851 | 0 |
| 1 | 30 | 0.005936 | 30 |
| 2 | 16 | 0.003397 | 46 |
| 3 | 12 | 0.001405 | 58 |
| 4 | 9 | 0.000568 | 67 |
| 5 | 4 | 0.000223 | 71 |
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.