Early this year, Jacob Hilton from Alignment Research Center (ARC) posted about the problem that his team had been tackling, primarily around really mechanistically understanding small (< 1,500 parameters) models. He released his team’s Algo Zoo, with model families of small models of various sizes, with the TLDR being:

They were able to mechanistically describe the smallest 2 models (10 and 32 parameters) but struggled with 432 and more parameters
The generalized difficulty of mechanistically describing models is coming up with a procedure that is at least as efficient as sampling

I really like the focus that is placed on really understanding small models, and find this problem rather interesting in general, so this post will be about walking through some of the post’s findings, deriving and exploring mechanistic descriptions of the small models, and if scope permits, postulating how one might find an efficient way to arrive at a mechanistic description at least as quickly as sampling, or if the comparing the 2 techniques is fundamentally wrong due to the different limitations of each approach (they tell you different kinds of information, after all).

What Counts as a Mechanistic Description?

Hilton’s post expatiates their inability to fully define what a “Mechanistic Description” entails. I will not strive for a formal definition either, but here are some examples of what counts; they all revolve around the idea of “the ability to control”:

You can outline conceptually what the various mechanisms of the model are and hand-define weights to construct the “ideal” mechanisms to achieve no loss.
You can describe precisely what classes of inputs the model will fail on (e.g. “input coordinates falling in this convex cone will produce the wrong answer”).
You can predict what the model’s loss will be given some test dataset sampled from a known distribution.

To me, mechanistic interpretability is all about being able to say with certainty what your model does (i.e. what it can and cannot do), because it is fundamentally a control problem - you understand the model so that you can control and improve it. But as we’ll see, “certainty” is frequently a qualifier that has to be relaxed, due to computational demands.

Problem 1: 2nd Argmax

One of the model families Hilton’s team wrote about was 1-layer RNNs trained to correctly identify the 2nd-argmax of a series of real numbers.

Architecture

All matrix math here is written in math convention instead of torch convention

“The Model $M_{d,n}$ is a 1-layer ReLU RNN with $d$ hidden neurons that takes in a sequence of $n$ real numbers and produces a vector of logit probabilities of length $n$. It has three parameter matrices:”

the input-to-hidden matrix $W^{hi} \in \mathbb{R}^{d \times 1}$
the hidden-to-hidden matrix $W^{hh} \in \mathbb{R}^{d \times d}$
the hidden-to-output matrix $W^{oh} \in \mathbb{R}^{n \times d}$

Hidden Size 2, Sequence Length 2

$M_{2,2}$ is the smallest model (and Sequence Length 2 is a special case since it’s “2nd Argmax” here is essentially “Argmin”), and the team has achieved full understanding of it. The post’s mechanistic description of $M_{2,2}$ offers almost no details, but they do say that an ideal, hand-constructed model, achieves 100% accuracy using these weights:

\[\begin{align*} W^{hi} = \begin{bmatrix} +1 \\ -1 \end{bmatrix}, \space W^{hh} = \begin{bmatrix} -1 & +1 \\ +1 & -1 \end{bmatrix}, \space W^{oh} = \begin{bmatrix} +1 & -1 \\ -1 & +1 \end{bmatrix} \end{align*}\]

This is enough information for me to extract my own mechanistic description of the model.

Mechanistic Description

For the sake of simplicity, I will normalize all weight columns so they have unit length. This will allow me to simply say things like “projection” and “change of basis” without caring about the scaling factor.

\[\begin{align*} W^{hi} = \frac{1}{\sqrt{2}} \begin{bmatrix} +1 \\ -1 \end{bmatrix}, \space W^{hh} = \frac{1}{\sqrt{2}} \begin{bmatrix} -1 & +1 \\ +1 & -1 \end{bmatrix}, \space W^{oh} = \frac{1}{\sqrt{2}} \begin{bmatrix} +1 & -1 \\ -1 & +1 \end{bmatrix} \end{align*}\]

Step 1: input $x_0$ is embedded (projected) onto the span of $W_{hi}$:

Step 2: ReLU

Because a ReLU follows immediately after, the vector space is segmented into a red half-space (first element is preserved through ReLU), and a blue half-space (second element is preserved through ReLU). There’s a cone where the 2 half-spaces overlap, as well as a cone outside either half-space. We can see that positive values of $x_0$ will embed into the blue cone, while negative values of $x_0$ will embed into the red cone. Crucially, this means that all possible values of $x_0$ are preserved.

Step 3: $W_{hh} ( \cdot )$

Next, we have to apply $W_{hh}$ to our current hidden state. This means calculating $W_{hh} W_{hi} x_0$. Visually, it’s just a change of basis, where the new basis vectors are simply the columns of $W_{hh}$ (highlighted blue and red); they take the role of the horizontal and vertical axes of the pre-change-of-basis space.

So this is the transformation (pre and post) of the basis via $W_{hh}$:

Step 4: input $x_1$ is embedded (projected) onto the span of $W_{hi}$ (this time I highlighted it orange to indicate that it’s specific to $x_1$):

Step 5: Add the current embeddings of $x_0$ and $x_1$ to get sum, followed by a ReLU:

Crucially, we see that:

Positive values of sum get embedded into the blue-green region
Negative values of sum get embedded into the brown region
No values of sum get embedded into the region of overlap; this the region of interference (in this particular case it’s ok, but in general interference is bad - out of scope of this post)
All possible values of sum are represented.

The partitioning of negative and positive sum values into different regions is really convenient for the problem. In this case, having a non-zero $h_2$ corresponds to $x_0 > x_1$ and having a non-zero $h_2$ corresponds to $x_1 > x_0$.

Step 6: $W_{ho}(\cdot)$

You’ll see that in this case, $W_{ho}$ doesn’t really do much. The partitioning remains unchanged, which is a good thing. This means that logit[0] will be non-zero iff $x_0 > x_1$ and logit[1] will be non-zero iff $x_1 > x_0$.

There are many solutions for $W_{ho}$ that would have preserved this property all the same, including the identity matrix.

Hidden Size 4, Sequence Length 3

What we did for $M_{2, 2}$ essentially amounts to Guess-And-Check, where we intuited what a plausible mechanism is, and checked to see if the weights actually explain the mechanism. The north-star here is to device a mechanism that allows us to avoid Guess-And-Checking, but this is a generally extremely hard problem in mech interp. All mech interp techniques now rely on intuiting / proposing some plausible mechanism, and searching to see if the mechanism exists. In attempting to work our way to a mechanical technique, we need more insight.

There are a few techniques we may use to interpret $M_{4, 3}$. Just to get things out of the way, I want to say that brute-forcing is possible, since this model is really small. This is something ARC considered “not adequate” - not that they did it (they explored using mathematical arguments instead, but even the approaches they wrote up turned out to be very incomplete and clunky). If a mechanistic description amounts to being able to enumerate all ranges of inputs that the model will classify wrongly, i.e. describing the mapping from input image to output classification, then here it is:

Blue: model thinks 2nd-argmax is 0, Grey: model thinks 2nd-argmax is 1, Red: model thinks 2nd-argmax is 2

matplotlib’s plotting / lack of 3D buffering leaves quite a bit to desire, but the above plot is a union of many convex regions, each a simple intersection of multiple halfspaces. This means I have the entire mathematical description of what points the model will classify as what class, including what will be correct / wrong. To read more about the technique I devised to perform this, you can read “Multi-Layer Latent Space Visualization”.

With that out of the way, let’s explore possible methods and more mathematical methods that may allow us to search for this mechanistic description more quickly. This next approach will still rely on Guess-And-Checking; this model adds more complexity to our analysis because it doesn’t achieve 100% accuracy. The explanations for the mechanisms are not clean because they obviously don’t work all the time.

Hilton’s team reports that $M_{4, 3}$ achieves an accuracy of 98.5%, and in evaluating its accuracy for 40,000 inputs uniformly sampled between $(-20, -20, -20)$ and $(20, 20, 20)$, I got these accuracies:

Accuracy by ground-truth answer

Ground-truth Answer (idx)	Accuracy
`0`	0.9865 (`n`=13349)
`1`	0.9776 (`n`=13230)
`2`	0.9689 (`n`=13421)

Accuracy by input ordering

Input Ordering (idxs)	Accuracy
`0 > 1 > 2`	0.9767 (`n`=6574)
`0 > 2 > 1`	0.9710 (`n`=6690)
`1 > 0 > 2`	0.9881 (`n`=6611)
`1 > 2 > 0`	0.9667 (`n`=6731)
`2 > 0 > 1`	0.9850 (`n`=6738)
`2 > 1 > 0`	0.9785 (`n`=6656)

There didn’t seem to be a major class imbalance, which would have suggested that the model just de-prioritized learning a certain input ordering or ground-truth answer, so I decided to simply visualize the logits of all the inputs, color-coded by ground-truth class:

We can see here that most of the points are correctly classified, with what seems to be 3 major decision groups:

0 > 2 > 1 (orange) and 1 > 2 > 0 (red)
0 > 1 > 2 (blue) and 2 > 1 > 0 (brown)
2 > 0 > 1 (purple) and 1 > 0 > 2 (green)

Notice that each major group contains 2 classes, and the 2 classes seem to be the strict opposite of each other. I initially thought that this implies that there is some feature (neuron + its ReLU) that classifies between forward and reverse ordering, for its specific ordering (e.g. 0 > 2 > 1 (forward) or 1 > 2 > 0 (reverse)). But I think this is an opinionated guess; the most obvious reason is that these major groups are united by the fact that their constituents contain the same answer to the 2nd argmax question.

For other reasons, working backwards was also rather cumbersome, so I will work forwards through the model for this problem.

Weights

Let’s take a look at $W_{hi}$:

\[\begin{align*} W_{hi} = \begin{bmatrix} 3.9653 \\ 4.6995 \\ -4.7163 \\ -3.9652 \end{bmatrix} \end{align*}\]

Notice that this is just a line:

whose direction is positive in dimensions 0 and 1, and negative in dimensions 2 and 3
whose dimensions (0, 1) seem to be a mirror image of dimensions (3, 2).

If we plot the embeddings $W_{hi} x_0$ in these pairs of dimensions, we get:

Notice that orange points correspond to data where $x_0$ was the biggest, while red points correspond to data where $x_0$ was the smallest. You can see this in play in the plots:

Left: orange points appear clustered in the positive orthant, which makes sense, since a larger (more positive) $x_0$ makes it more likely that $x_0$ is the max of the tuple. Conversely red points appear clustered in the negative orthant, since a smaller (more negative) $x_0$ makes it more likely that $x_0$ is the min of the tuple.
Right: The opposite effect takes place, since dimension pair (3, 2) is the mirror image of (0, 1). Note that the gradient of the lines in both plots are about the same due to this.
I’ve also shaded in the regions preserved by ReLU. You can see that dims (0, 1) preserve positive $x_0$, while dims (3, 2) preserve negative $x_0$.

Next, let’s take a look at $W_{hh}$:

\[\begin{align*} W_{hh} = \begin{bmatrix} 1.3387 & -2.1996 & 2.6352 & -1.8948 \\ 1.9277 & -2.4339 & 2.7868 & -2.3826 \\ -2.3877 & 2.8057 & -2.4339 & 1.9331 \\ -1.8931 & 2.6429 & -2.1912 & 1.3409 \end{bmatrix} \end{align*}\]

Things to note - more symmetry:

Columns (0, 1) seem to be a flip (permutation) of columns (3, 2)
Rows (0, 1) seem to be a flip (permutation) of rows (3, 2)

If we think of $W_{hh}$ as a change of basis (COB) just like we did in the case of $M_{2, 2}$, we can visualize the COB by looking at how the axis lines are mapped onto the columns of $W_{hh}$. There’s just one complication - because we’re in 4D now, and we can only really visualize a COB in 2 dimensions, we need to pick 2 source dimensions and 2 target dimensions to visualize. Because the natural pairings of dimensions have been (0, 1) and (3, 2) so far, we can visualize the COBs for each combination of these pairs as source and target:

Note (more symmetry):

The top left COB (src [0, 1] -> tgt [0, 1]) and bottom right COB (src [3, 2] -> tgt [3, 2]) are basically doing the same thing.
The other 2 COBs are also basically doing the same thing.

Let’s add in the embeddings $W_{hi} x_0$ to these COB visualizations:

Notice that the source dimensions (0, 1) contain positive $x_0$, while the source dimensions (3, 2) contain negaitve $x_0$.

Now, let’s add in the embedding of the next number in each input, $W_{hh} x_1$. I’ve plotted these embeddings in the same colors, just fainter and using a different marker, so:

solid orange and solid red round points are $W_{hh} W_{hi} x_0$
faint orange and faint red triangular points are $W_{hi} x_1$

Since they’ll be added together, which gives these embeddings just before the second ReLU:

Notice that in each pair of dims, the diagonal line acts as a “$x_1 - x_0$” and “$x_0 - x_1$” feature respectively. This is completely analogous to the $M_{2, 2}$ case, in that W_hh[:, [0, 1]], and also W_hh[:, [3, 2]], function to flip the embeddings along the span of $W_{hi}$, just like $W_{hh}$ in the $M_{2, 2}$ case. You can see this visually by noticing that the orange-red gradient flip orientation between this plot and this plot.

The only difference is that each of W_hh[:, [0, 1]] and W_hh[:, [3, 2]] is rank-2, whereas $M_{2,2}$’s $W_{hh}$ was rank-1, which is why the scatters span non-zero area now. This 2 dimensionality actually causes some issues for us, which the post-ReLU plots below make clear, so there must be some reason the model needs the 2 dimensionality.

Notice that post-ReLU, we largely “keep” (by this I mean that their embeddings are not the $0$ vector) only points where $x_1 > x_0$ and $x_0 < x_1$ respectively, but not very cleanly. Notice that there are some outliers along the axis lines, not just at the origin. From now on it is useful to treat these embeddings as features for $x_1 > x_0$ and $x_0 > x_1$, which I’ll call $\Delta(x_1, x_0)$.

Let’s transform by $W_{hh}$ again:

And now I’m going to make these $W_{hh} \Delta(x_1, x_0)$ embeddings really faint so that we can also plot $W_{hh} x_2$ in bold:

Notice that unlike the previous 2 times, there isn’t a natural ordering / gradient (red to orange) of the embeddings $W_{hh} x_2$. This is because these 2 categories of points are where $x_2$ was the 2nd-argmax. This means that it is just as likely for $x_2$ to be large as it is to be small (symmetric). If anything, the only pattern is that $x_2$ is more likely to be intermediate (near $0$).

Something rather interesting happens when you add $W_{hh} \Delta(x_1, x_0)$ to $W_{hi} x_2$:

You will notice that the shapes of the respective triangles aren’t the same. If we were to add points sampled uniformly from the diagonal line ($W_{hh} x_2$) to the $W_{hh} \Delta(x_1, x_0)$ triangles, we would expect to end up with something that looks like this:

The fact that we got the above resultant shape instead hints at some more hidden structure between $W_{hh} \Delta(x_1, x_0)$ and $W_{hi} x_2$. This part is not super relevant (feel free to skip) but in case you were curious about how the eventual triangle got formed, hopefully this graphic makes it somewhat clear.

Basically, where we are in the original $W_{hh} \Delta(x_1, x_0)$ triangle tells us 2 things: how large the difference between $x_1$ and $x_0$ is, and where in the number line they are. This imposes restrictions on where $x_2$ can be in order for $x_2$ to still be the 2nd-argmax, so we aren’t really adding uniformly sampled points from $W_{hi} x_2$ to the $W_{hh} \Delta(x_1, x_0)$ triangle.

Ok - back to looking at what’s important for our 2nd-argmax task. Let’s zoom into the origin:

Notice:

The 2 categories are cleanly separated by a linear boundary (not that it matters actually, eventually these 2 categories need to give the same answer (2))
Some red points in the left plot are in the (-, +) orthant
Some orange points in the right plot are in the (-, +) orthant

This has some implications for ReLU:

Look at the how the combined clouds have a “thumb” sticking out of what is otherwise a very nice cone of final embeddings. This “thumb” may or may not be the exact cause of the model’s inaccuracies, depending on how many classes there are and how much space we have.

By “space”, I mean given this number of hidden dimensions, and the properties of $W_{hh}, W_{oh}$ that seem necessary (e.g. symmetry) that may cut down the effective rank of our final embeddings (e.g. hidden_dim –> 0.5 * hidden_dim), can we arrange our final embedding hyperplane (determined by $W_{oh}$, and hence $W_{hh}$ as well), to achieve the desired model logit readouts?

Turns out, the orange and red thumbs are a problem indeed.

Next: Categories 0 > 1 > 2 (blue) and 2 > 1 > 0 (brown)

But first, let’s try to use the same technique to follow our 0 > 1 > 2 (blue) and 2 > 1 > 0 (brown) data points through the model.

Step 1: Plot embeddings $W_{hi} x_0$

Step 2: ReLU

Step 3: Rotate by $W_{hh}$

So far pretty equivalent to the previous categories.

Step 4: Introduce $W_{hi} x_1$

You’ll notice that whereas there was a strong orange-red gradient when we introduced $W_{hi} x_1$ for the previous categories, there isn’t a strong blue-brown gradient here. This is because the 2 categories here are for data samples where $x_1$ is the 2nd argmax. $x_0$’s value has no discriminating effect on $x_1$’s value; conditioned on knowing $x_0$, $x_1$ is still as likely to be small (negative) as it is to be big (positive). This is analagous to the introduction of $W_{hi} x_2$ for the orange-red case.

Step 5: Add $W_hh x_0 + W_{hi} x_1$

Perfectly linear separable. Again, pretty equivalent to the orange and red categories. The (+, +) direction is again analagous to $\Delta (x_1, x_0)$

Step 6: ReLU

Again, same as before; some outliers along the axis lines.

Step 7: Rotate by $W_{hh}$

Same as before.

Step 8: Introduce $W_{hi} x_2$

So far, everything follows the same shapes as in the orange-red case:

But when we try to add the embeddings, things start to change because the implicit structure ($x_2$ needs to be the largest, instead of between $x_0$ and $x_1$) is different:

Here’s a similar graphic illustrating how the addition of $W_{hi} x_2$ attains the final shape:

Step 9: ReLU

Inaccuracy

From here, we can already see why the thumbs are a problem. If we plotted the relu_2_post activations for all orange, red, blue, and brown points, we see that they literally inhabit the same space, so they have to be linearly separable:

You can see that:

the blue thumb is in the main orange cone
the red thumb is in the main brown cone
the orange thumb is in the main blue cone
the brown thumb is in the main red cone

Algo Zoo (WIP)

What Counts as a Mechanistic Description?

Problem 1: 2nd Argmax

Architecture