Optimal ablation for interpretability

Summary

This paper presents an alternative approach to perform component ablations for neural network interpretability. Specifically, the authors consider a neural network as a causal graph and propose *optimal ablations* which simulate component removal by setting the value of a node in the computational graph to a constant value. In contrast to alternative task-agnostic ablation methods (e.g. zero and mean ablations), this constant is learned using gradient descent. The authors argue that alternative task-agnostic ablation methods could confuse downstream computations as these specific values might not have been observed in the training data and then demonstrate how optimal ablations can lead to improvements across a range of different application domains: 1. They evaluate single-component ablations on the indirect object identification (IOI) task and find that optimal ablations perform better than other subtask-agnostic methods (although still worse than subtask-specific approaches); 2. They perform circuit discovery on the IOI and Greater-Than tasks using a new method, which they term uniform gradient sampling (UGS), which outperforms existing circuit discovery approaches; 3. They study factual recall, or more specifically localize where factual associations are stored within a language model, and in contrast to previous studies observe that adjacent components often have very different patching losses which suggests that factual associations are likely stored in single layers, as opposed to being spread across multiple layers as previously hypothesized; 4. Finally, the authors apply optimal ablations to the domain of decoding latent predictions. Here, they propose the optimal constant attention (OCA) lens, which sets all following components to their optimal constant value, instead of skipping these by learning a linear map such as in the tuned lens. They find that OCA scores better on KL-divergence and in terms of faithfulness to the model computations.

Strengths

This paper made a strong impression on me as it combines methodological contributions with extensive empirical studies. The authors take a simple but well-motivated idea and then demonstrate how it can lead to improvements across a range of different interpretability application domains. To this end, the authors propose a range of novel methods that leverage this idea, including UGS, and OCA. I believe that these results could have a significant impact on the (mechanistic) interpretability field and should be of interest to various other subfields.

Weaknesses

- Some of the applications lack depth as a result of the number of different methods studied. For example, it would have been interesting to further investigate the observations on factual recall. - Optimal ablations still appear to perform worse than subtask-specific ablation methods in single-component ablations. However, the authors make a reasonable argument that these are subjective and require human and require manual effort, making them hard to use in some settings.

Questions

1. Resample ablations are almost on par with optimal ablations in the single-component ablation setting (see Table 1), but perform significantly worse than mean and optimal ablations when it comes to making latent predictions (see Figure 3). Do you have a hypothesis as to why this might be the case?

Rating

8

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

I believe the limitations are properly addressed in the paper.

Summary

This paper explores a new approach to replacing features in the forward pass of a neural network, for purposes of interpretability. The approach is, rather than zeroing out features or replacing them with some pre-specified constant or random variable, to optimize a replacement constant to minimize the loss of the model over the subtask of interest. So when ablating an attention head or MLP layer, the ablated component’s outputs are replaced with a constant that is optimized to maximize model performance over some data of interest. This is done to be “minimally disruptive to model inference”, e.g. avoid disrupting model inference in the way that OOD replacement values are known to disrupt model inference, as has long been lamented by past work in interpretability. This new feature replacement method is demonstrated in three case studies, focused on circuit analysis, factual association localization, and something to do with the tuned lens that I could not quite understand. Results suggest that this method reveals clearer/sharper/stronger phenomena in circuit analysis and factual association localization, a promising result for interpretability research.

Strengths

- Very important: The paper’s main strength is that it does something new and interesting in what has been a very saturated space for years. For a long time — since at least 2007 by my count — people have wondered how to remove features (or the outputs of some feature extractor / neural model) from a model, in order to do something like estimate that feature’s effect on the model. The debate rages on, now re-enacted in the era of mechanistic interpretability. While I’m still not sure we have a general theory for this, this paper makes a reasonable proposal on how to do this in a task and model agnostic way, with some promising early results for downstream applications. More specifically, avoiding disrupting model computation “too much” has stood out as a goal of feature replacement methods, and this paper proposes a new method for doing so and provides some interesting case studies showing that this may produce more stronger results in different kinds of interpretability analyses. - Important: The paper is ambitious in a way rarely seen in ML papers, with a new method being applied across three case studies. The paper covers a lot of ground quickly. - Important: The paper is well-written and the notation is clear. - Important: I am sympathetic to the view that subtask-specific interchange interventions can be difficult to construct (a view argued for in the paper, that motivates the proposed method). At the very least, they take some manual effort to construct and are subjective. A drop-in automated replacement could be very useful. For what it’s worth, an argument I would add to the paper is that task-specific counterfactuals can sometimes be difficult to interpret as well. Does changing a name from Alice to Carol have *only* the effect of removing task-relevant information? We can imagine some hidden effects appearing with counterfactuals that we try to construct to remove only one variable, especially for more complex tasks. - Important: The results on circuit analysis and localization are pretty promising, in my view. (I have a very hard time interpreting the results regarding the tuned lens — I don’t understand that experiment’s setup or its broader purpose.) - Important: Experiments appear to be conducted with great care. Many small design choices seem good and clever, like the ReLU in Eq5 preventing rewarding the optimization process for finding a value that leads to over-restoration of model predictions during the causal intervention. - Important: Appendix D shows that the optimization process is unlikely to adversarially induce good task performance in the model, and therefore unlikely to produce interpretability artifacts, because the causal intervention does not seem expressive enough to so.

Weaknesses

- Very important: In my opinion, the approach presented cannot be the right one in general, preventing the proposal from providing a generic solution that circumvents the need to design subtask-specific interchange interventions / counterfactuals. There are two reasons for this, one a criticism of the argument in the paper and the other a counterexample. First, the argument in the paper, while being mathematically clear, is not philosophically rigorous. It is said that we want to “simulate the removal of” model components. What does that mean? I could imagine removing a layer from a Transformer by not running it — by virtue of the residual connection between layers, this is equivalent to zero-ablating the outputs of that layer. Yet the authors take issue with zero ablation, so removing a component must mean something else. Then, to remove a component, it is said that we should be “minimally disruptive to model inference”. What does that mean? Evidentally we are disrupting the model inference. What does it mean to not disrupt it too much? There are some other specific claims that are difficult to interpret precisely, such as avoiding an intervention that “disrupts this flow of information” or using an intervention that provides information that is “inconsistent with information the model derives from other vertices.” To be clear, I basically do like these kinds of arguments in ML papers, because we need arguments for what is right and wrong to do stated in plain English, but I am not sure this one is adequate for proving the point that optimal ablation is the right way to go about things. Second, a counterexample can build on an intuition hinted at in the paper that “there may not exist a constant value that perfectly conveys a lack of information.” The problem is that it is not clear that optimal ablation is the answer to an intelligible question. For example, suppose we have a neural model that classifies people into “give loan” and “reject loan” categories, and that we know that a neuron in this model represents a person’s annual income in USD. So, we know what 0 means, and we know what 100,000 means (it’s their income). Now let’s say there’s a component that takes in this value, and we want to know this value’s “importance” to that component (or later model outputs), so we replace this value with the optimal ablation constant c. What value of c would be correct for this? The question doesn’t really make sense, and it’s because it doesn’t make sense to ask how “important” a component is. Importance is not well-defined concept in causation. What would make sense is to ask what setting the income value to 100,000, when it was 60,000, for a particular datapoint $x$ does to the output of that downstream component. I would posit that the important thing about counterfactual patching / interchange intervention is not that it “selects uninformative counterfactual inputs“ but rather that it selects *known* counterfactual inputs, i.e. inputs we know the meaning of. This allows us to say precsely what the causal intervention means, as well as what its effect was. This is the argument in https://proceedings.neurips.cc/paper_files/paper/2021/file/4f5c422f4d49a5a807eda27434231040-Paper.pdf, as far as I can tell. - Important: while results are plenty promising on the first two case studies, optimal ablation seems to hardly improve over marginal resampling based on Table 1, and the improvement is extremely marginal looking at the rank correlation in Fig. 6 where a linear correlation seems inappropriate for the log-log relationships. So why then, does optimal ablation perform far better than counterfactual patching in Fig. 1, when counterfactual patching was just treated as something of a ground-truth in Sec. 2. This apparent over-performance in circuit analysis seems like it could be an interpretability artifact, in spite of Appendix D suggesting overfitting is unlikely. - Important: I really could not tell what the third case study was trying to show, or how it tried to show it. See next point. - Of some importance: The paper suffers from trying to do so much in nine pages. I think it would be difficult for someone without a deep background in interpretability to read this paper and get as much out of it as the authors would hope. Unless the authors feel strongly about the third case study, I would totally remove that and focus on providing more detail to earlier sections of the paper. - Of some importance: It is hard give any credit to the paper for the introduction of the uniform gradient sampling (UGS) method because it is introduced without nearly enough detail, and without basically any motivation, in the main body of the paper. And it is evaluated only in one case study. I can’t tell if the UGS direction is big enough to split off into another paper, but it doesn’t feel like it belongs in this one. - Of some importance: The organization of the paper is a little unorthodox, without a related work or conclusion section. While the introduction does a good job situating the paper in the literature, I would recommend discussing some earlier work on how this entire debate has played out with feature attribution methods before the intensified focus on “mechanistic” interpretability (i.e. methods focused on data space ablations and not only hidden feature ablations), including (1) https://arxiv.org/pdf/1806.10758 (2) https://arxiv.org/pdf/1910.13413 (3) https://arxiv.org/pdf/2106.00786.

Questions

- See questions in the above. - A note on the references: I think it’s extremely non-standard for the references section to be full of references that do not appear in the paper. I would say that this should be fixed. - In the intro, I think causal tracing is described as a noising ablation, when it is a denoising ablation. - (Goldowsky-Dill et al., 2023) recommends → Goldowsky-Dill et al. (2023) recommend - “constant to which to ablate” → constant for ablating? - "confluence of conflicting information may cause model representations to spiral out-of-distribution” — are the hidden states during a forward pass in circuit analysis with optimal ablation in-distribution? - The sentence beginning with “To assess the magnitude component” is not at all clear. - The paper is not in correct submission format, as it is missing line numbers.

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The limitations could discuss some higher level points around feature replacement and what it could be useful, vs. what the paper empirically shows it may be useful for. It could also include some of the caveats about the proposed method from Sec. 2.

Summary

- Introduces a notion of "optimal component ablation" for activation patching methods in mechanistic interpretability. Specifically, they propose an ablation method that involves setting a component’s value to the constant value that minimizes ablation loss over a chosen distribution. - Shows that optimal ablation can improve algorithmic circuit discovery—the identification of sparse subnetworks that recover low loss on interpretable subtasks.

Strengths

- Simple approach to compute 'optimal' ablations that simulate the removal of a given component/vertex. This fares better than over zero, mean, and resample ablation in terms of ablation loss. - The OCA lens is a clever idea to evaluate layerwise representations and an alternative to linear probing / tuned lens.

Weaknesses

- Writing and paper structure can be significantly better (motivation, problem setup, describing the applications in more cohesive manner, results etc). - Justification for why ablation loss is the "right" way to measure "optimality" could be better. Also, given that the proposed method explicitly optimizes for ablation loss minimization, it is unclear why making this the primary evaluation metric in the experiments is insightful. Are there other downstream metrics for localization that also improve as a result of minimizing ablation loss? For example, does it find smaller circuits that induce the same effect on model behavior? Is the second-order effect on unrelated tasks less (as intended) due to optimal ablations? These evaluations would make the overall story significantly stronger. - One of the main contributions (uniform gradient sampling method) is not described well. I would like to see a more step-by-step description of the proposed method. Most of the important details are deferred to the appendix - These applications are niche and specific to mechanistic interpretability subroutines. To me, factual recall (the second application) seems like a special case of the first one, as the goal is again to localize a given task down to a subset of model components. I would be interested to see optimal ablations improve model editing (e.g., https://arxiv.org/abs/2404.11534 show that accurate localization via zero ablations can be directly used for model editing). - Novelty concerns. I am not sure if the contributions in this paper are strong enough right now, especially given related work in this area, e.g., https://arxiv.org/abs/2309.16042 also focus on systematically evaluating activation patching - I like the OCA lens idea, but it's a bit unfair to compare it to tuned lens because OCA lens is still using additional computation via learned ablated layers. I think this experiment requires fair-er baselines that are more of an apples-to-apples comparison (e.g., smaller models with fewer learned layers?)

Questions

See strengths and weaknesses

Rating

4

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

See strengths and weaknesses

Summary

Different intervention techniques aim to "ablate" parts of the representation of the model to infer their causal function, e.g. by adding gaussian noise. the paper suggests to derive a notion of "optimal" ablation. Particularly, instead of zeroing out or replacing the ablated part with its mean, it is proposed to optimize with GD the constant value that minimizes the loss, I.e, the "perfect" ablation would be replacing the element with the "best" constant value. It is shown that the proposed techniques improves circuit discovery, identifying a subnetwork that reconstructs the original performance of the model in some task.

Strengths

The contribution is well defined and elegant. The proposed method is experimentally validated on several interpretability-related tasks and it improves over commonly used ablation method.

Weaknesses

I don't see major weaknesses in this work.

Questions

none.

Rating

9

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

see above.

Reply to rebuttal

Thanks to the authors for their reply above! Some comments below: > We agree and have rewritten this section (see global reply). The term “total ablation” clarifies what we consider ablation: replacing the value of component A with a random variable that is independent from the input X, so the new value of A cannot provide any information about X. We think the revised section 2.2 crystallizes why it’s desirable to reduce ∆. Thanks! I think this is a step in the right direction. I still do not totally agree, because I think there is something missing from the list. The list includes deletion and spoofing, but it could also include “insertion”. Similar to our other income example, changing someone’s income from 100k to 40k might be like deletion (I don’t think it is like spoofing, unless this new feature vector is totally OOD / logically impossible), but I think it is more like “inserting a new value.” For what it’s worth, some of my original criticisms still stand, like the idea of “inconsistency” with other information that is relied on in the “spoofing” definition. But I think the new 2.2 is clearer, and I don’t really expect this paper to present an unassailable definition of inconsistency anyway. >Interesting example! In this case, we think OA does, in fact, answer an intelligible question: “How much worse does the model perform if it had to treat everyone as having the same income?” which is arguably extremely similar to “How much worse would the model perform if it had no information to distinguish anyone’s income?,” i.e. "How important is information about income to the model's performance?” Interesting! I agree with your first claim here, and thanks for pointing it out. We do know the meaning of this operation, which is intervening on everyone’s income for it to be the same. I also understand that you treat the 2nd and 3rd question as the same. This is reasonable, and people used to do it years ago when estimating feature importance for decision trees. They would shuffle the column of a dataset, and check the accuracy of the model. That is, they would do marginal resampling. But the thing is, many ablation methods imply that the model has no information to distinguish anyone’s income, including marginal resampling, zero ablation, constant ablation and random ablation. So what is special about optimal ablation that it is the *right way* to make sure the model has no information about anyone’s income? My other question continues to be, why are we trying to estimate “importance” when it is ill-defined? But I see that comes up next. > We agree it’s not clear “importance” is uniquely defined, but we think it’s still helpful to discuss! If we know a component represents… Ultimately, I agree the whole argument here. Counterfactuals aren’t perfect, and I think we do need a method that “somehow aggregates causal effects to guide our search for relevant components.” This would be a nice big caveat to add to the paper when the word “importance” gets used, since it has haunted interpretability papers for years. >We updated Table 1 to adopt the first approach for every method (i.e. conditioning on sequence position), which shows a much larger improvement of OA over resampling. Ok great! >…We present Table 1 to confirm this, and do not mean to present CF as a ground truth Thanks makes sense. > Even for single components, OA achieves lower ∆ on average than CF (0.0035 vs 0.0296), in line with what we see for circuits. The disparity may be larger for circuits due to ablating many components at once. Thanks, this also makes sense. It would probably be worth adjusting some of the language in the paper to not present CF so much as a ground truth, but more as the closest thing we had to a ground-truth (to the extent the intro does this). Then it would help to present Table 1 and the single-components results on somewhat equal footing, to show that OA could even be better than CF in some ways, while also being easier/automatic. Right now the Table 1 result reads more like a sanity check/validation before moving into “real” results. --- Based on this discussion, I plan on keeping my score at 8. I leave it to the authors to clean up some of the motivation/argument for the method, and to figure out how to present some of the extra UGS / OCA results in a better way, but the core the paper is very good.