Validating Mechanistic Interpretations: An Axiomatic Approach

Thank you for your helpful comments, we respond in detail below.

Q1. More experiments are required to evaluate the framework by analyzing practical foundation models using the axioms from B.1

We sketch how the circuit for IOI in GPT-2 [1] may be expressed and analyzed in our framework in the response to Q1 of reviewer LJQi. Evaluating circuits in our framework is straightforward given a fully-specified mechanistic interpretation. However, in practice, we observe that popular mechanistic interpretations such as the IOI circuit leave many key details unspecified. Circuit analyses of this type can be strengthened by evaluating using our framework as it forces the analysts to spell out all the details; we hope that such studies will be conducted in future work.

[1] Wang et al. Interpretability in the Wild: A Circuit for Indirect Object Identification in GPT-2 Small. 2022.

Q2. Framework should be evaluated in comparison to baseline evaluation frameworks

We agree that evaluation of this type would be valuable, however, we note that we are the first to propose an axiomatic framework for the evaluation of mechanistic interpretations, presenting a natural and principled set of criteria; we believe that this already constitutes a valuable contribution to the community. We hope that our work inspires further conversation in this direction, and leave such an analysis to future work.

Q3. How do the axioms handle residual connections?

We can express residual connections via concrete functions which return the residual stream as part of a tuple of values. See the analysis of the 2-SAT model for an example: the second concrete component returns both the residual stream and the hidden layer of the MLP. Appendix B.1 discusses a variation of Axioms 1 through 4 which operate on nonlinear graphs, and in this case no change is necessary.

In general, we may treat the residual connection as computing the identity function in the abstract model; we discuss this in more detail in the response to Q1 for reviewer LJQi.

Q4. Are the axioms for some fixed $\epsilon$?

Yes, that is correct. In definition 3.1, we define an interpretation as $\epsilon$-accurate when Axioms 1 through 4 are satisfied with that fixed value of $\epsilon$; we will clarify this.

Q5. Axioms in high-dimensional vector spaces

We do state that our axioms require that the abstract states be discrete (lines 128-130), however we will improve clarity. No such restriction applies to the concrete model. While lines 410-420 illustrate an important tradeoff of discretization techniques, this is not a weakness of our axioms. Noting that our equivalence class formulation for discretization generalizes comparison up to a tolerance, any evaluation approach which considers equivalence of intermediate values (essential to have any hope at resolving the extensional equivalence problems of the type discussed in [2]) will observe the same behavior.

Also see response to Q2 of reviewer zQST.

[2] Scheurer et al. Practical pitfalls of causal scrubbing. AI Alignment Forum. 2023.

Q6. Why do the axioms capture explainability better than other approaches?

Firstly, our axioms define clear standards to evaluate. Approaches such as causal abstraction are overly broad, making it unclear to the analyst how their interpretation should be evaluated. If we view the equalities in our axioms as tests under causal interventions, we propose a specific and natural set of interventions which ensure that both internal and output behaviors of the interpretation and the concrete model match, and moreover that the corresponding abstract and concrete components are interchangeable.

Secondly, we were motivated to formulate our axioms in a manner which captures validation techniques already used informally in the mechanistic interpretation community to evaluate interpretations.

Finally, our axioms improve upon existing approaches by considering all of internal equivalence, output equivalence, and compositionality. We discuss the importance of compositional evaluation in appendices D and E. As shown by [2], evaluating only on outputs does not ensure that the intermediate values agree.

We discuss these differences in more detail in section 3. We should note that we do not claim that our axioms cannot be improved upon; rather, our work represents a first step at formalizing validation of mechanistic interpretations. We hope that work in this direction continues and that our work inspires further improvements to evaluation techniques.

Regarding the specific approach mentioned, as described above, techniques which neglect the equivalence of internal values invite the problem of only establishing extensional equivalence [2].

Q7. Why the non-homogeneous architecture?

Our goal was to identify the simplest architecture with high accuracy, which we observed with this architecture.

Q8. What is the loss function?

We use a next-token prediction objective.

Thank you for your thoughtful comments, we respond in detail below.

Q1. Case studies are too simple; it is unclear whether the approach is applicable to larger models

We note that our framework is already compatible with larger models and more complex architectures, and we emphasize that our main contribution is our evaluation framework (i.e., axioms) for validating mechanistic interpretations and not techniques to derive mechanistic interpretations. Hence, any increased difficulty of deriving mechanistic interpretations on larger models does not affect the applicability of our approach. See our response to Q1 by reviewer LJQi for more details.

Q2. How does the proposed axiomatic approach scale to larger, real-world models, particularly when intermediate states are high-dimensional continuous vectors rather than discrete symbols?

As noted above, our approach is compatible with larger models. We note that the intermediate states of the concrete model are, in all cases (including for 2-SAT and modular arithmetic models), high dimensional continuous-valued vectors. Our requirement that the intermediate states be discrete is only for the abstract model. In particular, we require this as we agree that mechanistic interpretations cannot hope to achieve equality of high-dimensional continuous vectors in Axioms 1 and 2.

However, we emphasize that our requirement is natural for useful mechanistic interpretations, in particular, as any interpretations which operate on high dimensional continuous vectors will not generally be human-interpretable and hence will be of limited utility.

Q3. Essential references not discussed

Thank you for the references, we will incorporate them into the paper. We should note that while [1] proposes an axiomatic framework for interpretability, it considers input interpretability only. [2] is primarily a survey of the architectural components interpreted and techniques such as probing for identifying the representations of abstract values, and reflects the class of concrete programs which $\lambda_T$ must represent.

[1] Sundararajan, Mukund, Ankur Taly, and Qiqi Yan. "Axiomatic attribution for deep networks." International conference on machine learning. PMLR, 2017. [2] Mueller, Aaron, et al. "The quest for the right mediator: A history, survey, and theoretical grounding of causal interpretability." arXiv preprint arXiv:2408.01416 (2024).

Thank you for your detailed comments, and for your strong support for our work! We respond in detail below.

Q1. What happens when only some of the axioms are adhered to?

In appendices D and E, we include a discussion of what happens when we consider component axioms (Axioms 2 and 4) alone; in this case, we do not consider compositionality of the interpretation. Likewise, considering replaceability alone invites the problem of the interpretation being extensionally equivalent (equivalent in terms of input-output behavior) to the model under analysis but lacking intensional equivalence (equivalence in terms of internal structure) [1].

[1] Scheurer et al. Practical pitfalls of causal scrubbing. AI Alignment Forum. 2023.

Q2. How should the $\epsilon$ values be contextualized?

We may view the $\epsilon$ as a form of error, in the sense that $1 - \epsilon$ is the accuracy with which the abstract model or intervened concrete model predicts the behavior of the original concrete model. In our experiments, we observe values of $\epsilon$ in the full range of 0 to 1. We observe values of up to 1 for bad interpretations, as well as values of 0 for interpretations which perfectly match the behavior of the model (the only reason the paper does not report any $\epsilon$ values of 0 is because we compute Clopper-Pearson confidence intervals). In this sense, $\epsilon$ describes the probability that the behavior of the concrete and abstract models differ, and is meaningful in itself. We should add that $\epsilon$ is, in addition, useful for relative comparisons, identifying which interpretation better matches the underlying behavior of the model.

Q3. Relationship of our framework with causal abstraction and causal interventions

We will make a note about Remark 35 of [2], thank you for the suggestion! However, it is important to note that $\tau^{-1}$ may not always be a feasible choice for concretization. In particular, $\tau$ will not be invertible in general, and hence this necessitates application of set semantics, which may be infeasible in the concrete domain, and hence for replaceability. In addition, we should note that it is not necessarily the case that abstract interpretations align individual high level with individual low level variables.

While our axioms do not exactly follow the structure of interchange intervention (in particular, the interventions are not derived by substituting intermediate states from different inputs), they may be regarded as performing a specific class of causal interventions, which are soft interventions of the type described in [2].

The replaceability axioms may be viewed as interventions on the original concrete model. If we consider component replaceability and $x_i$ is the input to the $i^{\text{th}}$ concrete component, the intervention performed in component replaceability replaces the constraint $x_{i+1} = d_t$ i $(x_i)$ with the constraint $x_{i+1} = \gamma_i \circ d_h$ i $\circ \alpha_{i-1} x_i$. The equivalence axioms may be defined as analogous soft interventions on the corresponding abstracted prefixes of the concrete model.

From that perspective, our axioms present a principled choice for a standardized set of causal interventions, which we believe is very valuable to the mechanistic interpretability community. In particular, while frameworks such as that of [2] are very broad, that breadth means that there is no clear standard choice for evaluation, hence causal abstraction analyses continue to emphasize interchange intervention accuracy, which does not directly evaluate the equivalence of internal representations.

[2] Geiger et al. Causal Abstraction: A Theoretical Foundation for Mechanistic Interpretability. 2024.

Q4. The framework is evaluated only on toy models.

We do not intend the case studies presented to serve as an evaluation of our framework, but as an illustration of how it may be applied. Our framework can indeed be used to validate interpretations of larger models; see the response to Q1 of reviewer LJQi for more details. As the reviewer notes, our primary contribution is the new framework for evaluation of mechanistic interpretations and not particular techniques to derive them.

Thank you for your thoughtful comments, we respond in detail below.

Q1. It is unclear how well this approach generalizes to larger models

We emphasize that our key contribution is our framework (i.e., axioms) for evaluation of mechanistic interpretations and not particular techniques to derive them. Hence, while deriving/finding mechanistic interpretations is more difficult on larger models, it does not impact the applicability of our approach. As an illustration, consider the application of our framework to IOI [1].

The IOI circuit for GPT-2, at a high level, identifies the indirect object by identifying all names in a sentence, filtering out any duplicated names, and selecting the remaining name.

The analysis in [1]. is at the level of attention heads, and consists of five broad categories of heads:

1. Previous token heads, which copy information from the prior token
2. Duplicate token heads, which identify whether there exists any prior duplicate of the current token
3. Induction heads, which serve the same function as duplicate token heads, mediated via the previous token heads
4. S-inhibition heads, which output a signal suppressing attention by name mover heads to duplicated names
5. Name mover heads, which copy names except those suppressed by the S-inhibition heads for output

Each of these attention heads computes a well-defined human-interpretable function and are representable in our framework; each may be modeled as an independent node in the computational graph. Each depends on the residual stream flowing into the corresponding layer; the subsequent residual stream is then a function of the preceding residual stream and all interpreted attention heads. As only attention heads are interpreted, the remainder of a block is an identity function in the abstract model. In this way, we can construct isomorphic concrete and abstract computational graphs for arbitrary circuits.

Appendix B.1 describes how the resulting computational graphs may be linearized, allowing the application of axioms 1 through 4, and how to extend our axioms to operate natively on nonlinear computational graphs.

Moreover, our framework provides a principled approach for evaluating a number of key questions left by [1] for future work. In particular, the authors of [1] hypothesize that S-inhibition heads output relative, and not absolute positions of duplicated names, but do not clearly demonstrate that this is the case. Applying our axioms to the corresponding abstract models would permit the analyst to identify the better hypothesis.

However, we again emphasize that our focus is on the evaluation framework, and not methods for deriving mechanistic interpretations. Thus, while the analysis of GPT-2 Small and other larger models is compatible with our axioms, conducting such an analysis and deriving mechanistic interpretations is out of scope for our work.

[1] Wang et al. Interpretability in the Wild: A Circuit for Indirect Object Identification in GPT-2 Small. 2022.

Q2. Significant manual effort may be necessary

We agree that significant effort is necessary for the analyst to derive a mechanistic interpretation, but it should be noted that mechanistic interpretability is in general a time-consuming and highly manual process. While automated approaches to deriving interpretations are an interesting direction of research in their own right, our work is about evaluation of derived mechanistic interpretations and not about techniques to derive an interpretation.

Q3. How sensitive are the axioms to the choice of abstraction and concretization functions?

The axioms may be sensitive to this choice, but we should note that these are essential components of the interpretation; interpretations which claim that the same abstract function is computed over different representations are not the same, and hence their evaluations should not be either. It is the analyst's responsibility to choose these functions appropriately, and our axioms are useful for identifying which choice more closely matches the behavior of the model.

Q4. Are the proposed axioms sufficient to validate mechanistic interpretations?

Yes, our axioms are sufficient, presenting a strong set of criteria to ensure that the behavior of the concrete model and the claimed mechanistic interpretation are interchangeable and behave in the same way. We recognize that stronger axioms are possible, and, in particular, we propose two additional axioms, Axiom 5 and 6 in the appendix. These strengthen the evaluation, but are difficult to operationalize, and we believe that Axioms 1 through 4, as presented in the main body, already constitute an effective evaluation valuable to the mechanistic interpretability community. Our work represents a first step in the direction of formally validating mechanistic interpretations. We hope that our work inspires a conversation in the community and leads to further improvements to evaluation techniques.