Compact Proofs of Model Performance via Mechanistic Interpretability

Thank you for your feedback! We have also rectified the typographical issues.

You only compare sequence lengths of up to four because the brute-force solution is infeasible beyond that. Could your proposed proof strategies be used to make guarantees beyond this length? I believe that demonstrating that your approach allows us to provide guarantees for settings where brute-force solutions are (nearly) impossible to derive would greatly strengthen the results.

Thanks for the suggestion, we agree that this is a serious concern. We have generated new results on Max-of-5, Max-of-10, and Max-of-20 (holding parameters other than n_ctx fixed). We ran the cubic proof on each of these models, giving bounds of 94% (Max-of-5), 93% (Max-of-10), and 83% (Max-of-20) in under two minutes. For comparison, the brute-force approach on Max-of-20 would take roughly $64^{20} \cdot 20 \cdot 64 \cdot 32 \approx 10^{40}$ FLOPs and it took about $10^{23}$ FLOPs to train GPT-3. We will add a section on the scalability of proof strategies to larger models, which includes these experiments as validation that our strategies result in lower complexity proofs.

The results appear to be merely an interesting case study in a highly simplified setting. The proposed proof strategies rely on various unrealistic simplifications, including the use of attention-only transformers, and the learned algorithm for the task allows for further simplifications, such as ignoring positional embeddings. These simplifications raise doubts about the transferability of the proposed proof strategies to other settings. For example, the observation that compact proofs, which rely on various simplifying assumptions, are already less faithful in this synthetic setting suggests that obtaining compact proofs in more realistic settings would be incredibly challenging.

Our goal in the project was to prototype verification of transformer models, lay out desiderata for verification (compactness of proof length for feasible cost, and tightness of bound for guarantees to be materially useful), and study what the key challenges to achieving the desiderata (faithfulness of interpretation, compounding error terms). Obstacles that arise in prototyping are likely to be even more challenging in complex, real-world settings. This approach is standard in the field of mechanistic interpretability: starting in simplified settings and pushing for extensive exploration [1]. Simplifications of the model architecture in our work are comparable to state of the art analysis of toy models [2]. And the challenges of decomposing even small models adequately are discussed in various papers in the field [3] [4].

However, we agree that we could have taken the alternative approach of targeting the verification of a large model, which would have been an impressive result.

The authors employ a variety of methods to analyse the model and cleverly use mechanistic insights to reduce proof length. To the best of my knowledge, this is the first paper that attempts to verify global model behaviour using insights from reverse-engineering the neural network.

We appreciate the highlights about our contributions! We want to add that to the best of our knowledge, we are the first to (1) verify the performance of a transformer-architecture model, (2) guarantee global performance bounds on specific models, and (3) incorporate mechanistic information to obtain shorter proofs. While we work in a simplified setting, we are excited about follow-up work building upon these milestones.

The proposed measure to assess mechanistic understanding (degrees of freedom that have to be evaluated using brute-force) is interesting, as current evaluations of interpretability results often rely on proxy metrics.

We appreciate this highlight as well. Whereas current work relies upon human judgment, with ambitious projects pushing towards procedural arguments, by pushing mechanistic interpretability to formal proofs, we were able to find that degrees of freedom remain remain even when we have decent informal interpretations.

Please don't hesitate to ask any further questions. We would love the chance to clarify our approach and provide details that might encourage you to increase your score!

[1]: Elhage, et al., "Toy Models of Superposition", Transformer Circuits Thread, 2022.

[2]: Elhage, et al., "A Mathematical Framework for Transformer Circuits", Transformer Circuits Thread, 2021.

[3]: Stander, D., Yu, Q., Fan, H. and Biderman, S., 2023. Grokking Group Multiplication with Cosets. arXiv preprint arXiv:2312.06581.

[4]: Chan, et al., "Causal Scrubbing: a method for rigorously testing interpretability hypotheses", AI Alignment Forum, 2022.

How would this method generalize to different models that implement different circuits or algorithms on the same task?

We agree that guarantees require taking into account the specific model – this is part of the mechanistic understanding that interp should provide. That being said, we believe that this specificity can be useful: for example, if a proof strategy derived for one model does not apply to a second model, this provides evidence that the two models are using different algorithms.

In line 73, why is pessimal ablation done $...with samples$ ?

We apologize for the typo. In line 73, “with values drawn from $\mathcal{D}$” should instead read “to values over $X$”. That is, instead of sampling, which might not find the actual worst-case, we use the actual global worst-case values as the counterfactuals for ablation. The counterfactuals commonly constructed in mechanistic interpretability are attempts to estimate the typical (/average-case) behavior of a component, which is insufficient for getting a worst-case guarantee on the overall behavior of the model. We’ll edit the paper to fix the typo and further clarify this point.

$T$ he authors denote the residual from the rank-1 approximation of the EQKE matrix as ''unstructured noise.'' $...$

We apologize for our lack of clarity. First, by “unstructured”, we mean that insofar as structure exists, it doesn’t help the model’s performance. Second, “If we sample elements randomly” should read “If we replace the entries of $E_{q,2}^\perp$, $E_{k,2}^\perp$, $Q^\perp$, and $K^\perp$ with randomly sampled values”. In order to validate our claim that the noise contains no helpful structure, we performed this replacement 100 times and compared the max-row-diff of these samples with the diff in the original model (lines 1074-7), and found that the diff of the original model is worse than the mean value by about 4σ. (This is analogous to the standard resampling or mean ablations that mech interp work uses to check for lack of structure.)

Why is the ratio of the largest to second-largest singular value a good metric for the faithfulness of the interpretation?

Our interpretation claims that $EQKE$ is approximately rank 1 (the size direction). If the first singular value is much larger than all other singular values, this suggests that the matrix is well-approximated by a rank-1 approximation. We’ll clarify this in the text.

Responses to more minor points:

Line 55, ''we restrict $f$ to be the 0-1 loss so our theorems bound the accuracy of the model.'' In Eq. 1, we seek to lowerbound $f$.

We had this terminology wrong, and indeed were intending to speak of 1 minus the 0-1 loss. We’ll correct this, and thank you for pointing it out.

Line 87, Line 106, inconsistent bold-facing. What is difference between $**E**$ and $E$. Moreover, what does $\mathbf{E}_q$ mean?

The boldfacing is significant and is defined on line 107:

$\hat{\mathbf{P}} = P - \overline{\mathbf{P}} \quad \text{and} \quad \mathbf{x}\overline{\mathbf{E}} = \mathbf{x}E + \overline{\mathbf{P}} \quad \text{and} \quad \mathbf{x}\mathbf{E}_q = \mathbf{x}E + \mathbf{P}_q \quad (\text{since } h^{(0)} = \mathbf{x}\overline{\mathbf{E}} + \hat{\mathbf{P}}).$

That is, we use the boldfaced versions of $\mathbf{E}$ to denote $E$ summed together with positional embeddings. We’ll try to clarify this more in the text.

The authors claim that Fig. 4 shows how less faithful interpretations lead to tighter bounds. It is extremely unclear to me that there is a positive correlation between the Singular ratio and accuracy bound. The authors should perform a statistical test here based on the $R^2$ to demonstrate that there is indeed a significant correlation.

We deeply apologize for not catching the bug in our code that led to an incorrect version of this figure being included in the original submission. We’ve included a corrected version of this figure in the accompanying pdf. From this figure, it is apparent that for the smallest values of $\sigma_1/\sigma_2$ (the three values below 350), this ratio is strongly correlated with proof bound. For the “svd” proof strategy, where $\sqrt{2}\sigma_1(\mathrm{EQKE\_err})$ is used to bound the error term, a linear regression gives an $R^2$ of about 0.71.

Please don't hesitate to ask any further questions! We appreciate the thoughtful engagement, and would love the chance to clarify results that might encourage you to increase your score.

Thank you very much for the extremely detailed review! We also appreciate the list of typographic errors, which we’ve corrected in a revised draft of the paper.

We’d like to start by defining what we mean by “compact proof of global robustness” by clarifying our definition on lines 48-69 of Section 2.

• By “compact proof”, we’re referring to a proof that is short. Specifically, in our paper, we measure the length of proofs by considering the length of a trace of the computational component of the proof (as described in lines 56-61). We agree that our use of the concept “compactness” could be made more explicit in the paper; we’ll add additional clarification. As for why we want compact proofs: proofs that are too long are both infeasible to find and also infeasible to verify – hence our interest in using mechanistic understanding to generate shorter/more compact proofs.
• By “global robustness”, we mean that the model performs correctly in expectation over an input distribution $\mathcal D$ (lines 48-52). This is in contrast to much of the prior work, which focuses on establishing correctness of model behavior for points close to a particular input $x$ (which we refer to as local robustness). We agree that the use of the word “global robustness” in the introduction is confusing, because we use “global performance” to refer to the desired form of our verification target in the title, on line 49 in Section 2, and on line 294 and 301 in Section 6. We’ll edit the paper to be more clear and consistent.

Based on Eq. 1, it does not seem that these guarantees fend against out-of-distribution or adversarial examples?

In addition, we’d like to clarify the flexibility inherent in equation 1. Normally, by OOD or adversarial inputs, we suppose that there is a distribution $\mathcal D_{\textrm{in}}$ that’s used for training and (in-distribution) validation, and another distribution $\mathcal D’$ that is the deployment distribution or generated by an adversary. If we had knowledge of $\mathcal D’$, we could compute the expected performance from inputs sampled from $\mathcal D’$. Even if we don’t have exact knowledge of $\mathcal D’$, we can still define a very broad distribution $\mathcal D$ that covers possible $\mathcal D’$s.

In our work, $\mathcal D$ is the distribution of all $64^4$ possible valid input sequences. In addition, as our proofs partition $\mathcal D$ into subdistributions, and bound the performance on each subdistribution, we can bound the model’s performance on any possible distribution over valid input sequences.

We agree that we don’t adequately discuss the flexibility inherent in our definition; we’ll add clarification to sec 2 as a well as a section to the appendix discussing this in more detail.

We believe that much of the motivation for the work (and indeed the answers to your questions) follow from the goal of producing compact proofs.

Why did the authors choose to investigate a task like max-of-$k$ instead of modular addition (or IOI, colored-objects, etc.) $...$ ?

We picked a very simple task because formal verification is quite challenging, even for Max-of-$k$, and believe that Max-of-$k$ captures some of the core difficulties that we might expect to see. But to address some of your specific examples:

1. Modular addition: In Nanda et al 2023, the authors don’t provide a mechanistic understanding of how the MLP layer performs multiplication – instead, they check this via brute force. This lack of understanding prevents us from compactifying our proof of the model’s correctness.

   In follow-up work, we extend our methodology to the modular addition task. However, compactifying the proof required us to reverse engineer and provide a mechanistic description of how the MLP computes its output – which is, as far as we know, not previously known in the relevant literature.

2. IOI/colored-objects/other small circuits in LMs: One of the concerns with using these small circuits is that they tend to only be valid on a very restricted subdistribution. For example, the original IOI paper considers a small set of sentences with names ordered ABBA, and the circuit metrics become invalid on as small a distributional shift as ordering the names ABAB[1]. If we chose to study these circuits, we’d run into a dilemma: either we restrict our distribution to be small enough to check feasibly with brute force (as the authors tend to do), or we’d have to perform additional interp that advances SOTA. As our focus is on the formal verification side, and not on the novel interpretability side, we chose instead to pick a task where the mech interp is relatively straightforward.

Given that these mechanistic interpretability techniques require models with almost perfect performance a priori why is this not tautological with respect to their formal verification?

We agree that on models where we can evaluate the behavior on all possible input-output pairs, we can indeed guarantee performance via the brute-force proof (lines 151-4). We study such a case because it allows us to have “ground truth” on the performance of the model.

In general, however, mech interp is often studied with samples from a distribution – for example, in IOI, the mean ablation is computed with respect to a small handful of other sequences, and in modular addition the authors use sampling to evaluate the 5-digit addition, repeated subsequence, and skip trigram tasks. In cases where the bad behavior is heavy tailed, sampling may not provide a viable estimate of performance – our hope is that understanding derived from studying the model with samples will allow us to construct formal guarantees that hold in the worst case. We prototype this problem by providing formal guarantees on a max-of-20 transformer, where brute forcing the input distribution is infeasible and our understanding is derived via samples.

[1]: Miller et al. Transformer Circuit Faithfulness Metrics are not Robust.

Thank you for your time, effort, and engagement.

Why is it the case that the role of every component in the network needs to be known for its proof to be compactified?

To be more precise: the role of every component must be known for the proof to be asymptotically compactified.

When measuring asymptotic behavior, it is standard to consider the leading order term(s) only. If the number of layers and attention heads is fixed, the number of parameters in a component is a constant fraction of (that is, linear in) the number of parameters of the overall network. Hence if we have to use the brute-force proof on any component, we will have an asymptotic cost equal to that for performing a brute-force analysis of the entire network.

As we are not able to compress the proof of a component's behavior in the absence of mechanistic understanding of the role and functioning of that component, the leading contribution to the asymptotic size will always be the components we don't understand at all.

Moreover, we want to clarify that compressing components using mechanistic understanding is how we obtain compact proofs with non-vacuous bounds at all, even if approximation errors from the compression can eventually tank the hoped-for bound. Without integrating mechanistic understanding, we are not able to compactify the brute-force proof.

Where are the results for these experiments?

Sorry, “attached” was left-over from an earlier draft of the response, and we meant to say “included”: our results were included in the sentence “the cubic proof achieves bounds of 94% (Max-of-5), 93% (Max-of-10) and 83% (Max-of-20) in under two minutes”.

However, we can provide some additional results for Max-of-10: The equivalent graph to Fig. 3 for Max-of-10 has the following points:

• importance sampling (acc: 0.9988 ± 0.0013), $1.6 \cdot 2^{85}$ FLOPs for brute force
• cubic (rel acc: 0.915 ± 0.021), $1.4 \cdot 2^{27}$ FLOPs
• subcubic (rel acc: 0.734 ± 0.019) $(1.49 ± 0.10) \cdot 2^{22}$ FLOPs
• attention-$d_{\\mathrm{vocab}}d_{\\mathrm{model}}^2$ (rel acc: 0.718 ± 0.019), $(1.609 ± 0.090) \cdot 2^{22}$ FLOPs
• direct-quadratic (rel acc: 0.657 ± 0.032), $(1.392 ± 0.013) \cdot 2^{22}$ FLOPs
• attention-quadratic (rel acc: 0.437 ± 0.028), $(1.384 ± 0.013) \cdot 2^{22}$ FLOPs
• attention-quadratic, direct-quadratic (rel acc: 0.398 ± 0.028) $(1.318 ± 0.014) \cdot 2^{22}$ FLOPs

The equivalent graph to Fig. 4 has ratios $\sigma_1 / \sigma_2$ starting at around 400, and the best-fit line for the SVD proof approach has $bound ≈ 0.56 + 8.3\cdot 10^{-5} (σ₁/σ₂)$, with $r^2 \approx 0.5$.

Also, as $k$ or $d$ changes how are the authors verifying that the underlying circuit implemented by the model is the same, as this could affect the assumptions of all subsequent theorems/algorithms?

We note that:

1. We are working with proofs as computations (lines 56--59), i.e. the theorem statement doesn't exist independently of the proof, it doesn’t presuppose a lower bound $b$ that we then prove.
2. Whatever bound $b$ we are able to compute is the theorem statement we have proven. Thus, the theorem template is valid regardless of what circuit is implemented by the model.

The faithfulness of the interpretation used in the proof mediates the bound (Section 5.2). In this way, the tightness of bound becomes a metric for how faithful the interpretation is. Thus, the recovered bounds can be seen as a measurement of how well the interpretation for Max-of-4 generalized to the new models trained to compute Max-of-5, Max-of-10, and Max-of-20. We are clarifying this further in our writeup, and this is valuable feedback on what concepts are worth highlighting from verification.

Note that if the EVOU circuit did not perform copying, or if the model did not compute its answer by paying more attention to the correct token, the bound would be vacuous. We want to highlight that in the Fig. 4 equivalent for Max-of-10, the smallest ratio $\sigma_1/\sigma_2$ that we saw was over 400, which is evidence that the interpretation that EQKE is approximately low-rank remains valid. We are making this point more clear in a new section on applicability of results to other models.

Thank you for your review!

The tightness of bounds seem to degrade very quickly as additional insights from mechanistic interpretability are used. Thus it seems difficult if future works would be able to use insights from mechanistic interpretability to derive tight bounds on performance of the model.

We agree that the degrading tightness of bounds is an issue. However, note that we evaluate proofs on the additional criteria of compactness of proof length, and using information about model internals is how we obtain shorter proofs than brute force. Thus, some degree of mechanistic interpretation is necessary for finding compact proofs in future work. We address this question in more detail in the top-level rebuttal.

One claim that we will add more clearly to the paper is that proof length can be used as an evaluation metric on the quality of mechanistic interpretation (see Section 5). The interpretations in our different proofs all appear to be sound and use the same high-level abstractions. However, by pushing them to rigorous formal proofs we get to see the various degrees of freedom that remain in the different interpretations. Therefore, we expect that future work will need mechanistic interpretations to get short proofs, and can use the gains in proof reduction as a way of distinguishing what interpretation techniques to use.

The author analyze a very simple task and techniques used by authors to decrease computational complexity are specifically tailored for this task. It therefore remains unclear how insights from this work can be applied to any other tasks. It would have been great if the authors would have attempted to investigate a few other tasks which are simple (fundamental) but still very different in nature, e.g., counting, sorting, etc.

Verification projects are quite intensive in human-labor, and our goal with this project was to prototype verification for transformer models, and focus on studying the core challenges to its viability. Thus, we limited ourselves in the number of different tasks we worked on, and instead focused on deriving multiple proofs that varied the two core desiderata: compactness and bound tightness.

However, we want to highlight that it should be straightforward (if labor-intensive) to generalize our work to other toy models where the model computes the answer by paying attention to a single token or position, i.e. general element retrieval, example tasks include: first element, last element, binding adjective to key phrase, etc. This is because our theorems in the appendices prove in full generality that any computation performed by a single attention head, when only one of the position and the token value matters, is extremized by pure sequences. This means that the worst-case behavior of the computation at a single token can always be found by setting all other tokens to be identical, regardless of what the task is. Furthermore, generalizing our convexity arguments to multiple attention heads whose error terms can be treated independently without degrading the bound too much is also straightforward.

We address the high-level version of this question in the top-level rebuttal where we talk about the general applicability of our approach to projects on larger input distributions, different tasks, more complex architectures, and larger models.

There isn't enough evidence to adequately support the noise hypothesis presented by the authors i.e. the noise in estimating rank-1 approximations of E, Q, K results in larger magnitude in the error term. Currently this seems based on speculation and needs more detailed investigation. Therefore it would be great if the authors could try to artificially inject noise (of varying magnitude) in these matrices and see its effect. Else, the authors should not claim this being the major cause for not so tight bounds obtained on using insights from mechanistic interpretability.

Thank you for pointing out our confusing use of terminology! By “noise” we meant error terms in our approximations of the model that we use to generate more compact proofs, not error terms that the model learned. So, we are not referring to model weights that we can ablate to improve performance bounds of our proofs. In the improved version of Table 1 (included in the accompanying pdf, and which will replace Table 1 in the paper), consider the difference in bound between “Low-rank QK” (0.806) and “Quadratic QK” (0.407). The only change between these results is in the strategy for bounding the error that comes from approximating QK as low rank. We provide a more detailed clarification in the top-level rebuttal.

Please don't hesitate to ask any further questions! We would love the chance clarify our approach, and provide further details that might encourage you to increase your score.

Thank you for your kind review!

how can this method be carried over to more realistic tasks?

It would be relatively straightforward to carry over the methodology to more realistic tasks if the following conditions were met:

1. A sufficiently detailed mechanistic understanding of the components of the model not being brute-force evaluated.
2. A model that is interpretable enough that the (not-understood) remainder of the model can be treated as independent error terms without blowing up the bound.

If these conditions were met, then the interpretation would broadly proceed as follows: Every expression $x$ in the model with interpretation $y$ can be re-expressed as $x = y + (x - y)$. The error terms $(x - y)$ must then be given worst-case bounds.

As far as we can tell, most realistic models do not meet these conditions. Consequently, in the project, we devote significant effort to getting reasonable bounds with minimal understanding (and it is surprising that, e.g., we can get a non-vacuous bound on the model that, modulo our lack of understanding of low-rank approximations to the identity matrix, is linear in the parameter count of the model!). In follow-up work, we are making progress on applying the approach to models with MLPs, to models with with multiple attention heads, to models with layer norm, and to two-layer models. We see three big obstacles. We believe all the obstacles are important for the field of mech interp to address (with the possible exception of the third one):

1. Unstructured approximation error terms (“noise”) accumulate quickly, necessitating either models with significantly less noise (more interpretable models) or more sophisticated error term analysis. In follow-up work, we have preliminary results suggesting that fine-tuning models can make them sufficiently more interpretable as to bypass this problem.
2. Path decomposition is exponential in the number of layers, necessitating adequate mechanistic understanding of how to simultaneously argue that multiple paths contribute negligibly. Existing approaches to solving instances of this problem in finding mechanistic interpretations -- such as presuming independence of irrelevant paths and evaluating them all simultaneously, or using gradient-based search methods that leave the path decomposition implicit -- unfortunately do not yield valid proof techniques. This problem is especially pressing for us, as it interacts with the central problem of compounding error approximations (1). However, we are currently investigating the possibility of matching prefixes of paths to allow mechanistic understanding to be used to simultaneously argue that multiple paths contribute negligibly.
3. Worst-case error bounds are both significantly looser by default and significantly more expensive to compute than average-case error bounds. We would be excited to see a formalization and analysis of compact guarantees of average-case performance! (Unfortunately, we are not aware of any adequate formalism of average-case performance that is independent of training method.)

If I understand correctly, much of the theory in the appendices relies on a careful analysis of the structure of this particular transformer model and the particular (arithmetic) task in question.

That is correct. However, we would like to highlight that it should be straightforward (if labor-intensive) to generalize our work to other toy models where the model computes the answer by paying attention to a single token or position, i.e. general element retrieval, example tasks include: first element, last element, binding adjective to key phrase, etc. This is because our theorems in the appendices prove in full generality that any computation performed by a single attention head, when only one of the position and the token value matters, is extremized by pure sequences. So the worst-case behavior of the computation at a single token can always be found by setting all other tokens to be identical, regardless of what the task is. Furthermore, generalizing our convexity arguments to multiple attention heads whose error terms can be treated independently without degrading the bound too much is also straightforward.

Generalizing the theory to handle coupled attention heads or more complex uses of attention will require additional theoretical work.

We address the high-level version of this concern in the top-level rebuttal where we talk about the general applicability of our approach to projects on larger input distributions, different tasks, more complex architectures, and larger models.

similarly, it is unclear whether re-introducing transformer components, and especially MLP layers, won't obliterate the accuracy bounds derived here to be practically useless (e.g. below chance levels, which would be 25% here)

In follow-up work, we conduct analysis of a modular addition model with an MLP where we achieve a bound of roughly 80%--90% on a sub-brute-force proof. However, it is indeed the case that integrating all the mechanistic understanding we have of that obliterates the accuracy bound. Ultimately, we hope that fine-tuning will allow bypassing this issue.

Separately, we want to note that 25% would be average case chance levels for argmax (predicting the position of the maximum). Average-case chance levels on max are actually $1/d_{\mathrm{vocab}} = 1/64 = 1.56\%$ for this task, since we are predicting tokens not positions. Furthermore, it's not entirely clear that the baseline for proofs should be average-case bounds rather than worst-case bounds. Indeed we typically see bound performance drop off to floating point error, not average-case chance, when error terms accumulate too quickly.

you mention: [...] I believe a standard reference is [...]

Thank you!

Please don't hesitate to ask any further questions! We'd love the chance to highlight strengths of our submission that might encourage you to further increase your score!