A lot of “capture-the-flag” style ML puzzles give you a black box neural net, and your job is to figure out what it does. When we were thinking of creating our own ML puzzle early last year, we wanted to do something a little different. We thought it’d be neat to give users a complete specification of the neural net, weights and all. They would then be forced to use the tools of mechanistic interpretability to reverse engineer the network—which is a situation we sometimes find ourselves facing in our own research, when trying to interpret features of complex models.
We published the puzzle last February. At the time, we weren’t even sure it was solvable. The neural network we’d designed would output 0 for almost all inputs. A reasonable solver might assume that the goal was to furnish an input that produced 1 or some other nonzero value. But we’d engineered the network in such a way, as you’ll soon see, that you couldn’t use traditional methods to brute force your way to an answer—say, by backpropagating a nonzero output all the way back to the input layer. You had to actually think about what the net was doing.
We were amazed by the response the puzzle got. Mostly by luck, it seemed like we’d calibrated the difficulty just so: it wasn’t so hard that no one could solve it, and wasn’t so easy that we were flooded with responses. In fact if you can solve this puzzle, there’s a decent chance you’d fit in well here at Jane Street.
We’ll restate the problem below, but be warned that the rest of this post contains huge spoilers. If you want to try solving the puzzle yourself, avert your eyes. The rest of this post will walk through the process that an actual solver took, with all the twists and turns before they finally cracked it.
The problem
Today I went on a hike and found a pile of tensors hidden underneath a neolithic burial mound! I sent it over to the local neural plumber, and they managed to cobble together this.
Anyway, I’m not sure what it does yet, but it must have been important to this past civilization. Maybe start by looking at the last two layers.
Model Input
vegetable dogModel Output
0If you do figure it out, please let us know.
That model.pt file is basically just a pickled PyTorch model.
A solution
Getting started
A senior at university named Alex was in his dorm room when a roommate told him about a puzzle that was making the rounds on Twitter. The roommate had tried it himself but given up after two nights. Alex, in his final winter at school, was looking for something to do and decided to have a look.
He started by downloading the model and poking around, focusing on the last layer in particular:
import torch
import plotly.express as px
model = torch.load('./model.pt')
linears = [x for x in model if isinstance(x, torch.nn.Linear)]
px.imshow(linears[-1].weight.detach())
Immediately it was plain that this was not an ordinary neural network. It clearly hadn’t been trained: all the weights had integer values. Instead, it had been designed by hand, probably to carry out some very specific computation.
The last layer was a 48x1 matrix, but apparently broken into three sections. And indeed if you looked at the activations from the previous layer, they were always three repetitions of the same thing. The second-to-last layer appeared to be three repetitions of the same weights, while its bias contained the same 16 bytes, but incremented by 1 each time, as if encoding a vector v, then v + 1, and v + 2. Here’s what the weights on that second-to-last layer looked like:
px.imshow(linears[-1].weight.detach())
and the biases:
px.imshow(linears[-2].bias.detach().unsqueeze(0))
Thinking about it some—and about the fact that the last layer emitted a single bit—Alex realized that this second-to-last ReLU layer must be computing whether two 16-byte integers were equal to one another (with one byte per neuron). The way it seemed to work is that it made three copies of the input vector v, a 16-byte number. It tried to check that against a reference number x (which was determined by the bias of the second-to-last layer). So the three copies would actually represent v - x - 1, v - x, and v - x + 1. The last layer applied weights 1, -2, and 1 to these cases respectively. We can do some casework on an individual value here: consider the value of ReLU(v-x-1) - 2ReLU(v-x) + ReLU(v-x+1). If v=x, then this is equal to 1. We won’t show the rest of the cases here, but they all result in 0. The bias on the last layer was -15, so the final neuron would only fire when v=x for all 16 bytes.
So now the question became, how do we get the activations of the second last layer to equal x?
Reverse-engineering the program at the heart of the network
Alex figured that if there’s some number that the network is checking against at the very end, then the rest of the network must be some sort of big equation. There indeed appears to be a lot of structure in the network, as you can see just from plotting the size of the 2500 linear layers (about half the full network):
px.line([l.out_features for l in linears])
So Alex began looking at various sub-networks, tracing their dependencies. This involved staring at a lot of graph structures:
