Compact Proofs of Model Performance via Mechanistic Interpretability

post by LawrenceC (LawChan), rajashree (miraya), Adrià Garriga-alonso (rhaps0dy), Jason Gross (jason-gross) · 2024-06-24T19:27:21.214Z · LW · GW · 3 comments

This is a link post for https://arxiv.org/abs/2406.11779

Contents

  Paper abstract
  Introduction
  Correspondence vs compression
    How to compact a proof
  Proofs on a toy model
  Reasoning about error in compressing the weights
  Our takeaways
  Citation Info
None
3 comments

We recently released a paper on using mechanistic interpretability to generate compact formal guarantees on model performance. In this companion blog post to our paper, we'll summarize the paper and flesh out some of the motivation and inspiration behind our work. 

Paper abstract

In this work, we propose using mechanistic interpretability – techniques for reverse engineering model weights into human-interpretable algorithms – to derive and compactly prove formal guarantees on model performance. We prototype this approach by formally proving lower bounds on the accuracy of 151 small transformers trained on a Max-of-$K$ task. We create 102 different computer-assisted proof strategies and assess their length and tightness of bound on each of our models. Using quantitative metrics, we find that shorter proofs seem to require and provide more mechanistic understanding. Moreover, we find that more faithful mechanistic understanding leads to tighter performance bounds. We confirm these connections by qualitatively examining a subset of our proofs. Finally, we identify compounding structureless noise as a key challenge for using mechanistic interpretability to generate compact proofs on model performance.

Introduction

One hope for interpretability is that as we get AGI, we’ll be able to use increasingly capable automation to accelerate the pace at which we can interpret ever more powerful models. These automatically generated interpretations need to satisfy two criteria:

1. Compression: Explanations compress the particular behavior of interest. Not just so that it fits in our heads, but also so that it generalizes well and is feasible to find and check.
2. Correspondence (or faithfulness): Explanations must accurately reflect the actual model mechanisms we aim to explain, allowing us to confidently constrain our models for guarantees or other practical applications.

Progress happens best when there are clear and unambiguous targets and quantitative metrics. For correspondence, the field has developed increasingly targeted metrics for measuring performance: ablations, patching, and causal scrubbing. In our paper, we use mathematical proof to ensure correspondence, and present proof length as the first quantitative measure of explanation compression that is theoretically grounded, objective, and avoids trivial Goodharting.

We see our core contributions in the paper as:

1. We push informal mechanistic interpretability arguments all the way to proofs of generalization bounds on toy transformers trained on the Max-of-$K$ task. This is a first step in getting formal guarantees about global properties of specific models, which is the approach of post-hoc mechanistic interpretability.
2. We introduce compactness of proof as a metric on explanation compression. We find that compactifying proofs requires deeper understanding of model behavior, and more compact proofs of the same bound necessarily encode more understanding of the model.
3. It is a common intuition that “proofs are hard for neural networks”, and we flesh this intuition out as the problem of efficiently reasoning about structureless noise, which is an artifact of explanations being lossy approximations of the model’s learned weights.

While we believe that the proofs themselves (and in particular our proof which achieves a length that is linear in the number of model parameters for the parts of the model we understand adequately) may be of particular interest to those interested in guarantees, we believe that the insights about explanation compression from this methodology and our results are applicable more broadly to the field of mechanistic interpretability.

Correspondence vs compression

Consider two extremal proof strategies of minimal compression with maximal correspondence and maximal compression with minimal correspondence.

1. Brute force proof: We can run our model on all possible inputs we care about and use the computational trace of the model as our proof. While comprehensive, this proof is infeasible for all but the most toy cases and inadequate for tasks like mechanistic anomaly detection.
2. Trivial proof: On the other hand, we can write proofs that make no reference to our model and only require looking at the property we care about explaining. For example, we can trivially prove that any model's accuracy is between 0% and 100%.

The central obstacle to good explanations is the trade off between making arguments compact, and obtaining sufficiently tight bounds. Broadly, being able to compress behavior shows that we have understanding. The thesis of our paper is that constructing compact proofs with good bounds requires and implies mechanistic understanding.

How to compact a proof

We can think of the brute force proof as treating every input to the model as a distinct case, where exhaustive case analysis gives us the final bound. Constructing a shorter argument requires reasoning over fewer distinct cases.

The naïve way to get shorter arguments is to throw away large parts of the input distribution (for example, only focusing on a distribution made of contrast pairs). In order to avoid throwing away parts of the input distribution, we abstract away details of the specific model implementation from our argument, and then group the inputs and construct a cheap-to-compute proxy for the model’s behavior on all data points in each group.

In order to ensure that the argument is a valid formal proof and that our proxies are valid, we bound the model’s worst-case behavior on each group when using the proxies. If our relaxations lose critical details then our bounds become much less tight, in contrast, if our relaxations abstract away only irrelevant details, then we might be able to still have tight bounds. Mechanistic understanding of model internals allows us to choose better groups and proxies, which in turn leads to shorter proofs with tighter bounds.

Proofs on a toy model

In our paper, we prototype the proofs approach on small transformers trained on a Max-of-$K$ task. We construct three classes of proof strategies, which each use an increasing amount of model understanding and correspondingly are cheaper to check.

The first is the brute force proof, which uses no mechanistic understanding and treats each possible input as its own case.

The second class of proof strategies breaks models into paths (i.e. circuits), and groups the inputs by their max token, the largest non-max token, and the final query token. Besides the path decomposition, these strategies use both the fact that our models generate the correct output by paying attention to the max token in a sequence. The amount of attention paid to a given max token is primarily determined by the identity of the max token together with the largest non-max token and the final query token.

The third class of proof strategies examines each path independently, and groups the inputs for each path by two of the three types of tokens mentioned previously. To get cheaper proxies for model behavior, strategies in this class use knowledge that we might associate with traditional “mech interp”, such as the fact that the QK circuit is approximately rank one, with the principal component measuring the size of the key token.

We then qualitatively and quantitatively examine the connection between mechanistic understanding and proof length/bound tightness:

Table 1: We report the proof complexity, normalized accuracy bound, and estimated flops required (Equation 2), as well as unexplained dimensonality (Section 5). We round the FLOP and unexplained dimension counts to the closest power of 2, and report the mean/standard deviation of the bound averaged across all 151 models. As we include more aspects of the mechanistic interpretation (reflected by a lower number of unexplained dimensions), we get more compact proofs (in terms of both asymptotic complexity and FLOPs), albeit with worse bounds. For space reasons, we use $k:=n_{ctx}$, $d:=d_{model}$, and $v:=d_{vocab}$.

Description of Proof	Complexity Cost	Bound	Est. FLOPs	Unexplained Dimensions
Brute force	$O\left(v^{k+1}kd\right)$	0.9992 ± 0.0015	$2^{40}$	$2^{30}$
Cubic	$O\left(v^3{k}^2\right)$	0.9845 ± 0.0041	$2^{25}$	$2^{14}$
Sub-cubic	$O(v^2\cdot k^2+v^2\cdot d)$	0.832 ± 0.011	$2^{21}$	$2^{13}$
(without mean+diff)		0.758 ± 0.039	$2^{21}$	$2^{13}$
Low-rank QK	$O(v^2{k}^2+v{d}^2$ $O(+\text{(EU\&OV)}{v}^{2}d)$	0.806 ± 0.013	$2^{21}$	$2^{12}$
(SVD only)		0.643 ± 0.044	$2^{22}$	$2^{12}$
Low-rank EU	$O(v^2{k}^2+vd$ $O(+\text{(QK\&OV)}{v}^{2}d)$	0.662 ± 0.061	$2^{21}$	$2^{13}$
(SVD only)		$(3.38\pm 0.06)\times {10}^{-6}$	$2^{21}$	$2^{13}$
Low-rank QK&EU	$O(v^2{k}^2+v{d}^2$ $O(+\text{(OV)}{v}^{2}d)$	0.627 ± 0.060	$2^{21}$	$2^{13}$
(SVD only)		$(3.38\pm 0.06)\times {10}^{-6}$	$2^{22}$	$2^{13}$
Quadratic QK	$O(v^2{k}^2+vd$$O(+\text{(EU\&OV)}{v}^{2}d)$	0.407 ± 0.032	$2^{21}$	$2^{12}$
Quadratic QK&EU	$O(v^2{k}^2+vd$$O(+\text{(OV)}{v}^{2}d)$	0.303 ± 0.036	$2^{21}$	$2^{13}$

Reasoning about error in compressing the weights

As we impose tighter demands on proof length, there is a steep drop-off in tightness of bound. This is perhaps even more apparent visually:

The fundamental issue seems to be that our compression is lossy and worst-case error bounds in compression add up quickly.  Going from $O\left(d_{vocab}^2{d}_{model}\right)$ to $O\left(d_{vocab}{d}_{model}\right)$ in the skip connection EU costs us about 15%. Going from $O\left(d_{vocab}^2{d}_{model}\right)$ to $O\left(d_{vocab}{d}_{model}\right)$ in the QK attention circuit EQKE costs us about 30%–40% in our best worst-case accuracy bound.

Focusing on the QK circuit, we claim that the EQKE is approximately rank one:

One way to cheaply check this without multiplying out EQKE (which would take $O\left(d_{vocab}^2{d}_{model}\right)$  time) is to factor out the principal components of each of the component matrices. Although the residuals for each of the four component matrices (after removing the first two principal components) are both small and seem to be noise, proving that there's no structure that causes the noise to interact constructively when we multiply the matrices and “blow up” is hard. 

In general, it's often hard to prove an absence of structure, and we believe that this will be a key limitation on scaling proofs, even with a high degree of mechanistic understanding. 

See our paper for more details.

Our takeaways

We think that there are two main takeaways from our work:

1. Mechanistic interpretability as compression: we think that the proofs frame highlights the fact that the purpose of mech interp is not to explain individual model components in isolation, but instead to compress the behavior of the entire model across a full data distribution, either to achieve useful safety guarantees directly, or in ways that are useful for downstream tasks^[1] such as mechanistic anomaly detection, scalable oversight, or adversarial training.
2. Structureless noise as a serious challenge to scaling proofs: For any behavior, asserting that the unexplained parts of the model are always unimportant will likely be infeasibly expensive. Since it seems likely that we'll need to relax our guarantees from worst case bounds to average case bounds, we're excited about ARC Theory's work on heuristic arguments. We're also excited to see work on other approaches to relaxing the guarantees, such as getting probabilistic bounds on unstructured noise via clever sampling methods. Finally, it's possible that if we can explain the central behavior of a model on all parts of the input distribution, then we can directly finetune the model against our explanation to remove unstructured noise – we'd also be excited for work in this vein.

---

The work on this project was done by @Jason Gross [AF · GW], @rajashree [AF · GW], @Thomas Kwa [AF · GW], @Euan Ong [AF · GW],  Chun Hei Yip, Alex Gibson, Soufiane Noubir, and @LawrenceC [AF · GW]. The paper write-up was done by @LawrenceC [AF · GW], @rajashree [AF · GW], and @Jason Gross [AF · GW]; @Adrià Garriga-alonso [AF · GW] assisted in writing up this blog post.

Citation Info

To reference this work, please cite our paper:

@misc{gross2024compact,
  author      = {Jason Gross and Rajashree Agrawal and Thomas Kwa and Euan Ong and Chun Hei Yip and Alex Gibson and Soufiane Noubir and Lawrence Chan},
  title       = {Compact Proofs of Model Performance via Mechanistic Interpretability},
  year        = {2024},
  month       = {June},
  doi         = {10.48550/arxiv.2406.11779},
  eprint      = {2406.11779},
  url         = {https://arxiv.org/abs/2406.11779},
  eprinttype  = {arXiv},
}

---

1. ^{^}

   We think that our work serves as an example of how to leverage a downstream task to pick a metric for evaluating mechanistic interpretations. Specifically, formally proving that an explanation captures why the model has a particular behavior can be thought of as a pessimal ablation of the parts the explanation claims are unimportant.^[2] That is, if we can replace the unimportant parts of the model with their worst possible values (relative to our performance metric) while maintaining performance, this provides a proof that our model implements the same behavior as in our explanation.

2. ^{^}

   Compare to the zero, mean, or resample ablation, where we replace the unimportant parts of the model with zeros, their mean values, or randomly sampled values from other data points.

3 comments

Comments sorted by top scores.