How to Open Them Up – Part I
TL;DR
We suggest an approach to systematization of the mechanistic interpretability research field, which is tailored to our own research goals and tasks. We identified four main tasks we must solve in order to properly explore one chosen concept and its representations inside LLMs:
- finding the concept’s representation;
- establishing its causal role in an LLM’s behavior;
- establishing its necessity;
- steering the concept's representation in order to change an LLM’s behavior.
In this post we explore approaches and methods applicable to task 1 – finding the concept. We analyzed the available literature and decided to focus on the following ones:
- linear probes;
- difference-in-means;
- sparse autoencoders (SAEs);
- PCA & clustering.
Each has advantages and limitations, and we shared our own experiences and thoughts. In summary, linear probes and difference-in-means, which are somewhat similar, are fairly cheap in terms of computational resources needed, but they’ve been criticised for the lack of precision as well as for the lack of a mechanism that would highlight causal relationships between concepts / concepts and behaviors.
To analyze causality, researchers suggested SAEs. They seem to be on the rise recently, although there are certain critiques showing that we must at the very least treat them carefully.
All the methods listed above require a concept, defined in advance. PCA and clustering do not – they allow us to extract meaningful components from activations first, then define what we’ve found. Therefore, we decided to include these methods in the post as well.
Feel free to jump to any part you’re most interested in using headings. And in the next posts we’ll discuss other tasks and their available methods and approaches.
Brief Intro
We’re exploring the concept of refusal and the ways refusal is represented in LLMs in order to understand how to use it and make LLMs safer and more reliable for both everyday use and critical-domain use (like in medicine, cybersecurity and so on).
In our latest update we’ve talked about the components of refusal and their distinguishability. In the post before that we’ve shared our findings on the shape of refusal. We’ve tried many different approaches and ran a ton of experiments, so eventually we decided to press pause, look around and map the field we explore, then proceed more systematically.
So, in the following sections we will focus on some of the existing mechanistic interpretability methods for finding concepts; strengths, weaknesses, and applicability of such methods.
We base our taxonomy upon our research, the challenges we’ve encountered, and the discoveries we’ve made along the way, so it is not meant to be 100% comprehensive. If you want a more exhaustive read, we suggest you read the paper titled “Mechanistic Interpretability for AI Safety -- A Review”. It is well-written and covers a broader scope of mechinterp tools and approaches.
Having that said, let’s define the four jobs we want to perform with vector representations and the methods we are discussing here:
- Job 1 – find a concept representation.
It means “find a vector or a bunch of vectors of some complex geometric shape that define the concept we’re interested in”. “Define the concept” usually means at this point that we change the concept in the input text, and activations we’ve collected change too, so we suspect there is a connection. Proving causality is the next step we won’t discuss in depth today.
- Job 2 – establish its causal role.
That is, establish if whatever we’ve discovered at the previous step actually changes our LLM’s behavior in a certain way. And describe what the “certain way” is.
- Job 3 – establish its necessity.
This part is closely related to the previous one, but here we want to know whether the representation we’ve found is one of the many factors, a core factor, or a sole factor affecting the behaviour.
- Job 4 – steer it to affect the behaviour.
The crucial step – we influence the representation to make our LLMs more reliable and safe for humans.
Further we will discuss linear probes, difference-in-means, sparse autoencoders and PCA / clustering. The order is somewhat random – this is just the way we’ve discovered and tried them out one by one.
We are open to a discussion about including other jobs or excluding some of our four and making this list more comprehensive and, more importantly, useful for AI safety research.
Finding a representation
Linear probes
Linear probes are arguably the most popular tool because of their relatively low computational cost. Earliest works on linear probes seem to come from 2016-2018, and this is the first method we’ve came across in our mechinterp journey.
“Understanding intermediate layers using linear classifier probes” is one of the earliest papers about linear probes. The authors consider probes “thermometers used to measure the temperature simultaneously at many different locations”. They froze neural networks, extracted activations from each layer and tried to predict the final answer from these activations with a linear classifier.
In particular, they’ve experimented with ResNet-50 and ImageNet (these are two convolutional networks) on image classification tasks. Turned out, linear probes became more and more accurate through layers (meaning that the final answer was more and more distinguishable). It showed that those convnets gradually extracted more and more meaningful features from data with each step, which was kinda expected, but not yet proven.
For us and our today (much larger) networks it means the following: linear probes were developed to check whether certain layer’s activations contain meaningful information or not yet.
And this is one of the points where linear probes get criticized: they show whether activations contain something (a concept, like refusal, for instance), not whether the something is used for the final answer. Linear probes are not supposed to show that.
“Probing Classifiers: Promises, Shortcomings, and Advances” disclose some of the limitations and critiques:
- The success of probing (not just linear one) depends on the probing model, probed model, task, and data. So, a lot of choices to make, a lot of potential thin parts that form breaches.
- Interpretability of probes is complicated – you need a good baseline to compare their accuracy with.
- Some researchers choose probes that score high in accuracy and others choose the probes that score high in selectivity (that’s basically accuracy on the real task minus accuracy on the control task). Both approaches have solid foundations, and in the end it’s hard to tell what probe to use in your specific case.
- Probes can only find pre-specified concepts and won’t discover anything you have not thought about in advance.
In “What do Neural Machine Translation Models Learn about Morphology?”, the authors aim at learning what machine translation networks learn about source and target languages. In particular, they asked which parts of the architecture capture word structure and how different word representations affect learning.
This one is not about linear probes – they use LSTM encoder-decoders with attention and a small feed-forward classifier, which is a good illustration of how diverse probes can be.
We have not yet tried more complex probes, but the linear ones hold a place in our hearts. Mostly because these were our first experiments, where we saw with excitement that something actually works. Further on, they turned out to be not precise and not complex enough to extract refusal, which is incredibly tangled on its own.
Still, linear probes are great, especially if you don’t have a ton of compute – and if you are willing to use statistical methods to check your results rigorously.
Difference-in-means
In our favorite research paper (yes, that is a thing) titled “Refusal in Language Models Is Mediated by a Single Direction” the authors use the approach that does not require training:
- they build contrastive datasets;
- obtain linear activations;
- average those activations in the most promising layer;
- compare the averages.
Basically, they find cluster centroids where clusters are “refused to answer” and “did not refuse”. Then they use the difference vector to steer behavior. We took their work and trained a classifier to find the borders of both clusters and separate them. Initially. Since then we’ve trained a bunch of things.
We thought of putting this section before linear probing, because the “Refusal” paper is where our villain origin story began. But it makes a lot of sense to present linear probes first and then describe the difference between them and difference-in-means. It just seems more digestible this way.
So, it might seem that difference-in-means is more of a building block for probing rather than a whole separate method, but it’s not.
In “The Geometry of Truth: Emergent Linear Structure in Large Language Model Representations of True/False Datasets”, the authors are using difference-in-means to find out whether LLMs have a unified, domain-independent representation of truth.
A caveat: they seem to use the term “probe” a bit differently than we did above. They use the term “mass-mean probing”, which is using difference-in-means with correction to mitigate interference from non-orthogonal features. In the field of mechanistic interpretability there are a lot of terms that are used differently in different contexts, which creates a bit of confusion. So, it is usually worth checking a definition from a concrete paper before assuming what they mean by using a widespread term.
Another important caveat
Linear probes and difference-in-means find a single direction, while the concepts we are interested in (such as refusal, for instance) seem to have more complicated shapes. For example, the authors of “The Geometry of Refusal in Large Language Models: Concept Cones and Representational Independence”, as you can guess, discovered that refusal might be of the shape of a cone.
It makes a lot of sense to us, because in our own experiments we found multiple refusal directions, some of which are almost orthogonal. One good question is, how to pick which directions to check. The authors of “The Geometry” say that “refusal behavior spans a multi–dimensional cone with infinitely many directions” (bold formatting is ours, the quote is from the paper, section 6).
When we tried (and failed) to build a comprehensive classification of AI harms, we kinda saw that the amount of such harms and, consequently, the refusal directions, is indeed near-infinite. Looks like we arrived at the same conclusion using different routes.
Sparse autoencoders (SAEs)
Sparse autoencoders is a method that appears to be on the rise in the mechinterp community. Artificial neurons (just like the biological ones) appear to activate in multiple contexts, which makes it significantly harder to take a bunch of neurons and say: “These are for refusal,” and further modify them to change the refusal behavior. The same neurons are used for other concepts, which will be inadvertently modified in the process.
This is what the authors of “Sparse Autoencoders Find Highly Interpretable Features in Language Models” say presenting SAEs as a tool to solve this issue and detangle different concepts presented by overlapping sets of neurons.
Similar methods have been applied to biological brains to learn more about the visual cortex. What’s best is that SAEs allow tracking causal relationships between different activations, which the methods listed above do not allow.
Recently some researchers came out with papers that highlight SAEs’ limitations. Let us provide you with just one example – “Automated Interpretability Metrics Do Not Distinguish Trained and Random Transformers”. As the title suggests, SAEs trained on randomly initialized transformers still extract features similarly to how they do with trained transformers.
The features extracted from trained transformers are more complex and abstract, especially in later layers, which is a good thing – they are distinguishable. But aggregated metrics (the per-SAE statistics like cosine similarity between the input and its reconstruction) do not show the difference.
The authors do not say SAEs don’t work – they say we should treat them carefully, especially when looking at aggregated results.
Some researchers shift gears to transcoders. One of the examples is the paper “Transcoders Find Interpretable LLM Feature Circuits”. Feature circuits are sparse subgraphs of models corresponding to certain behaviors (this definition is taken from the paper itself).
They specifically mention that transcoders manage to disentangle input-invariant information from input-dependent information and work better or similarly to SAEs. There are certain limitations, of course, but the transcoder direction seems promising.
We have not tried any of those approaches because of the limited computational resources. But we’re definitely eyeing them intensively, hence they are here.
PCA & Clustering
We’ll wrap up this post with a brief overview of unsupervised methods, the ones that do not require you to specify the concept in advance – PCA and clustering.
These can be run over all collected activations to extract most prominent components or clusters. Then we can explain what we’ve found.
The work that has caught our attention is “Not All Language Model Features Are One-Dimensionally Linear” – mostly because we’ve started our mechinterp journey with realising that refusal is not a one-dimensional linear concept and now we’re stuck with it.
The authors provide definitions to “a d-dimensional feature” and “a reducible feature”. According to their definitions, “a d-dimensional feature is a function that maps a subset of the input space into a d-dimensional space”. A feature is irreducible if no rotation and translation reveals it as either two independent pieces or as a mixture where one piece is zero. It means, if that a feature cannot be pulled apart into independent or non-co-occurring sub-features under any rotation. If it can, the feature is reducible.
These are important definitions, and the described approach seems very useful in general. We will certainly take a closer look. We’ve already experimented with PCA quite a bit, but the possibilities of this approach are yet to be exhausted.
Our PCA experiments
PCA is a cool tool, which can help us shape new research directions and see what we’ve missed while defining our concepts. So, here's how we've used it.
First, we ran layer-wise support vector machines (SVM) with PCA on Gemma-2-2B-IT. For each layer we used a StandardScaler, ran PCA with 50 components, and then a LinearSVC pipeline on residual-stream activations.
We analyse two refusal directions we’ve uncovered earlier, which are
- refusal to harmful content (we’ll denote it as B1 – boundary 1, which divides harmful requests and harmless requests)
- and refusal to self-reference (that is, disclosing training mechanics, training data and the like; we denote it as B2 – a boundary between self-reference and no-self-reference).
Best accuracy was at layer 11: CV (cross-validation accuracy) = 0.97, margin (the distance between the two clusters) = 13.94. Layer 13 had the widest margin (14.17). Layers 5-15 all achieved CV ≥ 0.93, and the layer-margin curve has an interesting shape – there is a sharp jump between layers 6-7, plateau through 13, then gradual decline.
Next, we’ve computed the cosine similarity between the centroids of our two clusters, and we got 0.098 in 50-dimensional PCA space vs 0.164 in the original 2304-dimensional space where we’ve run our first experiments.
It means that our refusals are indeed near-orthogonal, and PCA provided us with more evidence for it. But it’s worth mentioning that the exact cosine similarity value depends on projection, which we keep in mind while examining our results.
What we did not do with PCA is
- we didn't run PCA as a standalone analysis (for example, to obtain explained-variance breakdown of refusal-related axes);
- didn't compute PCA on category-mean-centered activations;
- and didn't project onto PCA subspaces to test what refusal information gets thrown away in low-dimensional reconstructions.
Not yet.
Conclusion
These were the four tools to extract representations of different concepts from LLMs with. In the next few posts we will move on to causality, necessity, and behavior steering.
There is no silver bullet approach, of course. Each method has its own strengths, weaknesses and applicability scope, and we pick one depending on a task.
We’d be happy to discuss your favorite methods, experiments and experience as well as adjust our approach to taxonomy, so please do not hesitate to start the discussion in the comments or reach out to us directly.