Algo Zoo (WIP)

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”:

1. 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.
2. 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”).
3. 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, we see that only the half-space that is pointed to by $W_{ih}$ is preserved (highlighted grey); everything else is zero-ed:

This means that if $x_0$ were a negative number, the model effectively treats it as a $0$.

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.