Reverse-engineering the Jane Street Tensor Puzzle

How I decompiled 5,000 layers, tried to solve the wrong problem, and found MD5 hiding in plain sight

Spoiler warning: This writeup gives strong hints about the winning input and describes the complete mechanism of the model from text input to final output.

Jane Street's tensor puzzle provided a large pickled PyTorch model, a text box, and one hint: start with the last two layers. Every input produced 0. It took some poking and prodding at the last part of the model to figure out that the goal was to find an input that produced 1.

The short solution is that the network computes a one-block MD5 hash and tests it against a fixed digest.

The long solution involved more pain and suffering and a lot of dead ends and quitting for months at a time. I spent a substantial amount of time recovering local layer operations, searching for partial successes, and trying to run the entire network backward before recognizing the program those layers implemented. Those failed approaches each left me with some new understanding of the model that made it easier to verify once I had landed on the right hunch.

Layer numbering

For analysis I converted the original pickle into a plain Sequential model so I could hook every leaf module. That artifact has 5,441 modules, numbered L0 through L5440: 2,721 Linear layers alternating with 2,720 ReLUs. It contains about 289 million parameters. The original model adds a final ReLU and also has a custom string-processing wrapper (which my conversion originally mangled adding a costly failure to reconcile against ground truth that plagued me until near the end).

Starting at the end

Bypassing the original final ReLU exposed a much more informative raw output. Ordinary inputs did not all produce the same value of zero; they mostly scored -15. I brute forced the network using a million random keywords and found occasional -14, -13, or -12 results. The final linear layer had a pretty noticeable form:

weights = [+1 repeated 16, -2 repeated 16, +1 repeated 16]
bias    = -15

The preceding layer contained 48 values I saw as three chunks of 16. I began thinking of the 16 positions as independent tumblers in a lock. A score of -14 meant that one tumbler had moved in the right direction; -13 meant two, and a winning +1 appeared to require all 16.

I began looking for hints in the better scoring keywords that would offer some clue about how to unlock more tumblers. Maybe the first tumbler asked something about the length of the input, the second about what character was in the third place, could a tumbler ask what type of object the input was, or what part of speech? (It would turn out that no, it could not!)

One million attempts

I looked through my million generated strings to find patterns in which 16 positions fired. Of those trials, 61,391 lit at least one position, spread across 165 distinct subsets. It dropped sharply from there. No input matched more than three positions so I was still starved for data.

There were no obvious patterns between entries with lower scores. This left me with a few questions:

1. Can I confirm nothing in one corner of the input ties tightly to one byte in the last layer? I knew about diffusion from cryptography, how much was the input diffusing throughout the network?

2. Where are the improved scores coming from? What math is producing them?

The observed counts were:

I implemented some toy gradient and evolutionary approaches to improve on the better results, but they were dead ends producing no meaningful improvement. (I mean they were bad, like, obviously flawed and produced junk theories and proposed inputs that were obviously wastes of time.)

Building a layer decompiler

I then attacked the model as a program expressed in Linear and ReLU layers. The basic workflow was:

1. Run a fixed suite of inputs and cache two adjacent activation layers.
2. Guess how each output span depended on the previous layer.
3. Use targeted mutations to distinguish correlation from causation.
4. Record confirmed formulas and add reusable detectors.
5. Repeat on the next unexplained span.

I started out by doing this by hand. I'd see that often the network would simply transpose the prior layer's, say, positions 32-64, to the next layer's position 128-160. Often it would add some constant to each value in the prior "input" layer.

I found these by hand at first then for 5440 layers realized I needed to automate the process.

This grew into a collection of tools:

• compare_layers.py rendered adjacent activations in a compact symbolic form, so I could visually scan for patterns and identify ops visually

• layer_guesser.py recognized common transformations automatically, I would build operations into it and it would look at each layer and ask "for the unexplained parts of this layer, what if I ran all known operations against chunks of the previous layer, would that explain this?" Then it would adversarially test these against a growing test suite of keywords designed to generate weird results (including some using high unicode, some that were all spaces, some that were the alphabet in order, etc.)

Eventually I wanted more precise confirmation. So:

• probe.py mutated input neurons in the previous layer to discover exact dependencies. This let me discover some rules and relationships that would never be triggered by any input I would ever possibly guess.

The recurring operations were not characteristic of a trained classifier. They were program primitives: copies, negations, decrements, A+B-1, 1-A-B, multiplexers, bit rearrangements, byte unpacking, and weighted binary packing. Values repeatedly cycled through small alphabets such as {-1, 0, 1, 2}, with ReLU acting as a threshold between logic stages.

By identifying classes of common constants my tests could more and more completely reconstruct the entire network. Eventually I had extracted what looked like a set of atomic ops, as if I had found the "instruction set" of the underlying chip.

By this time I had noticed that the layer widths also repeated in highly regular patterns:

From those two facts it felt clear that the cyclical patterns in the layer widths represented more complex operations and functions assembled from underlying operations.

I also noticed:

• Message values were unpacked into bits in 32-bit groups.
• The raw input was copied from each layer into the next, preserving it in a sort of stack where four bytes would be "extracted" at a time for operations. The input survived through almost the entire network and disappeared at L5326.
• An apparent end marker contained 128 followed by a length-like value.
• Near the output, 192 bit-like values were repacked into 16 integers.

The disadvantage of this approach was scale. I became very good at explaining tiny windows, 4-10 layers, without identifying the standard algorithm that fed into.

Trying to run the network backward

Once the final layers were mapped, I tried to synthesize a winning internal state and propagate its requirements backward.

The desired pre-ReLU comparator layer was clear:

L5438 = [-1 repeated 16, 0 repeated 16, +1 repeated 16]

I wrote a rewind engine for inverting known rules, followed by a SAT model for the underdetermined ReLUs and linear combinations. (I later read that one of the early winners, Alex, had tried something similar, which was reassuring since this ultimately ended in a humbling dead end).

The constraint-construction succeeded (in a sense) and discovered millions of independent constraints.

Constructing constraints is not the same as obtaining a valid input.

In hindsight, it explored lots of relationships between arbitrary synthetic states that would never be reachable through any known input, adding immense complexity that prevented this from being viable.

I went back and again tried to trace which early neurons had the biggest impact on late layer neurons. There was basically no signal to pull out, everything was perfectly diffuse, like a cryptographic hash.

I even discussed the puzzle with someone else who was working in a small startup using AI. He was confident that you should be able to trace influential neurons and probabilistically zoom in on the right ones. I told him I'd given up on that, everything was too diffuse. He played with it for a few days and came back and said "yeah, I hoped to see how things propagate, but it really is a bit like a hash."

But that stuck with me. It wasn't just "like" a hash. This was basically a hash function. Did they build one just for this puzzle? Did they invent one accidentally along the way?

That seemed like a massive effort for a simple online puzzle.

Maybe this was literally a hash?

Looking for a known hash

When I returned to the puzzle, the 16-byte target, bitwise logic, byte packing, and repeated layer-width motifs.... Instead of trying to assign semantic names to thousands more neurons, I built a signature search.

I started looking at constants in a new way. When I traced them again, not as individual operations, but sequences near each other, I found something telling. In row 5410, I no longer saw a series of random AFFINE() transformations, I saw the addition of 0x98badcfe being copied in. A quick search revealed that was not random, it was the C register in the IV for MD5.

We don't just resemble a hash function. WE ARE RUNNING MD5!

That means we could find the other IV registers, A-D with a similar search. Confirmed.

That means that there's a hard coded feature in the network near the end that runs an equality test against a target hash.

That target hash is the hash of our answer.

Examining constants in L5437 produces:

c7ef65233c40aa32c2b9ace37595fa7c

Finding the phrase

At this point I just needed to crack a hash.

For those who work in cybersecurity there are a few well known tools and methods to break MD5, one of them far quicker and less resource intensive than any other. It is the first tool you should reach for when you encounter any specific hard-coded string.

I'll leave the tool as an exercise for the reader, but I did not have to break out John the Ripper or cyberchef, it was a bit faster than that.

I had originally tried a brute force across pairs of random english words, after I had the confirmation from the hash I would have expanded that and burned a few more cycles on some local GPUs but otherwise would have eventually resolved it the same way.

It is a bit frustrating to have discounted brute force search only to later realize it was the right track all along!

But it was a rewarding (and sometimes punishing) puzzle. I learned a lot about neural nets simulating boolean operations and the barriers to interpretability in any sufficiently complex network.

I feel like a hash algorithm was one of the examples used about inscrutability in the discussion of Eliciting Latent Knowledge. This was only tractable because brute force was possible. If the keyword were randomly selected from the possible input space then it would be faster to invent a novel attack to break MD5 completely.

So here we are. Not just a fun toy puzzle, but a solvable puzzle illustrating how easy it would be to make an entire class of unsolvable puzzles. An amuse bouche for the limits of the interpretability and corrigibility of AI more generally.

Good luck to us all!