This is research I did in a short span of time, it is likely not optimal, but it's unclear whether the constraints are my skills, my tools, or my methods. Code and full results can be found here, but you'll need to download the model yourself to replicate it.

TL;DR

Using Gemma Scope, I am able to find connections between SAE features in the attention portion of Gemma 2 2B's layer 13. This can be done without running the model. These are often (somewhat) meaningful, with the OV circuits generally being easier to interpret than the QK circuits.

The Background

Anthropic has done some nice research decomposing attention layers into the QK and OV circuits. To summarize, the QK circuit tells an attention head where to look, and the OV circuit tells it what information to pass forward. When looking at a single-layer transformer, this looks like patterns in tokens of the form $A,...,B\to C$.

Recently, Google Deepmind released Gemma Scope, consisting of a large number of SAEs, trained on the Gemma 2 family of models. Neuronpedia has already generated a description for all of the "canonical" SAEs using GPT-4o.

This is an attempt to look at QK and OV circuits in Gemma 2 from the lens of features, rather than tokens.

The Setup

I chose to look at layer 13 (indexed from zero) roughly halfway through the model. For each decoder vector in the 16k-feature canonical residual stream SAE of layer 12 (henceforth SAE12), I calculated the element-wise product with the RMS norm coefficients, then calculated the products with $W_Q$, $W_K$, and $W_V$. I'll refer to these as the queries, keys, and values.

I then took the encoder vectors of the 16k-feature canonical residual stream SAE of layer 13 (henceforth SAE13), did an element-wise product with the RMS norm of the attention layer output, and multiplied by the transpose of $W_O$. This creates a vector I'll call the pre-output. This is working backward towards the attention layer.

This does ignore the MLP layer. I wish there were SAEs trained between the two components. It also ignores the rotary positional encoding. I may return to these later. It also ignores the attention-out prelinear SAE! This is because Neuronpedia doesn't have labels available to download for these yet.

I calculated, for each head, a $[16k,16k]$ matrix consisting of the dot-products of the queries and keys, and which contains the dot-products of the values and pre-outputs. For each head, I found the highest-weighted connections between features in the QK and OV circuits.

Decoder Bias vs Rotary Embedding

I take the decoder bias of SAE12 to represent the "average" value of the residual stream going into layer 13. By looking at the query-key dot product of this with itself at different positional distances, we can maybe kinda see how each head looks across tokens:

Seems like head 0 is mostly a previous-token head, whereas the others fall off more slowly over distance.

Interesting Findings

Head 0 has lower QK-weighting than the other heads. It does a few interesting functions, such as noticing when mathematical constants are present and down-weighting a feature related to social equity (presumably to point the model towards the correct concept of "equal to" in this context).

It also detects a feature relating to scientific context and up-weights a feature relating to  the past tense, specifically completed actions, which is indeed the most common tense for scientific writing to be in!

Lots of OV circuits are pretty interpretable, or at least seem that way: features up or down-weight later features appropriately.

Unfortunately, many of the feature labels are not very good. For example, I keep getting ones relating to "product descriptions" cropping up in unrelated text in the neuronpedia playground. I assume a more expensive model would do a better job. Also, the features in a 16k SAE for a 2k residual stream are not very monosemantic. It would be interesting to try the 256k SAEs once those are labelled.

It seems like many of the heads are deep in attention-head superposition. If the attention SAE gets labelled it would be cool to check that out.

Some features come up repeatedly in various heads. I think this is because they're up-weighted by the coefficients of the layernorm.

Diving Into Head 0

QK Circuits

Seems like several themes are coming up here: programming and numerical data, particularly lists; names and individuals; scientific research; legal proceedings. Overall it seems like this head is being used for a few different things. Makes sense, given what we know about superposition. Remember that head 0 might be mostly a previous-token head.

We also see a lot of what I call like-to-like QK-connections, in which the same feature appears as the query, and the key. This also makes sense intuitively. I'll show the first 20 connections which are not like-to-like, and which have positive QK values:

One thing to note here is that the QK values for this head are much lower in magnitude than the other heads. Perhaps this head takes the role of a general aggregator, picking up on vibes from lots of tokens, rather than passing specific information around. This kinda makes sense based on the sorts of things which crop up in the OV circuit:

OV Circuit

These are much richer and more interesting.

I really like 22 here, because it seems a bit weird at first, but what it's actually saying is that scientific work is almost always written in the perfect tense!

I find the negative values more interesting than the positive ones in a lot of cases. A lot of them seem to be "disambiguation". 8 seems to be telling the network "no, this is maths, we're looking at the mathematical definition of equality, not the social one!", as does 29. 16 seems to disambiguate Swift or Objective-C code from mathematical notations!

I don't know what 0 or 1 are doing! Why would "little" mean there are no prepositions, but we're in a scientific or mathematical context!

Head 7

I'll go through Head 7 here as well:

There are still a lot of like-to-like pairs. We some mutual-exclusion tetrads, like 8-11, or 26-29. Let's show the non like-to-like pairs with positive QK values:

Now let's take a look at the OV circuit:

Some of these seem totally nonsensical! 8-11 in particular, I mean what? These seem to be hopelessly lost in superposition. Perhaps head 7 is in a greater degree of superposition because it attends to more specific tokens than head 0.

Conclusions

With difficulty, it may be possible to reconstruct attention-based circuits in this way. It's unclear how much of the difficulties stem from technical limitations in the SAE and labels, and how much are fundamental to this method. I would like to try again someday.