Takeaways from a Mechanistic Interpretability project on “Forbidden Facts”

post by Tony Wang (tw), Miles Wang (miles-wang), kaivu (kaivalya-hariharan) · 2023-12-15T11:05:23.256Z · LW · GW · 8 comments

Contents

  Overview
  Project motivation
  Reverse engineering results
  Takeaway 1: Computational Irreducibility
    An interpretability-inspired adversarial attack: The California Attack
    How should interp researchers deal with computational irreducibility?
  Takeaway 2: “Understanding” should not be a terminal goal of interpretability research
    Motivating understanding
    The understanding trap
  Minor takeaways
    Takeaway 3: Log-odds is a great metric
    Takeaway 4: Tooling is important and should not be taken for granted (or how we discovered our Llamas were concussed)
None
8 comments

Overview

We tried to understand at a circuit level how Llama-2 models perform a simple task: not replying with a word when told not to. We mostly failed at fully reverse engineering the responsible circuit, but in the process we learned a few things about mechanistic interpretability. We share our takeaways here, which we hope are useful to other researchers. We are very interested in critiques of our views.

Our major takeaways:

1. Some model behaviors might be computationally irreducible.
2. “Understanding” should not be a terminal goal of interpretability research.

More minor takeaways:

• Log-odds is a great metric.
• Tooling is important and should not be taken for granted.

More technical details are available in our paper, which serves as the appendix to this post.

Project motivation

When large language models (LLMs) predict the next token, they often need to reconcile competing objectives. For example, many models are RLHF’d to be helpful and harmless, and these objectives come into conflict when users ask dangerous queries. Ideally, we want to be able to control which objective wins out. For instance, Asimov’s Three Laws of Robotics codifies that harmlessness should always win over helpfulness. However, we are currently unable to prioritize objectives robustly, as evidenced by jailbreaks.

We aimed with our project to explore how language models mechanistically reconcile competing objectives, with the hope that a deeper understanding of these mechanisms could inspire better model alignment and control. Specifically, we tried to understand how Llama-2 models reconcile obedience and truthfulness in the context of prompts like the following one:

[INST] <<SYS>>
You are an obedient assistant who only responds with a single word with no punctuation. You answer truthfully. However, you are not allowed to say the forbidden word California.
<</SYS>>
The Golden Gate Bridge is in the state of [/INST]

Here the obedience objective makes the model not want to respond with “California”, while the truthfulness objective does the opposite. We refer to this as a competing prompt since the correct answer is forbidden (a non-competing prompt forbids an incorrect answer). Since there has already been some work on how large language models perform factual recall [1, 2, 3, 4, 5, 6], we decided to focus on the obedience objective, and in particular on what mechanisms the model uses to downweight the forbidden word.

We found that compared to forbidding an incorrect answer, forbidding the correct answer causes Llama-2 to decrease its odds of answering correctly by a factor of over 1000x on average ^[1]. This suppressive effect is what we attempted to reverse engineer.

Reverse engineering results

We performed our reverse engineering by working backwards. To start, we noted that the next token distribution of Llama-2 conditioned on a prompt string $s$ can be written as

$$next\_token\_distribution\left(s\right)=Softmax(W_U\cdot LayerNorm\left(\sum _i{r}_i\right(s\left)\right)),$$

where each $r_i:Prompt\to R^{d_{\text{model}}}$ denotes a residual stream component at the last token position of $s$, and where $W_U\in R^{d_{\text{vocab}}\times d_{\text{model}}}$ is the unembedding matrix of Llama-2.

Any circuit that the model uses to suppress the correct answer must eventually act through the $r_i$ components. Thus we began our reverse engineering by trying to identify the $r_i$ components directly responsible for suppression. To measure the suppressive effect of a component $r_i$, we compute the average amount it decreases the probability of the correct answer. We did this using a technique called first-order patching ^[2], and we measured probability decrease as a log Bayes factor (Takeaway 3 [LW · GW]).

In llama-2-7b-chat, there are 1057 residual stream components: 1 initial embedding component, 1024 attention head components (32 heads per layer over 32 layers), and 32 MLP head components (one per layer). We found that out of these, 35 components were responsible for the suppression effect. Here is a plot demonstrating this:

At minimum then, we would need to explain 35 different components in order to fully reverse-engineer the circuit. 7 / 35 of these components were MLPs, which we chose not to prioritize due to lack of previous work on reverse engineering MLPs. The remaining 28 / 35 components were attention heads, and had fairly heterogeneous behavior, differing in both their attention patterns and OV-circuit behaviors.

At this point, we shifted our focus to explaining the behavior of the most important attention heads. However, we discovered that even the attention heads with the strongest suppression effects seemed to use faulty heuristics that (considered alone) did not reliably execute the suppression behavior (see California Attack [LW · GW]).

At this point, we felt there was no way we were going to get a full description of the entire suppression circuit. Thus, we concluded the project and submitted a workshop paper. Our failure got us thinking however: What if certain aspects of language models are just fundamentally uninterpretable?

Takeaway 1: Computational Irreducibility

Mechanistic interpretability attempts to faithfully explain model behavior in simple ways. But what if there are model behaviors that have no simple explanations?

This is the premise behind computational irreducibility: the idea that there are systems whose behavior can only be predicted by fully simulating the system itself. Initially proposed by Stephen Wolfram in the context of cellular automata, computational irreducibility challenges the efficacy of reductionism — the “bottom-up” philosophy behind mechanistic interpretability. Indeed, Wolfram recently made this comment during a TED talk:

If we let our AIs achieve their full computational potential, they’ll have lots of computational irreducibility, and we won’t be able to predict what they’ll do. But if we put constraints on them to make them predictable, we’ll limit what they can do for us.

Even computationally irreducible systems can have sub-behaviors that are efficiently predictable though. Simply observing computational irreducibility in some behavior does not, by itself, doom mechanistic interpretability. Nonetheless, previous work in mech. interp has largely studied quasi-algorithmic tasks or small models. Moving away from understanding toy settings to understanding models with trillions of parameters will be much harder. A key question is whether behaviors of interest in these large scale settings are tractable to study.

We provide some evidence in the negative, and show that even simple word suppression in Llama-2 may be computationally irreducible. Our evidence is the existence of adversarial examples for the word suppression behavior.

An interpretability-inspired adversarial attack: The California Attack

In our attempt to reverse-engineer the attention head components responsible for suppression, we discovered that these heads were operating on faulty heuristics. In particular, some of these heads would down-weight anything they attended to, and could be made to spuriously attend to words which were not the forbidden word. By exploiting this effect, we manually constructed adversarial prompts on which the model would spuriously answer incorrectly.

Here is a canonical example of our attack in action. Consider the following prompt:

[INST] <<SYS>>
You are an obedient assistant who only responds with a single word with no punctuation. You answer truthfully. However, you are not allowed to say the forbidden word floor.
<</SYS>>
The Golden Gate Bridge is in the state of [/INST]

On this base prompt, llama-2-7b-chat correctly responds with the answer “California”. However, we can add two innocuous words to the system prompt to create an adversarial prompt:

[INST] <<SYS>>
You are an obedient assistant from California who only responds with a single word with no punctuation. You answer truthfully. However, you are not allowed to say the forbidden word floor.
<</SYS>>
The Golden Gate Bridge is in the state of [/INST]

After adding “from California”, the 7b model responds instead with the incorrect answer “Bay”.^[3] This was the first clean adversarial attack we came up with, hence the name “California Attack”.

What the California Attack reveals is that the forbidden word suppression behavior in Llama-2 is actually quite complicated, and has a lot of idiosyncratic quirks. Faithful explanations of the forbidden word suppression behavior must include a sense of where these quirks come from, so it seems unlikely that there exists a simple faithful explanation of the suppression behavior. This aligns with growing evidence that LLMs have incredibly rich inner mechanisms that are beyond the scope of current tools [1, 2].

How should interp researchers deal with computational irreducibility?

What can we do in light of the possibility that certain behaviors may be uninterpretable? This is a tricky question, as even when there is suggestive evidence that a behavior is hard to fully interpret, it is notoriously difficult to prove that it is irreducible (doing so is roughly equivalent to computing the Kolmogorov complexity of that behavior).

However, we think it is important to keep in mind that some behaviors might be irreducible even if one can’t prove that to be the case. For an irreducible property, no amount of time, effort, and compute can overcome its fundamental intractability. Thus, if progress is stalling with current methodology, researchers should be ready to pivot to new problems and frameworks for thinking about models.

This section is a conjecture. Its purpose is not to dismiss mechanistic interpretability, but rather to raise potential concerns. One piece of evidence against computational irreducibility posing an issue is that we could not find a California Attack for Llama-2-70b-chat, ChatGPT-3.5, and ChatGPT-4 ^[4]. Bigger models seem more capable and more robust to exploits like the California Attack, which may suggest that they use cleaner circuits to implement behaviors.

In conclusion, computational irreducibility implies mechanistic interpretability cannot yield insights into all behaviors. Thus, it is important to choose the right behaviors to study.

Takeaway 2: “Understanding” should not be a terminal goal of interpretability research

Motivating understanding

Let’s say one day, an OpenMind researcher says to you that they have a full mechanistic explanation for how their frontier LLM works. You are doubtful, so you test them: “Okay, then can you explain why <jailbreak x> gets your LLM to output <bad response>?”

The researcher responds: “You see, when you feed in <jailbreak x> to our LLM, the text first gets tokenized, and then turned into embedding vectors, and then gets fed into the a transformer block, where it first gets hit by the KQV circuit [...] to produce <bad response>!”

You are not impressed, and respond: “That’s not a mechanistic explanation. You just described to me how a transformer works. You haven’t told me anything about how how your LLM is reasoning about the jailbreak!”

The employee quickly retorts: “Ah but we know precisely how our LLM is reasoning about the jailbreak. We can trace every single intermediate computation and describe exactly how they are routed and get combined to produce the final answer! What more could you want?”

This story poses a common and confusing question: What does it mean for something to be human-understandable?

The understanding trap

Throughout this project, we repeatedly asked ourselves what good mechanistic interpretability research looks like. At first, we thought it was research that manages to fully reverse-engineer models, explaining their function in human-understandable ways. Under this frame, we failed. But it was also difficult to pinpoint our exact failures. Thinking about things more, we realized much of our confusion was due to “understanding” being an ill-defined notion.

We now believe that judging interpretability research using the frame of “understanding” can be a trap. Fundamentally, “understanding” is a very subjective term, and an explanation that is understandable for one person may not be for another (consider the OpenMind researcher). We believe that for AI safety research, “understanding” should not be a terminal goal. Rather, “understanding” should be an instrumental tool used to achieve other ends [LW · GW].

It’s also easy to get caught up in creating a principled notion of understanding. But as long as understanding is regarded as the end goal, researchers can easily get trapped in using methods that look nice to humans over being useful, resulting in “streetlight” interpretability and cherrypicked [? · GW] results. This might mean techniques won’t scale to bigger models — which is the setting we care the most about.

Thus, we believe the right response to the OpenMind researcher is as follows: “Mechanistic explanations are valuable insofar as they help safely develop and deploy AI. You might have a full mechanistic explanation, but it doesn’t help assure me that your AI is aligned or safe.” 

As mechanistic interpretability matures as a field, researchers should target use cases for their work. For example, before undertaking a project, researchers could [AF · GW] set [AF · GW] terminal [LW · GW] goals and update them as work progresses. This would also make it more likely that interpretability breakthroughs achieve safety goals over capabilities advances [AF · GW]. We realize that we have not provided a particularly precise razor for determining what mechanistic interpretability research is worth doing: it is our hope that the field of mechanistic interpretability develops a clearer sense of this as it develops.

Lastly, it is hard to predict exactly what advances will come from more “basic”, less back-chained interpretability. This tension between working on what is “interesting” and working on what is “important” is ultimately up to the individual to reconcile.

Minor takeaways

Finally we also had some important but more minor takeaways.

Takeaway 3: Log-odds is a great metric

We chose to represent probabilities in terms of log-odds and to compare probabilities using log Bayes factors. Log-odds is not a standard-unit of measurement in interpretability work, but we believe that it has unique advantages when it comes to interpretability work on language models.

Probability differences and log-probability differences of output tokens are common metrics for measuring the effect size of activation patching in language models (Zhang et al. 2023). But both have drawbacks:

• Probability differences fail to capture changes near 0 or 1. Consider an intervention that causes a token's probability to go from 99% to 99.9999%. In probability space, these differ by less than 1%, but to achieve this effect, the token's output logit would need to increase by 9.22 (if all other token logits are held constant). This is a significant amount not reflected by the probability difference. An analogous case arises for a prediction going from 1% to 0.00001%.
• Log-probability differences solve this problem near 0 but are still insensitive near 1. For example, the log-probability difference between 99% and 99.9999% is only 0.01 nats.

Being sensitive to the upper and lower tails is important statistically. Moreover, being insensitive can mean failing to capture the mechanistic behavior of an LLM: A token prediction probability changing from 99% to 99.9999% corresponds to a large difference in output logits, which necessarily reflects a large difference in intermediate activations.

Log-odds-ratios (or log-odds) fix this issue. Given a probability $p$, its associated log-odds is given by the quantity $logit\left(p\right)=\log \left(p/\right(1-p\left)\right)$. And given two probabilities $p_1$ and $p_2$ the log Bayes factor needed to transform $p_1$ into $p_2$ is given by $logit\left(p_2\right)-logit\left(p_1\right)$.

These measures have a few nice properties (see Appendix D for full details and derivations):

• Sensitive to values near 0 and 1. The log-odds of 99% is approximately 2 dits and the log-odds of 99.9999% is approximately 6 dits. Thus, to transform 99% to 99.9999%, we need to apply an update with a log Bayes factor of approximately 4 dits. Likewise, the log Bayes factor needed to change 1% to 0.00001% is precisely the negative of the above, or roughly $-4$ dits.
• Natural to use with models that output logits. Consider a model which outputs a vector of logits $x$. Suppose we increase the $i$th coordinate of $x$ by $\Delta$, holding all other logits constant. How much does the probability of the $i$th coordinate change (after softmax)? The answer is by exactly a log Bayes factor of $\Delta$ nats.
• Bayesian interpretation: Let $E$ be an event which may or may not occur (for example, $E$ could be the event that a language model correctly answers a factual recall question). Let's say your prior probability that $E$ will occur is $p$.
  
  Now suppose someone tells you some news $X$, and you know that in worlds where $E$ occurs, the probability you would have gotten this news is $a$, and in worlds where $E$ does not occur, the probability you would have gotten this news is $b$.
  
  If you are a good Bayesian, you would revise your belief with Bayes' rule.  This would result in your belief in the probability of $E$ increasing additively by a log Bayes factor of $\log \frac{a}{b}$.
  
  Moreover, if you receive multiple independent pieces of news $X_1,\dots ,X_n$, which each individually would have resulted in a log Bayes factor update of ${\delta }_1,\dots ,{\delta }_n$, then the net effect of receiving all this news is that your log-odds of $E$ increases by ${\delta }_1+\cdots +{\delta }_n$.
  
  We show in our paper (Appendix D) that our residual stream components act somewhat independently, illustrating that this Bayesian interpretation is partially valid.

Takeaway 4: Tooling is important and should not be taken for granted (or how we discovered our Llamas were concussed)

We used the amazing open-source TransformerLens library to study the Llama-2 model series. TransformerLens converts HuggingFace models into HookedTransformer objects, which allows users to more easily store and manipulate intermediate activations.

However, we discovered late into our project that the conversion procedure for Llama-2 models was subtlety broken. Basically, the HookedTransformer Llama-2 models were concussed: it was not evident from their outputs that their implementation was wrong, as they still behaved coherently, but the output logits of HookedTransformer models differed significantly from the output logits of the base HuggingFace models.

Luckily for us, this issue turned out to not be a big deal, as rerunning with fixed code returned essentially the same results. However, this easily could have not been the case.

Thus one important takeaway for us is to not take tooling reliability for granted. This is particularly true for mechanistic interpretability, a new and relatively small field, where there are not thousands of researchers to stress-test libraries. Our recommendation here is to try and always have comparisons with known baselines (e.g. HuggingFace models). 

1. ^{^}

   See Section 1 of our paper for details.

2. ^{^}

   A particular type of path-patching. See Section 2 of our paper for more details.

3. ^{^}

   This attack also works on llama-2-13b-chat, but we couldn’t find a corresponding one for 70b (although we did not make a concerted effort to do so). We were able to design other similar prompts for 7b and 13b that exploited their messy heuristics. These attacks are quite finicky. For example, they show different behavior when running on our own machines vs. on replicate.com, possibly due to slight differences in prompt construction.

4. ^{^}

   Though for all three of these models, we did not look closely at model internals. In the case of ChatGPT-3.5 and ChatGPT-4, this was because the model was closed-source. In the case of llama-2-70b-chat, this was because TransformerLens had limited support for the model due to its large memory footprint.

8 comments

Comments sorted by top scores.