Paper deep dive

Decomposing and Editing Predictions by Modeling Model Computation

Harshay Shah, Andrew Ilyas, Aleksander Madry

Year: 2024Venue: arXiv preprintArea: Model EditingType: EmpiricalEmbeddings: 154

Models: CLIP ViT-B/16, GPT-2, Phi-2, ResNet-18, ResNet-50, ViT-B/16

Intelligence

Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%

Last extracted: 3/12/2026, 7:09:50 PM

Summary

The paper introduces 'component modeling', a framework to decompose machine learning model predictions into contributions from individual components (e.g., attention heads, filters). It proposes 'Coar' (Component Attribution via Regression), a scalable algorithm that estimates these contributions by treating them as a linear regression problem on counterfactual ablation data. The authors demonstrate that Coar-derived attributions enable effective model editing, including fixing errors, class forgetting, and improving robustness.

Entities (5)

Coar · algorithm · 100%Component Attribution · task · 95%Component Modeling · framework · 95%Model Editing · application · 95%Computation Graph · concept · 90%

Relation Signals (3)

Coar → performs → Component Attribution

confidence 100% · We then present Coar, a scalable algorithm for estimating component attributions

Component Attribution → enables → Model Editing

confidence 95% · Finally, we show that component attributions estimated with Coar directly enable model editing

Component Modeling → includes → Component Attribution

confidence 90% · We focus on a special case of this task, component attribution

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract:How does the internal computation of a machine learning model transform inputs into predictions? In this paper, we introduce a task called component modeling that aims to address this question. The goal of component modeling is to decompose an ML model's prediction in terms of its components -- simple functions (e.g., convolution filters, attention heads) that are the "building blocks" of model computation. We focus on a special case of this task, component attribution, where the goal is to estimate the counterfactual impact of individual components on a given prediction. We then present COAR, a scalable algorithm for estimating component attributions; we demonstrate its effectiveness across models, datasets, and modalities. Finally, we show that component attributions estimated with COAR directly enable model editing across five tasks, namely: fixing model errors, ``forgetting'' specific classes, boosting subpopulation robustness, localizing backdoor attacks, and improving robustness to typographic attacks. We provide code for COAR at this https URL .

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

153,801 characters extracted from source content.

Expand or collapse full text

bibliography/fmt_bib.bib Coar Harshay Shah harshay@mit.edu MIT Andrew Ilyas ailyas@mit.edu MIT Aleksander Mądry madry@mit.edu MIT Decomposing and Editing Predictions by Modeling Model Computation Harshay Shah harshay@mit.edu MIT Andrew Ilyas ailyas@mit.edu MIT Aleksander Mądry madry@mit.edu MIT Abstract How does the internal computation of a machine learning model transform inputs into predictions? In this paper, we introduce a task called component modeling that aims to address this question. The goal of component modeling is to decompose an ML model’s prediction in terms of its components—simple functions (e.g., convolution filters, attention heads) that are the “building blocks” of model computation. We focus on a special case of this task, component attribution, where the goal is to estimate the counterfactual impact of individual components on a given prediction. We then present Coar, a scalable algorithm for estimating component attributions; we demonstrate its effectiveness across models, datasets, and modalities. Finally, we show that component attributions estimated with Coar directly enable model editing across five tasks, namely: fixing model errors, “forgetting” specific classes, boosting subpopulation robustness, localizing backdoor attacks, and improving robustness to typographic attacks. We provide code for Coar at https://github.com/MadryLab/modelcomponents. 1 Introduction Despite their predictive power, large-scale machine learning (ML) models remain black boxes. In particular, the complex internal computation that these models perform to transform inputs into predictions makes it difficult to understand model behavior and, as a result, detect failure modes prior to deployment [beery2018recognition, sheng2019woman, geirhos2020shortcut]. In response to this difficulty, a line of work in ML interpretability aims to shed light on model computation by analyzing model components—intuitively “grouped” model parameters such as convolutional filters or attention heads. For example, feature visualization methods [simonyan2013deep, zeiler2014visualizing, ghiasi2022vision] identify components in vision models that detect visual concepts such as curves [olah2020overview] and objects [bau2020understanding]. Representation-based probes [alain2016understanding] identify groups of components in language models that encode sentiment [radford2017learning], part-of-speech tags [blevins2018deep], and syntactic structure [hewitt2019designing]. Finally, mechanistic interpretability analyses [wang2022interpretability, nanda2023progress] uncover specific components that encode a model behavior of interest, e.g., “knowledge neurons” [dai2021knowledge] and “induction heads” [olsson2022in]. Broadly, these works leverage different tools to answer the question: How do individual model components shape model behavior? In this work, we propose a new (and complementary) approach to studying this question. Our point of start is to rephrase the question, instead asking: How do changes to model components collectively change individual model predictions? We turn this rephrased question into a concrete task called component modeling. In this task, the goal is to build an interpretable counterfactual estimator of how a model’s output would change in response to interventions made to its components. In the rest of the paper, we present a general approach to building such estimators, which turn out to be highly predictive in large-scale settings. Beyond shedding light on how model components collectively contribute to a given prediction, these estimators enable effective model editing, allowing us to design targeted interventions that induce certain desirable model predictions. Figure 1: A summary of the component modeling framework. Roadmap & contributions. The main contribution of our work is a framework for decomposing model predictions in terms of model components, which we show has direct applications to model editing. Figure 1 summarizes these contributions. Specifically, in this paper we: 1. Introduce the component modeling framework: We formalize our goal of understanding how model components shape ML predictions as a concrete task called component modeling (Definition 1). The objective of this task is to learn a counterfactual estimator, or component model, that accurately predicts the effect of ablating a subset of model components on a given model prediction (Equation 1). Intuitively, this task operationalizes the idea that if we can “understand” how model components shape a prediction, we should also be able to estimate how the prediction would change if we were to ablate a subset of components. 2. Instantiate the framework via component attribution: We focus our attention on a special “linear” case of component modeling called component attribution, where we assign a score to each component, and estimate the counterfactual effect of ablating a set of components as the sum of their corresponding scores (Definition 2). Component attributions allow us to directly read off the “contribution” of every component to a prediction, abstracting away the complexity of the model’s internal computation. 3. Propose an algorithm for efficient component attribution: We develop Coar (component attribution via regression), a scalable way to estimate component attributions (Section 3). Through experiments on both image classifiers and language models, we show that Coar yields component attributions that can accurately predict how model predictions change in response to component-level ablations (Section 4). 4. Demonstrate that component attributions enable model editing: Component attributions from Coar directly enable edits to large-scale classifiers without additional training (Section 5). Specifically, we propose an editing procedure (Coar-Edit) that designs targeted ablations by using Coar attributions as a counterfactual estimator—given an objective, Coar-Edit finds ablations for which estimated model outputs perform well. We stress-test Coar-Edit through five editing tasks: fixing model errors (§ 5.1), selectively “forgetting” an entire class (§ 5.2), boosting subpopulation robustness (§ 5.3), localizing backdoor attacks (§ 5.4), and mitigating typographic attacks (§ 5.5). Paper organization. We begin by formalizing the component modeling task and its special case, component attribution, in Section 2. We then describe our method, Coar, for estimating component attributions in Section 3, and demonstrate its effectiveness on large-scale models in Section 4. Finally, we stress-test the practical utility of Coar attributions for model editing in Section 5. 2 Setup and Problem Statement Consider a typical supervised learning setup. We have a set S of input-label pairs (or examples) zi=(xi,yi)subscriptsubscriptsubscriptz_i=(x_i,y_i)zitalic_i = ( xitalic_i , yitalic_i ), and a trained model M that maps inputs x to predicted labels M⁢(x)M(x)M ( x ). We define the model output function fM⁢(z)∈ℝsubscriptℝf_M(z) _M ( z ) ∈ blackboard_R as any statistic that quantifies the correctness of model M on the example z. For instance, the model output fM⁢(z)subscriptf_M(z)fitalic_M ( z ) can be the cross-entropy loss in a classification task, or the squared loss in a regression task. In this work, we will think of the model M not as a black box, but instead as the output of a computation graph GMsubscriptG_MGitalic_M 1974computational. Each parameterized node of this graph—which we call a component—is a function mapping its incoming edges to an outgoing edge. For example, a d-dimensional linear model M naturally admits a computation graph GMsubscriptG_MGitalic_M with d components—one component Ci⁢(z)=wi⋅zsubscript⋅subscriptC_i(z)=w_i· zCitalic_i ( z ) = witalic_i ⋅ z for each parameter wisubscriptw_iwitalic_i—followed by a summation that combines the components into an output. For more complex models, there are often multiple valid computation graphs GMsubscriptG_MGitalic_M we could consider. For example, if M is a Transformer 2017attention, the components might be multi-head attention blocks, individual attention heads, or even individual parameters. In general, the component set depends on the model architecture and the level of granularity that we wish to study. Component modeling. Our goal in this work is to understand the behavior of the model M in terms of its components. By viewing the model M as a computation graph GMsubscriptG_MGitalic_M over a set of components C, we can restate our goal as follows: Given a model M and example z, how do individual components c∈Cc∈ Cc ∈ C combine to yield the model output fM⁢(z)subscriptf_M(z)fitalic_M ( z )? Of course, there is a trivial answer to this question: the components c∈Cc∈ Cc ∈ C combine through the very computation graph used to define C. This answer is correct but not satisfying, as it does not get us closer to our conceptual goal of understanding model behavior in terms of components. What we are truly after is a simple, interpretable function capturing how components in C impact fM⁢(z)subscriptf_M(z)fitalic_M ( z ). To make this more precise, we define the component counterfactual function fM⁢(z,C′)subscriptsuperscript′f_M(z,C )fitalic_M ( z , C′ ) as fM⁢(z,C′):=model output ⁢fM⁢(z)⁢ on example ⁢z⁢ after ablating components ⁢C′⊆C,assignsubscriptsuperscript′model output subscript on example after ablating components superscript′f_M(z,C ):=model output f_M(z) on example z % after ablating components C C,fitalic_M ( z , C′ ) := model output fitalic_M ( z ) on example z after ablating components C′ ⊆ C , (1) where “ablating” here corresponds to any intervention that overrides or patches the parameters corresponding to components c∈C′c∈ C c ∈ C′ (e.g., by setting them to zero [olsson2022in] or by adding random noise [meng2022locating]). Equation (1) allows us to operationalize our goal as a counterfactual estimation task. In this task, we want to estimate component counterfactuals fM⁢(z,C′)subscriptsuperscript′f_M(z,C )fitalic_M ( z , C′ ) using a much simpler function, which we call a component model. Definition 1 (Component modeling). Fix a model M with computation graph GMsubscriptG_MGitalic_M, component set C=c1,…,cNsubscript1…subscriptC=\c_1,…,c_N\C = c1 , … , citalic_N , and model output function fMsubscriptf_Mfitalic_M. For any subset of model components C′⊆Csuperscript′C C′ ⊆ C, let C′subscript0superscript′ 0_C 0italic_C′ be the corresponding ablation vector of C′, defined as a N-dimensional vector where (C′)i=0if ⁢ci∈C′1otherwise.subscriptsubscript0superscript′cases0if subscriptsuperscript′1otherwise. ( 0_C )_i= cases0&if c_i∈ C^% \\ 1&otherwise. cases( 0italic_C′ )i = start_ROW start_CELL 0 end_CELL start_CELL if citalic_i ∈ C′ end_CELL end_ROW start_ROW start_CELL 1 end_CELL start_CELL otherwise. end_CELL end_ROW Given an example z, a component model for z is a function g(z):0,1N→ℝ:superscript→superscript01ℝ g^(z):\0,1\^N ( z ) : 0 , 1 N → blackboard_R that maps ablation vectors of subsets C′ to estimates of the counterfactual fM⁢(z,C′)subscriptsuperscript′f_M(z,C )fitalic_M ( z , C′ ). In other words, the high-level goal of component modeling is to build an estimator that can directly answer counterfactual questions like “what would happen to my classifier’s prediction on a given image if I ablated a specific set of components C′⊆Csuperscript′C C′ ⊆ C?” without having to intervene on the computation graph GMsubscriptG_MGitalic_M and ablate components in C′. Component attribution. In this work, we consider a subcase of component modeling—which we call component attribution—where the function g(z)superscriptg^(z)g( z ) is linear in its input. That is, a component attribution for example z assigns a score i(z)superscriptsubscript w_i^(z)italic_witalic_i( z ) to each component ci∈Csubscriptc_i∈ Ccitalic_i ∈ C, and predicts the effect of ablating C′⊂Csuperscript′C ⊂ C′ ⊂ C as the sum of the scores corresponding to components in C∖C′ C C ∖ C′. Definition 2 (Component attribution). Fix a model M with output function fMsubscriptf_Mfitalic_M and component set C=c1,…,cNsubscript1…subscriptC=\c_1,…,c_N\C = c1 , … , citalic_N . A component attribution for example z is a set of parameters (z)≔1(z),…,N(z),b(z)≔superscriptsubscriptsuperscript1…subscriptsuperscriptsuperscript θ^(z) \ w^(z)_1\!\!,…, w^(z)_N% ,b^(z)\italic_θ( z ) ≔ italic_w( z )1 , … , italic_w( z )N , b( z ) which parameterize a linear component model, i.e., a function g(z)superscriptg^(z)g( z ) such that fM⁢(z;C′)≈g(z)⁢(C′)≔C′⊤⁢(z)+b(z)subscriptsuperscript′subscript0superscript′≔superscriptsubscript0superscript′topsuperscriptsuperscriptf_M(z;C )≈ g^(z)( 0_C ) 0_C^% w^(z)+b^(z)fitalic_M ( z ; C′ ) ≈ g( z ) ( 0italic_C′ ) ≔ 0italic_C′⊤ italic_w( z ) + b( z ) Component attribution satisfies our goal of finding a simple, interpretable account of how model components combine to form predictions. In particular, a component attribution for example z decomposes a model’s output on z into the contributions i(z)superscriptsubscript w_i^(z)italic_witalic_i( z ) of each individual component cisubscriptc_icitalic_i. Remark 1 (Linearity and misspecification). Modern ML models comprise complex computation graphs with highly non-linear interactions among model components. For such models, it is unclear a priori why the effect of ablating components on model outputs (i.e., component counterfactuals (1)) should be well-approximated by linear component attributions (2), which sum fixed additive effects of individual components. Still, despite this evident misspecification, our results on large-scale vision and language models in Section 4 show that component attributions can accurately predict component counterfactuals. 3 Component attribution with Coar In Section 2, we formalized our high-level goal of understanding how models internally process examples into a counterfactual estimation task called component modeling (Definition 1), of which we study a special (linear) case called component attribution (Definition 2). Now, we show how to estimate component attributions (z)superscript θ^(z)italic_θ( z ) by casting the counterfactual estimation task as a regression problem. Specifically, we now describe Coar (component attribution via regression), a general component attribution method for models ranging from random forests to deep neural networks. Approach. Consider a fixed model output fM⁢(⋅)subscript⋅f_M(·)fitalic_M ( ⋅ ) of interest, and a corresponding computation graph GMsubscriptG_MGitalic_M that encodes the model components C at the desired level of granularity. Additionally, we fix an ablation method, i.e., a procedure for “overriding” or patching any given subset C′⊂Csuperscript′C ⊂ C′ ⊂ C of the model components in the computation graph GMsubscriptG_MGitalic_M. Our method Coar takes in an example z and outputs a corresponding component attribution vector (z)∈ℝ|C|+1superscriptsuperscriptℝ1 θ^(z) ^|C|+1italic_θ( z ) ∈ blackboard_R| C | + 1 (Definition 2). To do so, Coar casts the task of predicting component counterfactuals as a supervised learning problem, which we solve in two steps: 1. Construct a component dataset. We construct a dataset D(z)superscriptD^(z)D( z ) of component counterfactuals for the example z. Each “datapoint” in D(z)superscriptD^(z)D( z ) consists of a component subset Ci⊆CsubscriptC_i CCitalic_i ⊆ C and its corresponding counterfactual fM⁢(z,Ci)subscriptsubscriptf_M(z,C_i)fitalic_M ( z , Citalic_i ) (see (1))—we evaluate the latter by simply ablating the components in CisubscriptC_iCitalic_i and evaluating the model on example z. In this work, we choose the component subsets CisubscriptC_iCitalic_i to be random αtrainsubscripttrain _trainαtrain-fraction subsets of the component set C, for a ablation fraction hyperparameter αtrain>0subscripttrain0 _train>0αtrain > 0.111 We opt for this random α-fraction sampling method for simplicity—it may be possible to make Coar more statistically efficient by choosing the subsets CisubscriptC_iCitalic_i more carefully. The output of this step is a component dataset D(z)superscript D^(z)D( z ) :=(C1,fM⁢(z,C1)),…,(Cm,fM⁢(z,Cm)),assignabsentsubscript1subscriptsubscript1…subscriptsubscriptsubscript :=\(C_1,f_M(z,C_1)),…,(C_m,f_M(z,C_m))\,:= ( C1 , fitalic_M ( z , C1 ) ) , … , ( Citalic_m , fitalic_M ( z , Citalic_m ) ) , (2) where Ci∼Uniform⁢(C′⊂C⁢ : ⁢|C′|=αtrain⁢|C|)similar-tosubscriptUniformsuperscript′ : superscript′subscripttrainC_i (\C ⊂ C : |C |= _% train|C|\)Citalic_i ∼ Uniform ( C′ ⊂ C : | C′ | = αtrain | C | ). We study the effect of varying the ablation fraction α on Coar in Section E.1. 2. Fit a linear estimator. We then use the dataset D(z)superscriptD^(z)D( z ) to fit component attribution parameters (z)superscript θ^(z)italic_θ( z ) for each example z (see Definition 2). More specifically, for each example z, we minimize the squared loss between the component counterfactuals from Step 1 and their corresponding attribution-based predictions by solving the following linear regression problem: (z)≔arg⁡minb∈ℝ,∈ℝ|C|⁢∑(Ci,fM⁢(z,Ci))∈D(z)(b+Ci⊤⁢−fM⁢(z,Ci))2,≔superscriptsubscriptformulae-sequenceℝsuperscriptℝsubscriptsubscriptsubscriptsubscriptsuperscriptsuperscriptsuperscriptsubscript0subscripttopsubscriptsubscript2 θ^(z) _b ,\, w ^|C|% _(C_i,f_M(z,C_i))∈ D^(z) (b+ 0_C_i w% -f_M(z,C_i) )^2,italic_θ( z ) ≔ arg minitalic_b ∈ blackboard_R , italic_w ∈ blackboard_R| C | ∑( C start_POSTSUBSCRIPT i , fitalic_M ( z , Citalic_i ) ) ∈ D( z ) end_POSTSUBSCRIPT ( b + 0italic_C start_POSTSUBSCRIPT i end_POSTSUBSCRIPT⊤ italic_w - fitalic_M ( z , Citalic_i ) )2 , (3) where again Cisubscript0subscript 0_C_i0italic_C start_POSTSUBSCRIPT i end_POSTSUBSCRIPT is the ablation vector of CisubscriptC_iCitalic_i (Definition 1). Our component model is then g(z)⁢(C′)≔C′⊤⁢(z)+b(z).≔superscriptsubscript0superscript′subscript0superscript′topsuperscriptsuperscriptg^(z)( 0_C ) 0_C w^(z)+b^% (z).g( z ) ( 0italic_C′ ) ≔ 0italic_C′⊤ italic_w( z ) + b( z ) . (4) We provide pseudocode for Coar in Section A.1. As we discussed in Section 2, the resulting component attribution (z)≔((z),b(z))≔superscriptsuperscriptsuperscript θ^(z) ( w^(z),b^(z))italic_θ( z ) ≔ ( italic_w( z ) , b( z ) ) is interpretable in that the coefficient j(z)superscriptsubscript w_j^(z)italic_witalic_j( z ) estimates how the model output on example z would change if we were to ablate component cjsubscriptc_jcitalic_j. We can thus view this coefficient as the (estimated) additive contribution of component cjsubscriptc_jcitalic_j to the model output. The above two-step approach is simple and highly scalable—we can construct the dataset D(z)superscriptD^(z)D( z ) with just forward passes on the given model to compute component counterfactuals, and optimize the linear regression problem (eq. 3) with off-the-shelf GPU-based solvers—see Section A.4 for details. This enables us to apply Coar on large-scale models (e.g., ViT 2021image) and datasets (e.g., ImageNet 2009imagenet), as shown in the next section. Instantiating Coar for classifiers. Our method Coar is general in that we can use it to study any machine learning model M that has a corresponding output function fMsubscriptf_Mfitalic_M and computation graph GMsubscriptG_MGitalic_M. In this work, we primarily use Coar to analyze models trained on classification tasks. Although the computation graph GMsubscriptG_MGitalic_M will vary based on the specific model architecture we are studying, across all models we use the standard correct-class margin [ilyas2022datamodels] as the model output fMsubscriptf_Mfitalic_M, i.e., fM⁢(z):=(logit for correct class)−(highest logit for incorrect class).assignsubscript(logit for correct class)(highest logit for incorrect class).f_M(z):=(logit for correct class)-(highest logit for incorrect % class).fitalic_M ( z ) := (logit for correct class) - (highest logit for incorrect class). (5) a quantity whose sign indicates the correctness of model M on the example z. For the latter, we choose to ablate component subsets C′⊂Ssuperscript′C ⊂ SC′ ⊂ S by simply setting the parameters of the components in C′ to zero [wang2022interpretability, olsson2022in]. We use Coar with alternative model output functions and ablation methods in Section E.3 and Section E.2 respectively. Remark 2 (Ablation is not removal). As noted in prior work 2022causal, ablation methods (e.g., setting weights or activations to zero) do not “remove” model components from the computation graph. Instead, such ablations shift the activations off-distribution in a systematic way—the goal of component attribution (Definition 2) is to predict the change in model outputs induced by this shift. We use zero-ablation as our ablation method because it is a common choice in the literature [olsson2022in, wang2022interpretability]. In Section E.2, we show that Coar can estimate component attributions with alternative ablation methods as well. 4 Does Coar learn accurate component attributions? We now evaluate whether Coar-estimated attributions accurately predict component counterfactuals (as in (1)) for deep neural networks trained on image classification and language modeling. Datasets, models, and components. We apply Coar to compute component attributions in three different setups: • Setup A: A ResNet-18 [he2015deep] trained on the CIFAR-10 dataset [krizhevsky2009learning], with a computation graph GAsubscriptG_AGitalic_A comprising |C|=2,3062306|C|=2,306| C | = 2 , 306 components. Specifically, each model component ci∈Csubscriptc_i∈ Ccitalic_i ∈ C corresponds to a convolutional filter in the model, and ablating a set of components C′⊂Csuperscript′C ⊂ C′ ⊂ C means setting all the weights in the corresponding filters to zero. • Setup B: A ResNet-50 trained on the ImageNet dataset [deng2009imagenet], with a computation graph GBsubscriptG_BGitalic_B comprising |C|=22,72022720|C|=22,720| C | = 22 , 720 components. Again, each component here corresponds to a convolutional filter in one of the 49494949 convolution layers of the ResNet-50. • Setup C: A Vision Transformer (ViT-B/16) [dosovitskiy2021image] trained on ImageNet, whose computation graph GCsubscriptG_CGitalic_C comprises 82,9448294482,94482 , 944 components. Each component here corresponds to a row of a weight matrix in one of 12121212 transformer blocks of the ViT, and ablating a set of components means setting the corresponding rows to zero. We provide additional details on the models and datasets in Section A.2. Applying Coar. We use Coar to obtain component attributions (one for each test example) in each setup. Specifically, for a given model, we first construct a component dataset D(z)superscriptD^(z)D( z ) for each example z (as in Step 1 of Section 3) by randomly ablating αtrainsubscripttrain _trainαtrain fraction of all components and evaluating the resulting correct-class margin (5) on z, where αtrain=10%,5%,5%subscripttrainpercent10percent5percent5 _train=\10\%,5\%,5\%\αtrain = 10 % , 5 % , 5 % for setup A,B,CABC\A,B,C\ A , B , C above. We repeat this m times, yielding a component dataset D(z)superscriptD^(z)D( z ) of size m for each example z—we use m=50000,100000,20000050000100000200000m=\50000,100000,200000\m = 50000 , 100000 , 200000 for setup A,B,CABC\A,B,C\ A , B , C above. We then run linear regressions on the resulting datasets (as in Step 2 of Section 3) to yield the final component attributions. We defer implementation details to Section A.4 and study the effect of the dataset size m and ablation fraction αtrainsubscripttrain _trainαtrain on the resulting attributions in Sections C.5 and C.4. Evaluation metric. We evaluate component attributions based on their ability to estimate unseen component counterfactuals (1), i.e., the result of ablating component subsets C′ not observed at training time. Specifically, we sample a new collection of k component subsets Dtest(z)≔C1′,C2′,…,Ck, where Ci′∼Unif⁢(C′⊂C:|C′|=αtest⁢|C|),formulae-sequence≔superscriptsubscripttestsuperscriptsubscript1′subscript2′…subscript where similar-tosuperscriptsubscript′Unifconditional-setsuperscript′subscripttestD_test^(z) \C_1 ,C_2 ,…,C_k\,% where C_i (\C ⊂ C% :|C |= _test|C|\),Dtest( z ) ≔ C1′ , C2′ , … , Citalic_k , where Citalic_i′ ∼ Unif ( C′ ⊂ C : | C′ | = αtest | C | ) , where 0<αtest<10subscripttest10< _test<10 < αtest < 1 is the ablation fraction used at evaluation time. Varying αtestsubscripttest _testαtest allows us to evaluate attributions on in-distribution (αtest=αtrainsubscripttestsubscripttrain _test= _trainαtest = αtrain) and out-of-distribution (αtest≠αtrainsubscripttestsubscripttrain _test≠ _trainαtest ≠ αtrain) component counterfactuals. To quantify the predictiveness of component attributions, we use Dtest(z)superscriptsubscripttest D_test^(z)Dtest( z ) to measure the Pearson correlation between component counterfactuals fM⁢(z,Ci′)subscriptsuperscriptsubscript′f_M(z,C_i )fitalic_M ( z , Citalic_i′ ) and their corresponding attribution-based estimates g(z)⁢(Ci′)superscriptsubscript0superscriptsubscript′ g^(z)( 0_C_i )g( z ) ( 0italic_C start_POSTSUBSCRIPT i′ end_POSTSUBSCRIPT ) (eq. 4.), i.e., ρ⁢(z)≔Pearson-⁢ρ⁢(fM⁢(z,C1′),…,fM⁢(z,Ck′)⏟ground-truth counterfactuals,g(z)⁢(C1′),…,g(z)⁢(Ck′)⏟attribution-based estimates).≔Pearson-subscript⏟subscriptsuperscriptsubscript1′…subscriptsuperscriptsubscript′ground-truth counterfactualssubscript⏟superscriptsubscript0superscriptsubscript1′…superscriptsubscript0superscriptsubscript′attribution-based estimatesρ(z) -ρ ( \f_M(z,C_1 % ),...,f_M(z,C_k )\_ground-truth counterfactuals% , \g^(z)( 0_C_1 ),...,g^(z)( 0_C_k^% )\_attribution-based estimates ).ρ ( z ) ≔ Pearson- ρ ( under⏟ start_ARG fitalic_M ( z , C1′ ) , … , fitalic_M ( z , Citalic_k′ ) end_ARGground-truth counterfactuals , under⏟ start_ARG g( z ) ( 0italic_C start_POSTSUBSCRIPT 1′ end_POSTSUBSCRIPT ) , … , g( z ) ( 0italic_C start_POSTSUBSCRIPT k′ end_POSTSUBSCRIPT ) end_ARGattribution-based estimates ) . (6) Baselines. We use the evaluation metric described above (Equation 6) to compare Coar attributions with four baselines, two adapted from related work and two natural baselines. We defer implementation details to Section A.3. • Adapted baselines (NC, I): We adapt neuron conductance (NC) [dhamdhere2018important] and internal influence (I) [leino2018influence] to the component attribution setting. Both methods use integrated gradients [sundararajan2017axiomatic] (an input-space feature attribution method) to compute importance scores for each component cisubscriptc_icitalic_i. To compare these methods to Coar, we apply NC and I to model outputs on example z, and interpret the resulting scores as the attribution coefficients wi(z)superscriptsubscript w_i^(z)witalic_i( z ). • Specialized baselines (LOO, GP): We also consider two other baselines. First, leave-one-out (LOO) ablates individual components cisubscriptc_icitalic_i and estimates the corresponding coefficient based on the effect of the ablation, setting i(z)=fM⁢(z,ci)−fM⁢(z,∅)subscriptsuperscriptsubscriptsubscriptsubscript w^(z)_i=f_M(z,\c_i\)-f_M(z, )italic_w( z )i = fitalic_M ( z , citalic_i ) - fitalic_M ( z , ∅ ). We also consider gradient-times-parameter (GP), which approximates the leave-one-out effect of each component using a first-order Taylor approximation, setting i(z)=∇cifM⁢(z,∅)⋅δcisubscriptsuperscript⋅subscript∇subscriptsubscriptsubscriptsubscript w^(z)_i= _c_if_M(z, )· _c_iitalic_w( z )i = ∇c start_POSTSUBSCRIPT i end_POSTSUBSCRIPT fitalic_M ( z , ∅ ) ⋅ δitalic_c start_POSTSUBSCRIPT i end_POSTSUBSCRIPT, where δcisubscriptsubscript _c_iδitalic_c start_POSTSUBSCRIPT i end_POSTSUBSCRIPT is the parameter-space change in cisubscriptc_icitalic_i induced by the ablation method of choice. Remark 3 (Relation between baselines and patching). Readers familiar with the field of mechanistic interpretability may observe that the leave-one-out (LOO) baseline resembles activation patching [vig2020investigating, meng2022locating, zhang2023towards] and the gradient-times-parameter (GP) baseline resembles attribution patching [syed2023attribution, kramar2024atp*]. Indeed, activation patching ablates individual activations to estimate their effect, and attribution patching considers a first-order gradient-based approximation. The key difference is that the ablations we consider are in parameter space rather than activation space, and that we study a different model output. 4.1 Results We now use the setup described above to test whether Coar learns accurate component attributions for setups A,B,CABC\A,B,C\ A , B , C . For each task, we first use Coar to estimate a component attribution for every example z in the corresponding test set. We then evaluate these component attributions using the correlation metric ρ⁢(z)ρ(z)ρ ( z ) defined in Equation 6. Figure 2 depicts our results. Figure 2: Evaluating Coar attributions. We evaluate whether component attributions computed using our procedure Coar accurately predict component counterfactuals (1). We compare Coar to four baselines (described in Section 4) on three image classification setups (one per row). The subfigures on the left each focus on a single example z (visualized in the bottom-right corner of each plot), and show that for each setup, the ground-truth component counterfactuals fM⁢(z,⋅)subscript⋅f_M(z,·)fitalic_M ( z , ⋅ ) (x-axis) and attribution-based estimates g(z)⁢(⋅)superscript⋅g^(z)(·)g( z ) ( ⋅ ) (y-axis) exhibit high correlation ρ⁢(z)ρ(z)ρ ( z ). On the right, we observe that Coar attributions exhibit high average correlation z⁢[ρ⁢(z)]subscriptdelimited-[]E_z[ρ(z)]blackboard_Ez [ ρ ( z ) ] over test examples, outperforming all baselines in each task and for all ablation fractions αtestsubscripttest _testαtest. The asterisk (*) in each legend denotes αtrainsubscripttrain _trainαtrain, the ablation fraction used to fit the component attributions. Example-level analysis. On the left side of each row in Figure 2, we focus on an individual test example z from each task. For each example z, we ablate random component subsets C′⊂Csuperscript′C ⊂ C′ ⊂ C of size αtest⋅|C|⋅subscripttest _test·|C|αtest ⋅ | C | (for αtest=αtrainsubscripttestsubscripttrain _test= _trainαtest = αtrain) from the model and estimate the correlation ρ⁢(z)ρ(z)ρ ( z ) from Equation 6. Across all three tasks, we observe that Coar learns accurate component attributions for the selected test examples. In Section C.6, we provide additional (randomly selected) example-specific correlation plots, as well as analogous plots for all baselines described above. Aggregate analysis. The right side of each row in Figure 2 plots the average correlation between the ground-truth counterfactuals and attribution-based estimates over test examples, i.e., z⁢[ρ⁢(z)]subscriptdelimited-[]E_z[ρ(z)]blackboard_Ez [ ρ ( z ) ]. We also analyze the effect of ablation fraction αtestsubscripttest _testαtest on the average correlation, finding that: (a) Coar outperforms baselines by a large margin across datasets, models, and ablation fractions αtestsubscripttest _testαtest. For example, when ablating αtest=5%subscripttestpercent5 _test=5\%αtest = 5 % of components in the ImageNet ResNet50 (setup B), attribution-based estimates using Coar and the best-performing baseline (LOO) exhibit 0.650.650.650.65 and 0.340.340.340.34 correlation with ground-truth counterfactuals, respectively. Additionally, the adapted baselines, NC and I, exhibit low correlation with the ground-truth counterfactuals in all three setups. (b) The correlation between ground-truth counterfactuals and attribution-based estimates decays gracefully on larger out-of-distribution component subsets, i.e., as αtestsubscripttest _testαtest increases. For example, increasing αtestsubscripttest _testαtest from 10%percent1010\%10 % (equal to αtrainsubscripttrain _trainαtrain) to 12.5%percent12.512.5\%12.5 % and 15%percent1515\%15 % on CIFAR-10 (setup A) only decreases the average correlation of Coar-based estimates from 0.740.740.740.74 to 0.700.700.700.70 and 0.680.680.680.68 respectively. Applying Coar to language models. Although we focus on vision models in this work, our attribution method Coar is general and modality-agnostic. In Appendix B, we show that Coar, without any modification, yields predictive component attributions for language models as well. First, we apply Coar to GPT-2 [radford2019language] evaluated on the next-token prediction task using the TinyStories dataset [eldan2023tinystories] (§B.1). Then, we turn to the zero-shot classification setting and apply Coar to Phi-2 [li2023textbooks] evaluated on the BoolQ question-answering task [clark2019boolq] (§B.2). Additional analysis. In Appendix C, we show that Coar attributions are predictive for out-of-distribution inputs (§C.1), additional architectures (§C.2), additional tasks (§C.3), and different train-time ablation fractions (§C.4). We also show that Coar outperform baselines when trained with 2222-5×5×5 × fewer samples in Section C.5, and provide qualitative analysis in Section C.7. 5 Do Coar attributions enable model editing? In the last section, we showed that Coar attributions accurately predict how model outputs change in response to component-level interventions. We now evaluate the practical utility of Coar by applying it to the problem of model editing. That is, we ask: Is ablating model components identified via Coar attributions an effective way to edit models? To answer this question, we first define model editing in our context and provide a simple method, Coar-Edit, for translating component attributions into model edits. We apply this approach to edit model behavior on individual examples (§5.1), classes (§5.2), subpopulations (§5.3), and concepts (§5.4, §5.5). Our findings indicate that Coar directly enables model editing. Problem setup. Consider a machine learning model M, a target distribution over examples TsubscriptD_TDitalic_T, and a reference distribution RsubscriptD_RDitalic_R. A model edit on M is an intervention that aims to modify performance on the target examples z∼Tsimilar-tosubscriptz _Tz ∼ Ditalic_T in a specific way, while leaving behavior on reference examples z∼Rsimilar-tosubscriptz _Rz ∼ Ditalic_R unchanged. In its most general version, model editing can involve additional training (e.g., constrained fine-tuning 2020modifying or rank-one parameter updates 2020rewriting), targeted modifications to model parameters (e.g., weight pruning 2021sparse or hypernetworks 2021fast), or even architectural modifications (e.g., adaptors [hartvigsen2022aging] or adding neurons 2023transformer) Since our goal is to study model predictions in terms of model components, we restrict model edits to ablation-based interventions in this work. That is, we only consider interventions whose output can be expressed as component counterfactuals (see Equation 1). The goal of an editing method, then, is to identify a subset of model components whose ablation changes performance on a given target set of examples STsubscriptS_TSitalic_T, without impacting model behavior on a reference set of examples SRsubscriptS_RSitalic_R. Definition 3 turns this intuition into a precise definition of the ablation-based model editing problem. Definition 3 (Editing models by ablating components). Consider a model M with computation graph GMsubscriptG_MGitalic_M, component set C=c1,…,cNsubscript1…subscriptC=\c_1,…,c_N\C = c1 , … , citalic_N , and model output function fMsubscriptf_Mfitalic_M. Let TsubscriptTD_TDT and RsubscriptRD_RDR denote target and reference distributions over examples, respectively. An (ϵ,δ)italic-ϵ(ε,δ)( ϵ , δ )-effective model edit for RsubscriptRD_RDR and TsubscriptTD_TDT is an intervention that ablates a subset of components Cedit∈2Csubscripteditsuperscript2C_edit∈ 2^CCedit ∈ 2C such that R⁢[|fM⁢(z,Cedit)−fM⁢(z,∅)|]⏟Effect of edit on reference examples is small≤ϵ and T⁢[|fM⁢(z,Cedit)−fM⁢(z,∅)|]⏟Effect of edit on target examples is large≥δ,formulae-sequencesubscript⏟subscriptsubscriptRdelimited-[]subscriptsubscripteditsubscriptEffect of edit on reference examples is smallitalic-ϵ and subscript⏟subscriptsubscriptTdelimited-[]subscriptsubscripteditsubscriptEffect of edit on target examples is large E_D_R[|f_M(z,C_ % edit)-f_M(z, )|]_Effect of edit on reference examples is % small≤ε and E_D% _T[|f_M(z,C_edit)-f_M(z, )|]_Effect of % edit on target examples is large≥δ,under⏟ start_ARG blackboard_ED start_POSTSUBSCRIPT R end_POSTSUBSCRIPT [ | fitalic_M ( z , Cedit ) - fitalic_M ( z , ∅ ) | ] end_ARGEffect of edit on reference examples is small ≤ ϵ and under⏟ start_ARG blackboard_ED start_POSTSUBSCRIPT T end_POSTSUBSCRIPT [ | fitalic_M ( z , Cedit ) - fitalic_M ( z , ∅ ) | ] end_ARGEffect of edit on target examples is large ≥ δ , (7) where fM⁢(z,C)subscriptf_M(z,C)fitalic_M ( z , C ) denotes the component counterfactual function (1), i.e., model output function fMsubscriptf_Mfitalic_M evaluated on example z after ablating components C. As per Definition 3, each component subset C′⊂2Csuperscript′2C ⊂ 2^CC′ ⊂ 2C defines a potential model edit. That is, effectively editing the model (as in (7)) requires identifying a subset C′ that, when ablated, significantly changes model outputs on the target distribution but not on the reference distribution. A naive approach to this task would thus require searching over the (combinatorial) space of all possible component subsets. Is there a better way? An attribution-based approach to model editing. In this section, our goal is to show that an effective component attribution method (such as Coar) can directly serve as a guide for identifying effective model edits. Key to this utility is a fundamental connection between the attribution problem and the editing problem. In particular, the former answers questions of the form, “how would the model outputs change if we were to ablate a subset of components?” while the latter inverts this question to ask “which components, when ablated, would change model outputs in a specific way?” By identifying the model components that are most “important” to the desired model outputs, an attribution method can thus identify a subset of model components to target through ablation-based editing (see Definition 3). To make this concrete, we propose Coar-Edit, a simple three-step editing approach based on Coar attributions. Specifically, given a model M with a set of model components C, a set of target examples STsubscriptS_TSitalic_T sampled from TsubscriptD_TDitalic_T, and a set of reference examples SRsubscriptS_RSitalic_R sampled from RsubscriptD_RDitalic_R, Coar-Edit identifies a model edit (7) in three steps: 1. Estimate Coar attributions (z)≔((z),b(z))≔superscriptsuperscriptsuperscript θ^(z) ( w^(z),b^(z))italic_θ( z ) ≔ ( italic_w( z ) , b( z ) ) where (z)∈ℝ|C|superscriptsuperscriptℝ w^(z) ^|C|italic_w( z ) ∈ blackboard_R| C | and b(z)∈ℝsuperscriptℝb^(z) ( z ) ∈ blackboard_R for every target and reference example z∈ST∪SRsubscriptsubscriptz∈ S_T∪ S_Rz ∈ Sitalic_T ∪ Sitalic_R. 2. For each model component ci∈Csubscriptc_i∈ Ccitalic_i ∈ C, use a simple t-test in order to quantify the “importance” of component cisubscriptc_icitalic_i to set of target examples STsubscriptS_TSitalic_T relative to the set of reference examples SRsubscriptS_RSitalic_R: τ⁢(ci)≔μi⁢(ST)−μi⁢(SR)σi2⁢(ST)|ST|+σi2⁢(SR)|SR|, where μi⁢(S)=1|S|⁢∑z∈Swi(z)σi2⁢(S)=1|S|⁢∑z∈S(wi(z)−μ⁢(S))2.≔subscriptsubscriptsubscriptsubscriptsubscriptsuperscriptsubscript2subscriptsubscriptsuperscriptsubscript2subscriptsubscript where casessubscript1subscriptsubscriptsuperscriptotherwisesuperscriptsubscript21subscriptsuperscriptsubscriptsuperscript2otherwise τ(c_i) _i(S_T)- _i(S_R) % _i^2(S_T)|S_T|+ _i^2(S_R)|S_R|,% where cases _i(S)= 1|S| _z∈ Sw% ^(z)_i\\ _i^2(S)= 1|S| _z∈ S(w^(z)_i-μ(S))^2. casesτ ( citalic_i ) ≔ divide start_ARG μitalic_i ( Sitalic_T ) - μitalic_i ( Sitalic_R ) end_ARG start_ARG square-root start_ARG divide start_ARG σitalic_i2 ( Sitalic_T ) end_ARG start_ARG | Sitalic_T | end_ARG + divide start_ARG σitalic_i2 ( Sitalic_R ) end_ARG start_ARG | Sitalic_R | end_ARG end_ARG end_ARG , where start_ROW start_CELL μitalic_i ( S ) = divide start_ARG 1 end_ARG start_ARG | S | end_ARG ∑z ∈ S w( z )i end_CELL start_CELL end_CELL end_ROW start_ROW start_CELL σitalic_i2 ( S ) = divide start_ARG 1 end_ARG start_ARG | S | end_ARG ∑z ∈ S ( w( z )i - μ ( S ) )2 . end_CELL start_CELL end_CELL end_ROW (8) 3. To increase model outputs on target examples, ablate a set of components CeditsubscripteditC_editCedit that contains the k most negative scores τisubscript _iτitalic_i, i.e., set Cedit=arg⁡bottom-⁢k⁢(τ⁢(ci):ci∈C),subscripteditbottom-conditional-setsubscriptsubscriptC_edit= -k(\τ(c_i):c_i∈ C\),Cedit = arg bottom- k ( τ ( citalic_i ) : citalic_i ∈ C ) , (9) where the number of ablated components k is a hyperparameter which one can set, e.g., by cross-validation. Similarly, if the goal is to decrease model outputs on STsubscriptS_TSitalic_T, we replace bottom-⁢kbottom-bottom-kbottom- k with top-⁢ktop-top-ktop- k in (9). To make sense of the approach above, note that for every component cisubscriptc_icitalic_i and a set of examples S, μ⁢(S)μ(S)μ ( S ) in Equation 8 leverages attributions to directly estimate μi⁢(S)subscript _i(S)μitalic_i ( S ), the average effect of ablating cisubscriptc_icitalic_i on model predictions for samples in set S. Similarly, the term σi2⁢(S)superscriptsubscript2 _i^2(S)σitalic_i2 ( S ) captures the variation (across examples) of this effect. As a result, the score τ⁢(ci)subscriptτ(c_i)τ ( citalic_i ) in Equation 8 exactly corresponds to the two-sample t-test statistic, with a null hypothesis that the component cisubscriptc_icitalic_i has an equal average effect on the target distribution TsubscriptD_TDitalic_T and the reference distribution RsubscriptD_RDitalic_R. We then use these scores τ⁢(ci):ci∈Cconditional-setsubscriptsubscript\τ(c_i):c_i∈ C\ τ ( citalic_i ) : citalic_i ∈ C in Equation 9 in order to identify components that, when ablated, would change outputs on target examples the most relative to the change in the outputs of reference examples. Remark 4 (Ablation-based edits). Our restriction of model edits to ablations of component subsets in Coar-Edit is a significant one. In particular, simply ablating model components is likely not the most effective way of editing a model. Our goal in this section, however, is not to propose a state-of-the-art model editing method, but rather to answer two questions. First, we want to assess the practical utility of Coar for informing model edits. Second, we want to shed light on whether large-scale models are amenable to zeroth-order editing, i.e., without gradient information. Leveraging Coar attributions in conjunction with editing techniques based on fine-tuning in order to make finer-grained model edits is an interesting direction for future work. Next, we use Coar-Edit to edit large-scale models trained on classification tasks, where we use correct-class margin (5) as the model output function. Specifically, we conduct five experiments to evaluate the effectiveness of Coar-Edit in editing model behavior using only a few examples and without requiring additional training. (a) In Section 5.1, we correct individual model errors without impacting overall performance; (b) In Section 5.2, we selectively “forget” a specific class while preserving model performance on other classes; (c) In Section 5.3, we start with a model that performs disparately across a set of subpopulations, and edit the model to improve its accuracy on underperforming subpopulations; (d) In Section 5.4, we localize (known) backdoor attacks and mitigate them by ablating a small number of components; (e) In Section 5.5, we edit CLIP classifiers to be more robust to typographic attacks. For all experiments, we provide additional details and analyses in Appendix D. 5.1 Editing individual model predictions In this section, we test whether Coar-Edit can modify individual predictions of an ImageNet ResNet50 classifier (Setup B in Section 4) without impacting its overall performance. Specifically, we study the case where the target distribution TsubscriptD_TDitalic_T is a singleton example on which we want to improve performance. An effective model edit in this context (Definition 3) would increase the model’s margin (5) on z to be greater than zero without affecting aggregate model performance. Results. We apply Coar-Edit to edit individual misclassified examples z, setting ST=zsubscriptS_T=\z\Sitalic_T = z and SRsubscriptS_RSitalic_R to be a small set of random samples from the ImageNet dataset. We present our findings in Figure 3. Figure 3a illustrates a single such edit, where we correct the model’s prediction on a specific ImageNet example from “keyboard” to “ballpoint pen” by ablating k=33k=3k = 3 components (0.01%percent0.010.01\%0.01 % of all components). Specifically, increasing the number of ablated components k consistently improves the correct-class margin on target example z (red) without changing the average margin over the training set (light blue) or validation set (dark blue). Figure 3b then visualizes (again, for the specific example being edited in Figure 3a) the examples on which model outputs changes most (and least) drastically. A qualitative inspection here suggests that examples with unchanged margins are dissimilar to z (first row), whereas examples that are most positively (second row) or negatively (third row) impacted share similar visual features with z, e.g., pen-like objects. Finally, Figure 3c shows that we can individually fix every misclassification in the ImageNet validation set while incurring a median accuracy drop of 0.2%percent0.20.2\%0.2 % on the training set (top row) and validation set (bottom row). We defer experiment details and additional results to Section D.1. Figure 3: Editing individual model predictions with Coar-Edit. We edit a ResNet50 model to correct a misclassified ImageNet example z shown on the left. Specifically, ablating a few components via Coar-Edit (see (9)) increases the correct-class margin (5) on example z (red) without changing the average margin on the train set (light blue) or validation set (dark blue). In the center panel, we observe that the examples on which model outputs change the least (top row) due to the edit are visually dissimilar to example z as well as examples on which model outputs change most positively (middle row) and negatively (bottom row). On the right, we find that individually performing model edits to correct every misclassified example in the validation set incurs a median accuracy drop of at most 0.2%percent0.20.2\%0.2 % on the train set (top row) and validation set (bottom row). Figure 4: “Forgetting” a class with Coar-Edit. We edit an ImageNet-trained ResNet-50 (Setup B from Section 4) to selectively degrade performance on the “chain-link fence” class. On the left, we observe that increasing the number of components k ablated via Coar-Edit decreases model accuracy on the “chain-link fence” class (red) while preserving overall accuracy on the train and validation set. In the center panel, we compare class-wise accuracies before and after performing the model edit and observe a significant accuracy drop on the “chain-link fence” class but not on other classes. On the right, we find that the edit transfers to distribution-shifted versions of ImageNet (ImageNet-Sketch [wang2019learning] and ImageNet⋆ ⋆ [vendrow2023dataset]) as targeted, i.e., degrading performance on the “chain-link fence” class without changing average performance. 5.2 “Forgetting” a class We now consider “selective forgetting” problem [wang2023comprehensive], where the goal is to impair model performance on (only) a specific set of examples. In this experiment, we edit the same ImageNet ResNet-50 model (Setup B) as in Section 5.1, with the goal of forgetting the entire “chain-link fence” class. Like before, we use our editing approach Coar-Edit (see (8) and (9)) to identify components that, when ablated, decrease the model’s correct-class margin on examples from the “chain-link fence” class, but not on reference examples from other classes. Results. Figure 4 summarizes our findings. In Figure 4a, we show that ablating just eight (out of 22,7202272022,72022 , 720) model components degrades accuracy on the “chain fence” class from 66%percent6666\%66 % to 20%percent2020\%20 % while preserving overall accuracy on the train and validation set. Then, in Figure 4b, a comparison of class-wise accuracies before and after the edit shows that our approach specifically targets the “chain fence” class without impacting performance on any other class. Finally, in Figure 4c, we evaluate model performance on ImageNet-Sketch 2019learning (top) and ImageNet⋆ ⋆ 2023dataset (bottom) datasets to show that the our edit is robust to distribution shifts in both the target and reference distribution. Through additional experiments in Section D.2, we highlight that (a) our approach is sample-efficient, not needing many samples from the target and reference distributions to find effective edits; and (b) our findings are robust to the choice of class to forget. 5.3 Improving subpopulation robustness Machine learning models often latch onto spurious correlations in the training dataset [geirhos2019imagenet, shah2020pitfalls, hermann2023foundations], resulting in subpar performance on subpopulations where these correlations do not hold [buolamwini2018gender, oakden2020hidden]. In this section, we test whether our editing approach can boost performance on such underperforming subpopulations without degrading overall performance. In particular, we evaluate Coar-Edit on two benchmark datasets—Waterbirds [sagawa2020distributionally] and CelebA [liu2015deep]—where models fare poorly on subpopulations that are underrepresented in the training data. On both datasets, our goal is to improve a given model’s worst-subpopulation accuracy—we defer experiment details to Section D.3. Results. On both datasets, Coar successfully identifies component subsets that correspond to effective model edits. Figure 5 depicts our results. On Waterbirds (Figure 5a), ablating 210210210210 components (0.9%percent0.90.9\%0.9 % of all components) improves worst-subpopulation accuracy from 64%percent6464\%64 % to 83%percent8383\%83 % (red) without degrading its accuracy averaged uniformly over examples (light blue) and subpopulations (dark blue). On CelebA, Figure 5b demonstrates that zeroing out 26262626 of 22,7202272022,72022 , 720 model components improves worst-subpopulation accuracy from 47%percent4747\%47 % to 85%percent8585\%85 % and average-subpopulation accuracy from 84%percent8484\%84 % to 90%percent9090\%90 % while only incurring a 5%percent55\%5 % drop in test set accuracy. Before continuing, we make two observations. First, on both datasets, our editing-based approach is sample-efficient—it does not require subpopulation-level annotations for the training set, and only uses 20202020 random training examples from each subpopulation to find effective model edits. Second, our results indicate that simply ablating a few components from models trained via “standard” empirical risk minimization (ERM) can achieve worst-subpopulation accuracy improvements that are comparable to gains from specialized methods (e.g., based on robust optimization [sagawa2020distributionally], dataset balancing [idrissi2022simple], and generative modeling [goel2020model]). Figure 5: Improving subpopulation robustness with Coar-Edit. We edit pre-trained ResNet-50 models to improve their worst-subpopulation accuracy on two benchmark datasets: Waterbirds [sagawa2020distributionally] and CelebA [liu2015deep]. Before applying Coar-Edit, models fine-tuned on Waterbirds and CelebA attain 87%percent8787\%87 % and 96%percent9696\%96 % test accuracy but only 64%percent6464\%64 % and 47%percent4747\%47 % accuracy on their worst-performing subpopulations, respectively. On the left, applying Coar-Edit by ablating 210210210210 of 22,7202272022,72022 , 720 components in the Waterbirds model increases worst-subpopulation accuracy from 64%percent6464\%64 % to 83%percent8383\%83 % without degrading its accuracy averaged over examples (light blue) and subpopulations (dark blue). Similarly, on the right, editing the CelebA model by ablating a targeted subset of 26262626 components improves worst-subpopulation accuracy from 47%percent4747\%47 % to 85%percent8585\%85 %. 5.4 Localizing backdoor attacks We now use Coar-Edit to analyze the sensitivity of model predictions to backdoor attacks [biggio2012poisoning, gu2017badnets], where an adversary plants a spurious correlation in the training dataset and uses it as a trigger to override predictions at test time. In this experiment, we first train a ResNet18 model on a modified CIFAR-10 dataset in which half of the training examples in the “airplane” class contain a planted blue-squared trigger, as shown in Figure 6a. Then, using Coar-Edit, we evaluate whether the effect of this trigger on predictions can be localized to a few components, which, if ablated, induces model robustness to backdoor attacks without degrading overall performance. Results. Figure 6 summarizes our findings. Figure 6a shows that prior to editing, the model trained on the modified CIFAR-10 dataset (top row) is sensitive to backdoor attacks—simply adding the “airplane” trigger to test examples drops model accuracy from 89%percent8989\%89 % (middle row) to 37%percent3737\%37 % (bottom row). To localize the effect of the trigger, we use Coar-Edit over ten paired examples—i.e., examples with and without the backdoor trigger—to identify a few components that, when ablated, correct the misclassifications induced by the trigger without impacting predictions on test examples without the trigger. In Figure 6b, we find that editing the model by ablating 25252525 components (1%percent11\%1 %) is sufficient to boost accuracy on test examples with the trigger (red) from 37%percent3737\%37 % to 84%percent8484\%84 % without impacting accuracy on examples without the trigger (blue) by more than 1%percent11\%1 %. Figure 6c shows that the edit suppresses the effect of the trigger at the example level as well, improving correlation between model outputs on examples with and without the trigger from 0.410.410.410.41 (top row) to 0.920.920.920.92 (bottom row). We defer additional details and analyses to Section D.4. Figure 6: Localizing backdoor attacks with Coar-Edit. We edit a ResNet18 model trained on a backdoored CIFAR-10 dataset in which half of all training examples in the “airplane” class contain a planted blue-squared trigger. On the left, we find that the model is sensitive to the trigger—backdoor attacks that add the trigger to examples drop test accuracy from 89%percent8989\%89 % (middle row) to 37%percent3737\%37 % (bottom row). In the center panel, we apply Coar-Edit to identify 25252525 backdoor-specific components that, when ablated, boost accuracy on examples with the trigger (red) from 37%percent3737\%37 % to 84%percent8484\%84 % without impacting accuracy on examples without the trigger (blue). On the right, we find that the edit suppresses sensitivity to the trigger even at the example level—the correlation between model outputs on paired examples with and without the trigger increases from 0.410.410.410.41 to 0.920.920.920.92. Figure 7: Improving robustness to typographic attacks with Coar-Edit. We edit a zero-shot CLIP ViT-B/16 classifier to improve its robustness to typographic attacks [goh2021multimodal]. On the left, we find that predictions on images of household objects (top row) can be manipulated to “taxi”, “twitter”, or “EU” via synthetic (middle row) and real (last row) typographic attacks. In the center panel, we use Coar-Edit to identify components that, when ablated, improve average accuracy on examples with synthetic typographic attacks (red) from 51%percent5151\%51 % to 89%percent8989\%89 % while maintaining accuracy on examples without attacks (blue). Similarly, on the right, we find that the edit transfers robustness to real typographic attacks as well, improving average accuracy from 54%percent5454\%54 % to 86%percent8686\%86 %. 5.5 Improving robustness to typographic attacks Zero-shot CLIP classifiers are vulnerable to typographic attacks [goh2021multimodal] that simply insert text to images in order to induce misclassifications. In this experiment, we evaluate whether our editing approach can improve the robustness of a zero-shot CLIP ViT-B/16 classifier using a dataset [materzynska2022disentangling] comprising 180180180180 images with and without multiple typographic attacks. Specifically, we use Coar-Edit to identify a subset of components that, when ablated, correct the misclassifications induced by synthetic typographic attacks without impacting predictions on images without attacks. We defer additional details to Section D.5. Results. Figure 7 summarizes our findings. In Figure 7a, we show that the predictions of a zero-shot CLIP ViT-B/16 classifier on images of household objects (top row) can be manipulated to “taxi”, “twitter”, or “EU” via synthetic (middle row) or real (last row) typographic attacks. More quantitatively, we find that the zero-shot accuracy on images with synthetic and real typographic attacks drops from 95%percent9595\%95 % to 51%percent5151\%51 % and 54%percent5454\%54 %, respectively. Figure 7b shows that ablating a subset of 300300300300 components (0.4%percent0.40.4\%0.4 %) identified via Coar-Edit improves the accuracy on held-out images with synthetic typographic attacks from 51%percent5151\%51 % to 89%percent8989\%89 % on average (red), without impacting accuracy on images without attacks (dark blue). Furthermore, in Figure 7c, we find that our edit transfers robustness to real typographic attacks as well, improving accuracy on held-out images from 54%percent5454\%54 % to 86%percent8686\%86 % on average. Similar to previous experiments, our approach is sample-efficient in that it only requires 15151515 pairs of target and reference examples with and without synthetic attacks to identify the edit described above. To summarize, simply ablating targeted subsets of components identified via Coar-Edit can induce specific model behavior without requiring additional training. More broadly, our findings highlight how accurate component attribution alone can directly inform model editing. 6 Related work Our work relates to multiple topics in interpretability, which we categorize into works that focus on localizing model behavior, interpreting individual components, editing models, and approximating functional behavior via interpretable proxies. Localizing model behavior. One line of work in mechanistic interpretability attempts to identify “circuits” 2020zoom or “subnetworks” 2021low within neural networks that are responsible for specific capabilities or biases such as factual recall [meng2022locating], gender bias [vig2020investigating], or arithmetic reasoning [stolfo2023understanding]. Building on this, several works focus on automated localization of behavior using techniques such as fine-tuning [panigrahi2023task], activation patching [conmy2023towards, goldowsky2023localizing], or differentiable masking [bayazit2023discovering, de2021sparse, chang2023localization]. More recent studies evaluate these methods in terms of their sensitivity to design choices [zhang2023towards], usefulness for model editing [hase2023does], and ability to characterize functional behavior [wen2023transformers]. Other works develop metrics to quantify the “importance” of individual components (e.g., 2018important,leino2018influence,ghorbani2020neuron), which we adapt as baselines in Section 4. Rather than localizing human-defined behavior to specific components, the component modeling task (Definition 2) aims to explicitly model the collective function mapping component ablations to predictions. Interpreting specific model components. The works discussed above aim to localize specific model behavior to components; another line of work studies the reverse direction, and introduces methods for uncovering the functionality corresponding to a specific model component. Such methods include feature visualization [zeiler2014visualizing, ghiasi2022vision], activation maps [bau2017network, mu2020compositional], ablations [zhou2018revisiting], saliency maps [olah2018building], probing [dalvi2019one, durrani2020analyzing], and natural language descriptions [hernandez2021natural, oikarinen2022clip, bills2023language]. Subsequent analyses use these methods to identify and ascribe meaning to specific model components by labeling them as, e.g., “curve detectors” [cammarata2020curve], “knowledge neurons” [dai2021knowledge], “multimodal neurons” [goh2021multimodal], or “syntax units” [lakretz2019emergence]. More recent work revisits the reliability and robustness of such methods [geirhos2023dont, bolukbasi2021interpretability, hooker2018benchmark, shah2021input, hewitt2019designing, antverg2021pitfalls, huang2023rigorously]. Here, our goal is not to interpret specific model components, but rather to study how all components jointly influence model predictions through the lens of component modeling (Definition 1). Editing model behavior. Another related line of work focuses on model editing, the goal of which is to make small, targeted changes to model representations in order to induce or suppress a specific behavior. Model editing methods include “hypernetworks” [de2021editing, mitchell2021fast], rank-one updates to model parameters [bau2020rewriting, santurkar2021editing, meng2022locating], constrained fine-tuning 2020modifying, and weight interpolation [ilharco2022editing, zou2023representation], among other methods. Recent work has also studied erasing concepts and suppressing spurious correlations from models using layer-wise linear probing [ravfogel2022linear], CLIP-specific methods [gandelsman2023interpreting, chen2023interpreting], and fine-tuning variants [gandikota2023erasing, kirichenko2022last]. In this work, we introduce Coar-Edit to show that effective component attribution can directly enable model editing at the level of examples (§5.1), classes (§5.2), subpopulations (§5.3), and spurious concepts (§5.4, §5.5) by zeroing out targeted subsets of components. Understanding machine learning models by proxy. More generally, our approach connects to a line of research that aims to understand machine learning models by constructing interpretable proxies. For example, feature attribution methods like LIME 2016 approximate a given ML model with a linear model in input space. Similarly, datamodeling 2022datamodels approximates a given learning algorithm by a linear model in “dataset space.” Another line of work analyzes the behavior of deep networks via high-level causal abstractions 2021causal,geiger2023causal, user-specified causal graphs over task-specific variables. 7 Discussion In this section, we discuss connections between component attribution and model editing, discuss directions for future work, and outline key limitations of Coar. Does localization help with model editing? The extent to which localizing specific model behavior to a subset of model components helps with model editing remains contested. On one hand, 2023does show that localizing factual associations in language models does not necessarily help with editing these associations. More broadly, recent evaluation studies suggest that model editing can degrade robustness to distribution shifts [brown2023robustness] and may not modify model behavior in a consistent manner [cohen2023evaluating]. On the other hand, recent work shows that some localization methods can in fact recover ground-truth localization in controlled settings [chang2023localization] and improve calibration of fine-tuned language models [panigrahi2023task]. Our findings in Section 5 substantiate the latter view, as Coar-Edit uses component attributions to identify a target subset of components that, when ablated, modify model behavior as intended. Future work. We highlight three directions that, while outside the scope of this work, may be interesting avenues for future work. • Analyzing neural network representations. An interesting direction for future work could be to use component attribution (and component models, more generally) to study empirically documented phenomena in deep learning. There are a plethora of questions to ask here which, although beyond the scope of this work, are natural applications of component attributions. For example, extending our results from Section 5.1, can we use component attribution to better isolate “opposing signals” 2023outliers for a given task, and to understand their role in shaping model predictions? Can we use component attributions to study how model predictions change due to adversarial perturbations 2015explaining, or over the course of training 2019sgd? Similarly, can we develop improved methods for localizing memorized inputs to specific model components 2020neural,maini2023can? Given that component attributions are causally meaningful, can we use them as a kernel with which to compare different models 2019similarity or learning algorithms 2023modeldiff? • Attributing generative models. While we focus on vision models in this work, Coar is a general method that can estimate component attributions for any machine learning model. Future work might thus explore possible model output functions (and their corresponding component attributions) for generative models. For diffusion-based generative models, one might study the denoising error for a fixed timestep, as in 2023journey,zheng2023intriguing. For language models, a possible point of start (following 2023trak) would be to use the average correct-class margin (5) of a sequence of tokens as the model output function. Our preliminary experiments in Appendix B show that Coar learns accurate component attributions for language models such as GPT-2 [radford2019language] and Phi-2 [li2023textbooks]. In general, estimating and applying component attributions for generative models is a promising avenue for future work. • Beyond linear component attribution. The fact that component attributions’ predictiveness decreases on out-of-distribution component subsets, i.e., when αtest≠αtrainsubscripttestsubscripttrain _test≠ _trainαtest ≠ αtrain, suggests that the linear form of component attributions might not be expressive enough to fully capture the map between model components and outputs. Given the generality of Coar, an interesting avenue for future work would be to explore whether non-linear component models such as decision trees or kernel methods predict component counterfactuals more accurately, and as a result, improve model editing. Limitations. Our attribution-based approach for decomposing and editing model predictions is not without its limitations. First, estimating Coar attributions involves constructing datasets of ground-truth counterfactuals (Equation 2), which can require a large number of forward passes through the model. In Section C.5, we show that simply using 2222-5×5×5 × fewer samples can significantly speed up Coar without impacting the quality of the resulting attributions. Mitigating this computational bottleneck further through better sampling or approximation techniques is an interesting avenue for future work. Second, Coar requires specifying an ablation method (Equation 1). While we use zero ablations in this work (Remark 2), one could also use Coar with ablation methods (e.g., 2022causal) that account for distribution shifts in activation space. For example, in Section C.4, we show that Coar yields predictive attributions with another ablation method that simply scales down the activations of ablated components by a small constant factor instead of setting them to zero. Finally, as noted in Remark 4, our model editing approach Coar is coarse-grained in that it modifies model behavior by simply ablating a targeted subset of components. Using Coar in conjunction with gradient-based editing techniques in order to make finer-grained model edits is an interesting avenue for future work. 8 Conclusion We first formalize the problem of decomposing model predictions in terms of model components through the component modeling task. We specifically focus on a special case of component modeling, component attribution, where the goal is to predict the counterfactual impact of every component on a given prediction. We then propose Coar, a scalable method for estimating predictive component attributions, and demonstrate its effectiveness across model architectures, datasets, and tasks. Finally, through a series of five experiments, we also stress-test the utility of Coar attributions in directly editing model behavior without requiring additional training. Acknowledgements The authors would like to thank Benjamin Cohen-Wang, Logan Engstrom, Alaa Khaddaj, and Kristian Georgiev for helpful discussions and comments on an earlier draft of this manuscript. Work supported in part by the NSF grant DMS-2134108. This material is based upon work supported by the Defense Advanced Research Projects Agency (DARPA) under Contract No. HR001120C0015. Appendix A Evaluation setup In this section, we outline the experiment setup—datasets, models, baselines, implementation details—used in Section 4 to evaluate whether Coar attributions can accurately estimate ground-truth component counterfactuals. A.1 Pseudocode Algorithm 1 An outline for estimating component attributions with Coar. 1:procedure Coar(example z, model M with output function fMsubscriptf_Mfitalic_M and components C, ablation frac. α) 2: Set D(z)←[]←superscriptD^(z)←[]D( z ) ← [ ] ▷ ▷ initialize component dataset (2) 3: for i∈1,…,m1…i∈\1,…,m\i ∈ 1 , … , m do ▷ ▷ m denotes dataset size 4: Sample a subset Ci⊂CsubscriptC_i⊂ CCitalic_i ⊂ C from CsubscriptD_CDitalic_C where |Ci|=α⋅|C|subscript⋅|C_i|=α·|C|| Citalic_i | = α ⋅ | C | 5: Set yi←fM⁢(z,Ci)←subscriptsubscriptsubscripty_i← f_M(z,C_i)yitalic_i ← fitalic_M ( z , Citalic_i ) ▷ ▷ compute component counterfactual (1) 6: Define Ci∈0,1|C| 0_C_i∈\0,1\^|C|0italic_C start_POSTSUBSCRIPT i end_POSTSUBSCRIPT ∈ 0 , 1 | C | as (Ci)j=0subscriptsubscript0subscript0( 0_C_i)_j=0( 0italic_C start_POSTSUBSCRIPT i end_POSTSUBSCRIPT )j = 0 if cj∈Cisubscriptsubscriptc_j∈ C_icitalic_j ∈ Citalic_i else 1111 7: Update D(z)←D(z)+[(Ci,yi)]←superscriptsuperscriptdelimited-[]subscript0subscriptsubscriptD^(z)← D^(z)+[( 0_C_i,y_i)]D( z ) ← D( z ) + [ ( 0italic_C start_POSTSUBSCRIPT i end_POSTSUBSCRIPT , yitalic_i ) ] ▷ ▷ update component dataset 8: θ(z),b(z)←superscriptsuperscriptabsentθ^(z),b^(z)←θ( z ) , b( z ) ← LinearRegression(D(z)superscriptD^(z)D( z )) ▷ ▷ estimate component attributions via Equation 3 9: return θ(z),b(z)superscriptsuperscriptθ^(z),b^(z)θ( z ) , b( z ) A.2 Datasets and models. We now outline the datasets and models used to evaluate Coar (§4) and Coar-Edit (§5). CIFAR-10. We use the standard CIFAR-10 2009learning image classification dataset to evaluate Coar attributions (Section 4, Section C.2) and for an editing task (Section 5.4). We train ResNet, ViT, and MLP models that attain test accuracies of 91%percent9191\%91 %, 83%percent8383\%83 % and 56%percent5656\%56 % respectively. We specify a computation graph over 2,34423442,3442 , 344 components for the ResNet-18 model, 31,7283172831,72831 , 728 components for the ViT model, and 3,07230723,0723 , 072 components for the MLP model. Each component in the ResNet-18 model corresponds to a convolution filter. Similarly, each component in the ViT and MLP models corresponds to a neuron. ImageNet. We use the standard ImageNet 2009imagenet image classification dataset to evaluate Coar attributions in Section 4 and for editing tasks in Section 5.2. We use ImageNet-Sketch 2019learning and five random shifts from ImageNet⋆ ⋆ 2023dataset—“in the water”, “at dusk simple”, “orange”, “pencil sketch”, “green”— to evaluate the out-of-distribution performance of edited ImageNet models in Section 5.2. We use the pre-trained ResNet50 and ViT-B/16 models222Model and pre-trained weights taken from torchvision: https://pytorch.org/vision/stable/models.html that attain test accuracies of 75.4%percent75.475.4\%75.4 % and 80.7%percent80.780.7\%80.7 % respectively. For the ResNet-50 model, we specify a computation graph over 22,7202272022,72022 , 720 components, each corresponding to a convolution filter. Similarly, for the ViT-B/16 model, we specify a computation graph over 82,9448294482,94482 , 944 components, each corresponding to a neuron. Waterbirds. The Waterbirds dataset [sagawa2020distributionally] comprises images of birds taken from the CUB dataset [wah2011caltech] and pasted on backgrounds from the Places dataset [zhou2017places]. The task here is to classify “waterbirds” and “landbirds” in the presence of spurious correlated “land” and “water” backgrounds in the training dataset. 2020distributionally introduce Waterbirds as a benchmark to improve model performance under subpopulation shifts induced by spurious correlations. We use this dataset to evaluate whether Coar-Edit can improve subpopulation robustness via model editing. In this experiment, we fine-tune an ImageNet ResNet50 model and use a computation graph over 22,7202272022,72022 , 720 components, each corresponding to a convolution filter. CelebA. The CelebA dataset [li2020celeb] comprises images of celebrities with binary attributes such as “smiling”, “wearing hat”, “wearing lipstick”, etc. Similar to previous work on subpopulation robustness (e.g., [sagawa2020distributionally]), we repurpose CelebA as a binary classification task where the goal is to predict whether a person in a given image has blond hair. The attributes “hair color” and “gender” are spuriously correlated in the training dataset, resulting in models that latch on to a “gender → blond hair” shortcut and underperform on the “blond males” subpopulation. Similar to the Waterbirds setting, we fine-tune an ImageNet ResNet50 model and specify a computation graph over 22,7202272022,72022 , 720 components, each corresponding to a convolution filter. Typographic attacks dataset. We use a dataset of typographic attacks [materzynska2022disentangling] for an editing task in Section 5.5. This dataset comprises 180180180180 images of household objects with and without eight typographic attacks such as “taxi”, “twitter”, “EU”, and “iPad”. We visualize some examples from this dataset in Figure 7. Our experiment in Section 5.5 uses this dataset along with a zero-shot CLIP ViT-B/16 classifier [radford2021learning]. For this model, we specify a computation graph over all 82,9448294482,94482 , 944 components, corresponding to the set of all weight vectors (individual rows in weight matrices) in all self-attention and MLP modules. See Section D.5 for more details. TinyStories. We use the TinyStories dataset [eldan2023tinystories] to evaluate Coar attributions over the GPT-2 language model (Appendix B). This dataset contains short stories synthetically generated by GPT-3.5 and GPT-4. To compute component attributions for GPT-2, we specify a computation graph over 64,5126451264,51264 , 512 components, which correspond to the set of all weight vectors, i.e., in every self-attention module and feed-forward module of the model. See Section B.1 for experiment details and findings. BoolQ. We use the BoolQ dataset [clark2019boolq] to evaluate Coar attributions for the Phi-2 model [li2023textbooks]. Each example in this dataset comprises a passage of text, a question, and a binary answer. We evaluate the zero-shot performance of Phi-2 using the prompting and evaluation procedure from 2023framework333https://github.com/EleutherAI/lm-evaluation-harness/. Given the size of the Phi-2 model, we specify a computation graph over 55,5525555255,55255 , 552 components, each corresponding to a contiguous block of 10101010 weight vectors in every self-attention module and feed-forward module of the model. See Section B.2 for experiment details and findings. A.3 Baselines In Section 4, we compare Coar against four baseline methods for estimating component attributions: Leave-One-Out (LOO), Gradient-times-parameters (GP), Neuron Conductance (NC), and Internal Influence (I). Each baseline computes an attribution vector ()∈ℝ|C|superscriptsuperscriptℝ w^(z) ^|C|italic_wbold_( italic_z bold_) ∈ blackboard_R| C | for a given example z by assigning an “importance” score wj(z)subscriptsuperscript w^(z)_jw( z )j to each component cj∈Csubscriptc_j∈ Ccitalic_j ∈ C. Then, as per Equation 4, we estimate a component counterfactual fM⁢(z,C′)subscriptsuperscript′f_M(z,C )fitalic_M ( z , C′ ) as the sum of importance scores of components in C∖C′ C C ∖ C′, i.e., scores of components that are not ablated. We describe each baseline in more detail below: • Leave-One-Out (LOO): This method ablates each component cj∈Csubscriptc_j∈ Ccitalic_j ∈ C and sets the coefficient θj(z)subscriptsuperscript θ^(z)_jθ( z )j to the change in model output fM⁢(z)subscriptf_M(z)fitalic_M ( z ) before and after ablation: wj(z)=fM⁢(z,cj)−fM⁢(z,∅)subscriptsuperscriptsubscriptsubscriptsubscript w^(z)_j=f_M(z,\c_j\)-f_M(z, )w( z )j = fitalic_M ( z , citalic_j ) - fitalic_M ( z , ∅ ) • Gradient-times-Parameters (GP): This method approximates the leave-one-out estimate described above. Specifically, it estimates the leave-one-out effect of each component cj∈Csubscriptc_j∈ Ccitalic_j ∈ C using a first-order Taylor approximation of fM⁢(z,cj)subscriptsubscriptf_M(z,\c_j\)fitalic_M ( z , citalic_j ) around fM⁢(z,∅)subscriptf_M(z, )fitalic_M ( z , ∅ ): wj(z)=∇cjfM⁢(z,∅)⋅δcjsubscriptsuperscript⋅subscript∇subscriptsubscriptsubscriptsubscript w^(z)_j= _c_jf_M(z, )· _c_jw( z )j = ∇c start_POSTSUBSCRIPT j end_POSTSUBSCRIPT fitalic_M ( z , ∅ ) ⋅ δitalic_c start_POSTSUBSCRIPT j end_POSTSUBSCRIPT where δcjsubscriptsubscript _c_jδitalic_c start_POSTSUBSCRIPT j end_POSTSUBSCRIPT is the parameter-space change in cjsubscriptc_jcitalic_j induced by the ablation method of choice. • Neuron Conductance (NC) [dhamdhere2018important]: This method extends the Integrated Gradients method [sundararajan2017axiomatic]—an input-space feature attribution method—to compute importance scores for each component cj∈Csubscriptc_j∈ Ccitalic_j ∈ C. Intuitively, NC modifies the computation in Integrated Gradients in order to quantify the “flow” through each component cj∈Csubscriptc_j∈ Ccitalic_j ∈ C. See Equation 3 in [dhamdhere2018important] for a formal description. • Internal Influence (I) [leino2018influence]: Similar to NC, this method also adapts Integrated Gradients [sundararajan2017axiomatic] to compute importance scores. At a high level, I directly applies Integrated Gradients to layerwise activations by treating the output of each layer as an input to subsequent layers. See Definition 1 in [leino2018influence] for a formal description. We implement the first two baselines (LOO and GP) from scratch444Our code is available at https://github.com/MadryLab/modelcomponents and use the captum library [kokhlikyan2020captum] 555Github repository: https://github.com/pytorch/captum to implement NC and I. As per Definition 2, we estimate the component counterfactual fM⁢(z,C′)subscriptsuperscript′f_M(z,C )fitalic_M ( z , C′ ) using these baselines by setting the bias term b(z)superscriptb^(z)b( z ) to zero and taking the sum over attribution scores of components that are not ablated. A.4 Implementation details Sample size for component attribution estimation. The computational cost of our approach linearly scales with the sample size m used to estimate component attributions (see Algorithm 1). Each sample in the component dataset D(z)superscriptD^(z)D( z ) corresponds to a single forward pass through the model M in order to compute the counterfactual fM⁢(z,C′)subscriptsuperscript′f_M(z,C )fitalic_M ( z , C′ ) (1), i.e., model output fM⁢(z)subscriptf_M(z)fitalic_M ( z ) after ablating a subset of components C′⊂Csuperscript′C ⊂ C′ ⊂ C. The setups A,B,CABC\A,B,C\ A , B , C considered in Section 4 use sample size m=50000,100000,20000050000100000200000m=\50000,100000,200000\m = 50000 , 100000 , 200000 respectively. In Section C.5, we show that the sample size m used in Section 4 can be reduced by 2222-5×5×5 ×, resulting in a direct speedup while only reducing the predictive power of Coar attributions by a small amount. Data loading. We use the FFCV library666Github repository: https://github.com/libffcv/ffcv 2022ffcv to train and evaluate models. FFCV removes the data loading bottleneck for small models, gives a 3-4×3-4×3 - 4 × improvement in throughput compared to standard PyTorch data loading. Speeding up regression. The second step of Coar—fitting component attributions to the component dataset (2)—requires solving a linear regression problem (Equation 3) for each example z. We parallelize this step by using the fast-l1 package777Github repository: https://github.com/MadryLab/fast_l1, a SAGA-based GPU solver for linear regression. Computing resources. We train our models and compute Coar attributions on a cluster of machines, each with 9999 NVIDIA A100 or V100 GPUs and 96969696 CPU cores. We also use half-precision to increase training speed. Appendix B Applying Coar to language models In Section 4 and Appendix C, we showed that our proposed method Coar attributions accurately estimate component counterfactuals (1) on large-scale vision tasks across several datasets and model architectures. In this section, we apply Coar to language models. Specifically, we consider two experiments: (a) GPT-2 2019language evaluated on the next-token prediction task and (b) Phi-2 2023textbooks evaluated on a zero-shot classification task. In both cases, we show that Coar attributions accurately predict how model outputs change in response to component ablations. B.1 Evaluating GPT-2 on the TinyStories dataset Task and model output function. We apply Coar to the next-token prediction task. Following 2023trak, we interpret this task as a sequence as a v-way classification problem, where v is the vocabulary size, and set the model output function to be the average correct-class margin (5) over all tokens in a given sequence. Model and dataset. In this experiment, we consider the GPT-2 model888https://huggingface.co/gpt2 2019language, with a computation graph over 64,5126451264,51264 , 512 components. These components correspond to the set of weight vectors in every self-attention module and feed-forward module in the model. We evaluate model performance on the next-token prediction task using the TinyStories dataset999https://huggingface.co/datasets/roneneldan/TinyStories 2023tinystories, where each sequence corresponds to a synthetically generated short story. Computing Coar attributions. We apply Coar (without any modifications) to compute component attributions for a random subset of 1000100010001000 examples in the TinyStories validation set using a component dataset of 200,000200000200,000200 , 000 component counterfactuals (2) and a ablation fraction of α=2.5%percent2.5α=2.5\%α = 2.5 %. Evaluating Coar attributions. Similar to the results in Section 4, Coar attributions are predictive in the language modeling setting as well. Specifically, these attributions accurately predict the effect of ablating components on the average correct-class margin of GPT-2 on examples from the TinyStories validation set. In Figure 8a, we pick a random example z from the TinyStories validation set and compute the correlation between ground-truth component counterfactuals fM⁢(z,⋅)subscript⋅f_M(z,·)fitalic_M ( z , ⋅ ) and the corresponding estimate (4) using its Coar attributions ()superscript θ^(z)italic_θbold_( italic_z bold_), as defined in Equation 6. In Figure 8b, we plot a histogram over example-level correlations of 1000100010001000 examples and find that Coar attributions attain an average correlation of 0.83,0.85,0.890.830.850.89\0.83,0.85,0.89\ 0.83 , 0.85 , 0.89 with ground-truth component counterfactuals sampled using ablation fraction α=5%,2.5%,1%percent5percent2.5percent1α=\5\%,2.5\%,1\%\α = 5 % , 2.5 % , 1 % respectively. B.2 Evaluating Phi-2 on the BoolQ dataset Task and model output function. We now turn to a reading comprehension task, where the goal is to answer a question given a passage of text. We evaluate this classification task in a zero-shot manner: the language model is prompted with a passage of text and a question, and the goal is to output the correct answer from yes, no. Like in vision tasks (Section 4), we use the correct-class margin (5) as the model output function for this zero-shot binary classification task. Model and dataset. We consider the Phi-2 model101010https://huggingface.co/microsoft/phi-2 2023textbooks and specify a computation graph over 55,5525555255,55255 , 552 components. Here, each component corresponds to a contiguous block of 10101010 weight vectors in the model. We evaluate this model on the BoolQ dataset111111https://huggingface.co/datasets/google/boolq 2019boolq, where each example consists of a passage of text, a question, and a binary yes, no answer. Using the prompting and evaluation procedure from the 2023framework121212https://github.com/EleutherAI/lm-evaluation-harness, Phi-2 attains an 83.6%percent83.683.6\%83.6 % accuracy on this task. Computing Coar attributions. Like in Section B.1, we apply compute Coar attributions for a random subset of 500500500500 examples in the BoolQ validation set using a component dataset of m=100,000100000m=100,000m = 100 , 000 component counterfactuals (2) and a ablation fraction of α=0.0250.025α=0.025α = 0.025. Evaluating Coar attributions. We find that Coar attributions are predictive of unseen component counterfactuals on this task as well. Figure 9a plots the correlation between ground-truth component counterfactuals fM⁢(z,⋅)subscript⋅f_M(z,·)fitalic_M ( z , ⋅ ) and the corresponding Coar estimate (4) of a random BoolQ example z. The histograms in Figure 9b show that Coar attributions attain an average correlation of 0.58,0.66,0.660.580.660.66\0.58,0.66,0.66\ 0.58 , 0.66 , 0.66 with ground-truth component counterfactuals sampled using ablation fraction α=5%,2.5%,1%percent5percent2.5percent1α=\5\%,2.5\%,1\%\α = 5 % , 2.5 % , 1 % respectively. Figure 8: Evaluating Coar on GPT-2. We apply Coar to the GPT-2 model 2019language on the TinyStories dataset 2023tinystories. The resulting component attributions are predictive of component counterfactuals. The left plot shows that component attributions can estimate the effect of ablating components on the average correct-class margin (over tokens in a sequence) of GPT-2 on a random TinyStories example with high correlation. The histograms in the right plot show that Coar attributions attain high average correlation for multiple values of ablation fraction α. Figure 9: Evaluating Coar on Phi-2. We apply Coar to the Phi-2 model 2023phi on the BoolQ dataset 2019boolq. The resulting component attributions are predictive of component counterfactuals. The left plot shows that component attributions can estimate the effect of ablating components on the average correct-class margin of Phi-2 on a random BoolQ example with high correlation. The histograms in the right plot show that Coar attributions attain high average correlation for multiple values of ablation fraction α. Appendix C Additional evaluation of Coar In this section, we first show that Coar learns accurate component attributions on additional datasets, model architectures, and tasks (Sections C.1, C.2 and C.3). This supplements our findings in Section 4, where we showed that Coar learns component attributions that accurately predict component counterfactuals (1) on three image classification setups: CIFAR-10 ResNet-18, ImageNet ResNet-50, and ImageNet ViT-B/16. Then, we show that Coar attributions retain its predictive power when estimated with fewer samples (Section C.5) or with different ablation fractions (Section C.4). Finally, we supplement our example-level evaluation of Coar attributions in Section 4 with additional example-level comparisons of ground-truth component counterfactuals and attribution-based estimates (Section C.6). C.1 Evaluating Coar on additional datasets Our experiments in Section 4 evaluated the predictiveness of Coar attributions corresponding to in-distribution test examples from the CIFAR-10 and ImageNet datasets. Now, we show that Coar attributions remain predictive on training examples as well as out-of-distribution examples. Specifically, we apply Coar to compute attributions of ResNet-18 predictions on the CIFAR-10 training set and on six corrupted versions of the CIFAR-10 test set [hendrycks2019benchmarking]. as shown in Figure 10, Coar attributions exhibit high correlation on average (between 0.60.60.60.6 and 0.80.80.80.8) depending on the ablation fraction α used to ablate random α-fraction sized components subsets. Note that the correlation is maximum when α=0.050.05α=0.05α = 0.05 because the component attributins are estimated with the same ablation fraction, i.e., αtrain=0.05subscripttrain0.05 _train=0.05αtrain = 0.05. Figure 10: Do Coar attributions generalize to out-of-distribution examples? Coar attributions remain predictive on the CIFAR-10 training set and on six corrupted versions of the CIFAR-10 test set [hendrycks2019benchmarking] over a range of ablation fractions α. See Section C.1 for more details. C.2 Evaluating Coar on additional model architectures Recall that Coar is model-agnostic in that it is not tied to any specific model architecture. In Section 4, we applied Coar to ResNets trained on CIFAR-10 and ImageNet and a ViT-B/16 model trained on ImageNet. In this section, we apply Coar to two additional model architectures: a ViT model trained on CIFAR-10 (83%percent8383\%83 % accuracy) and a one-layer fully-connected network trained on CIFAR-10 (56%percent5656\%56 % accuracy). Figure 11 shows that Coar attributions on both architectures yield accurate estimates of how model outputs change in response to ablating random α-fraction sized components subsets, with correlation 0.650.650.650.65 and 0.850.850.850.85 for the ViT and MLP models when α=αtrainsubscripttrainα= _trainα = αtrain respectively. Figure 11: Do Coar attributions generalize to other model architectures? Coar attributions yield accurate estimates of component counterfactuals on two additional model architectures: a ViT-based model (left) and a one-layer fully-connected network (right) trained on CIFAR-10. See Section C.2 for more details. C.3 Evaluating Coar on additional tasks We now evaluate Coar attributions on four additional tasks: • First, we apply Coar to pre-trained ImageNet ResNet50 model fine-tuned on two datasets—Waterbirds and CelebA—that we use in Section 5.3—see first row of Figure 12. We find that Coar attributions are predictive on both datasets, attaining higher correlation with ground-truth component counterfactuals when α is closer to αtrain=0.05subscripttrain0.05 _train=0.05αtrain = 0.05. • Second, we apply Coar to a pre-trained ImageNet ResNet50 model fine-tuned on MIMIC-CXR [johnson2019mimic], a dataset of labeled chest radiographs. In this case, we set the model output function to be the logit of the “Cardiomegaly” class instead of correct-class margin that we use in Section 4. Figure 12 shows that Coar attributions attain a correlation of 0.70.70.70.7 and 0.60.60.60.6 with ground-truth logits when α=αtrain=0.05subscripttrain0.05α= _train=0.05α = αtrain = 0.05 and α=0.100.10α=0.10α = 0.10 respectively. • The fourth plot in Figure 12 corresponds to the CLIP setting considered in Section 5. In this setting, we take the zero-shot CLIP ViT-B/16 classifier and evaluate it on a dataset of images with and without typographic attacks [materzynska2022disentangling]. As shown in the plot, the correlation between Coar attributions and ground-truth margins is close to 0.70.70.70.7 when α=αtrain=0.03subscripttrain0.03α= _train=0.03α = αtrain = 0.03, i.e., ablating 3%percent33\%3 % of the components in the CLIP model. Figure 12: Evaluating Coar attributions on additional tasks. We find that component attributions estimated using Coar are predictive on four additional tasks: fine-tuning ImageNet ResNet50 on Waterbirds, CelebA and MIMIC-CXR, and a zero-shot CLIP ViT-B/16 classification task on a dataset containing typographic attacks (Section 5.5). Note that the MIMIC-CXR setting uses the logit of the “Cardiomegaly” class as the model output function. See Section C.3 for additional information about these tasks. C.4 Comparing Coar attributions estimated with different ablation fractions We now analyze how changing the ablation fraction αtrainsubscripttrain _trainαtrain used to fit Coar attributions affects their predictiveness over different ablation fractions at test time. Specifically, we consider the ImageNet ResNet-50 setting from Section 4 and compute two sets of Coar attributions, corresponding to two values of αtrainsubscripttrain _trainαtrain: 0.050.050.050.05 and 0.100.100.100.10. Then, for each of these two sets of attributions, we evaluate its correlation with ground-truth component counterfactuals over a range of ablation fractions α. As shown in Figure 13, the correlation “profile” over α depends on the value of αtrainsubscripttrain _trainαtrain used to fit the attributions. When α is small, the correlation is higher for attributions estimated with αtrain=0.05subscripttrain0.05 _train=0.05αtrain = 0.05. Analogously, when α is large, the correlation is higher for attributions estimated with αtrain=0.10subscripttrain0.10 _train=0.10αtrain = 0.10. This is because the component attributions fare better as counterfactual predictors on component counterfactuals that are “similar” to the ones used to fit them—i.e., when αtest≈αtrainsubscripttestsubscripttrain _test≈ _trainαtest ≈ αtrain. C.5 Comparing Coar attributions estimated with different sample sizes In Section 4, we computed Coar attributions using sample sizes m=5000050000m=50000m = 50000 for the ResNet-18 model trained on CIFAR-10 and m=100000100000m=100000m = 100000 for the ResNet-50 model trained on ImageNet. Recall that the sample size m here corresponds to the number of component counterfactuals used to fit the component attributions. In this section, we vary the sample size m and show that Coar attributions remain predictive even when trained on k×k×k × fewer examples, where k∈2,5,102510k∈\2,5,10\k ∈ 2 , 5 , 10 . Specifically, the left column of Figure 14 shows that Coar attributions estimated on CIFAR-10 and ImageNet data with sample size m and m/km/km / k have high cosine similarity on average, with the similarity increasing as k decreases. The right column of Figure 14 shows that decreasing the sample size m by a factor of k∈2,5,102510k∈\2,5,10\k ∈ 2 , 5 , 10 does not significantly impact the correlation between Coar attributions and ground-truth component counterfactuals. For example, reducing the sample size by 5×5×5 × only reduces the correlation from 0.70.70.70.7 to 0.650.650.650.65 in the CIFAR-10 ResNet-18 setting. Additionally, we observe that Coar attributions fare better than attributions estimated with the best-performing baseline (LOO) even when trained on 10×10×10 × fewer examples on CIFAR-10 and 5×5×5 × fewer examples on ImageNet. Figure 13: Comparing Coar attributions estimated with different ablation fractions α. Coar attributions estimated with different ablation fractions αtrainsubscripttrain _trainαtrain attain a different correlation “profile” over α at test time. The correlation between ground-truth component counterfactuals and attribution-based estimates is higher for attributions estimated with αtrain=0.05subscripttrain0.05 _train=0.05αtrain = 0.05 when α is small, and higher for attributions estimated with αtrain=0.10subscripttrain0.10 _train=0.10αtrain = 0.10 when α is large. This empirically shows that Coar attributions are more predictive on component counterfactuals that are “similar” to the ones used to fit them—i.e., when αtest≈αtrainsubscripttestsubscripttrain _test≈ _trainαtest ≈ αtrain. See Section C.4 for more details. Figure 14: Comparing Coar attributions estimated with different sample sizes. Coar attributions for CIFAR-10 ResNet-18 and ImageNet ResNet-50 (Setup A and B respectively in Section 4) estimated with smaller sample sizes m are still predictive of component counterfactuals. On the left, we show that Coar attributions estimated with sample size m and m/km/km / k have high cosine similarity on average, with the similarity increasing as k decreases. On the right, we show that decreasing the sample size m by a factor of k∈2,5,102510k∈\2,5,10\k ∈ 2 , 5 , 10 does not significantly affect the correlation between Coar attributions and ground-truth component counterfactuals. In particular, Coar outperforms the best-performing baseline (LOO) even with 10×10×10 × fewer samples on CIFAR-10 (top row) and 5×5×5 × fewer samples on ImageNet (bottom row). C.6 Analyzing Coar attributions at the example level To supplement our evaluation in Section 4, we provide additional example-level scatterplot comparisons between ground-truth component counterfactuals and the corresponding estimates obtained using component attributions estimated with Coar and all baselines from Section 4. We plot these comparisons on CIFAR-10 examples in Figure 15 and on ImageNet examples in Figure 16. Our findings further substantiate that Coar attributions exhibit higher correlation with ground-truth component counterfactuals than all four baseliens on both CIFAR-10 and ImageNet. Figure 15: Additional example-level evaluation of component attributions on CIFAR-10. Each row corresponds to a different example z randomly picked from the CIFAR-10 test set and each column corresponds to a different attribution method. The left-most subfigure in each row shows that Coar attributions and the corresponding ground-truth component counterfactuals exhibit high correlation on example z. In comparison, the other subfigures in each row, one for baseline method, consistently exhibit lower correlation. See Section C.5 for more details. Figure 16: Additional example-level evaluation of component attributions on ImageNet. Similar to the results in Figure 15, each row corresponds to a different example z randomly picked from the ImageNet test set. The left-most subfigure in each row shows that Coar attributions and the corresponding ground-truth component counterfactuals exhibit high correlation on example z. In comparison, the other subfigures in each row, corresponding to a baseline method, consistently exhibit worse correlation. See Section C.5 for more details. C.7 Qualitatively analyzing Coar attributions We qualitatively analyze Coar attributions using two visualization techniques: Visualizing component-specific attributions across examples. Given examples z1,…,znsubscript1…subscript\z_1,…,z_n\ z1 , … , zitalic_n with corresponding component attributions θ(z1),…,θ(zn)superscriptsubscript1…superscriptsubscript\θ^(z_1),…,θ^(z_n)\ θ( z1 ) , … , θ( zitalic_n ) , we analyze how the attribution estimates of individual components vary across the set of examples. Specifically, for a component ci∈Csubscriptc_i∈ Ccitalic_i ∈ C, we visualize the examples with the most positive attribution values θi(z)subscriptsuperscript θ^(z)_iθ( z )i for component cisubscriptc_icitalic_i. In this experiment, we visualize a random subset of components from the ImageNet ResNet-50 model (setup B in Section 4). As shown in Figure 17, the examples with the most positive attributions for a given component exhibit high visual similarity at different levels of granularity: • The first, third and fifth row in Figure 17 show that the examples with the most positive attributions for layer4.0.conv3[477] and layer4.2.conv3[53] contain purple flowers, watch faces, and glass-shaped objects respectively. • However, consistent with recent work on superposition in deep networks [elhage2022toy], we observe that some components such as layer4.2.conv2[336] in the second row as well as layer3.1.conv3[655] in the last row can surface dissimilar subsets of examples and do not readily map to a single semantic concept. Visualizing nearest neighbors in attribution space. We also use component attributions as feature embeddings in order to visualize the nearest neighbors of a given example in “component attribution” space. Intuitively, this technique allows us to identify examples on which model outputs change similarly in response to component ablations. In this experiment, we visualize a random subset of examples from the CelebA dataset along with their 5555 nearest neighbors using Coar attributions of a fine-tuned ImageNet ResNet-50 model. Figure 18 shows that the nearest neighbors of a given example in attribution space high visual similarity, i.e., similar facial attributes such as background (first row), hair color (second and fourth row), accessories (third row), or even the same person in different poses (last row). Figure 17: Visualizing component-specific attributions across examples. We sample a random set of components from the ImageNet ResNet-50 model (setup B in Section 4) and visualize the examples with the most positive attributions for each component. In general, the examples with the most positive attributions for a given component exhibit visual similarity at different levels of granularity. For example, the first, third and fifth row in Figure 17 show that the examples with the most positive attributions for layer4.0.conv3[477] and layer4.2.conv3[53] contain purple flowers, watch faces, and glass-shaped objects respectively. However, consistent with recent work on superposition in deep networks [elhage2022toy], we observe that some components such as layer4.2.conv2[336] (second row) and layer3.1.conv3[655] (last row) can surface dissimilar subsets of examples or do not readily map to a single semantic concept. Figure 18: Visualizing nearest neighbors in Coar attribution space. We also use component attributions as feature embeddings in order to visualize the five nearest neighbors of examples from the CelebA dataset in “component attribution” space. Intuitively, this technique allows us to identify examples on which model outputs change similarly in response to component ablations. In general, we observe that the nearest neighbors of a given example in attribution space high visual similarity, e.g, similar facial attributes such as background (first row), hair color (second and fourth row), accessories (third row), or even the same person in different poses (last row). Appendix D Additional evaluation of Coar-Edit We use Coar-Edit in five different editing tasks: correcting misclassifications (§5.1); forgetting a class (§5.2); improving subpopulation robustness (§5.3); localizing backdoor attacks (§5.4); and improving robustness to typographic attacks (§5.5). In this section, we provide additional details and/or supplementary experiments for each task. D.1 Editing individual predictions Experiment details. In Section 5.1, we use Coar-Edit to correct misclassifications of a ResNet-50 model on ImageNet examples. In this experiment, we set the “target” example to be a misclassified ImageNet example and the “reference” example to a set of 50505050 randomly selected ImageNet examples. Then, we use these examples to identify and ablate components (9) that increase the correct-class margin (5) of the target example without impacting the average margin over the reference examples. Additional experiments. We first show that Coar-Edit is not sensitive to the choice of misclassified examples, model, or dataset. In Figure 20, we reproduce the experiment in Section 5.1 on additional ImageNet examples misclassified by a ResNet-50 model. In Figure 19, we use Coar-Edit to similarly fix misclassifications of a ResNet-18 model on the CIFAR-10 dataset. In Figure 21, we show that Coar-Edit can also be used to adversarially induce misclassifications on ImageNet examples by ablating the top-k components corresponding to the “target” example. Similar to our findings in Section 5.1, we observe that ablating a few components via Coar-Edit is sufficient to change the individual example-level prediction without changing overall model performance. Additional analysis. Which components does Coar-Edit ablate to correct misclassifications? To answer this question, we first aggregate all components ablated by Coar-Edit in order to (individually) correct ImageNet examples misclassified by a ResNet-50 model. Then, we plot the most common convolution layers corresponding to these ablated components in Figure 22. We find that Coar-Edit primarily targets convolution filters from the last few layers (closet to the output) of the ResNet-50 model in order to make fine-grained edits that do not impact overall model performance. For example, more than 25%percent2525\%25 % of the ablated components belong to layer4.0,1,2.conv3—the last convolution layer in the first three residual blocks of the last layer group of the ResNet-50 model. Figure 19: Correcting misclassified CIFAR-10 examples via Coar-Edit. We reproduce the Coar-Edit experiment from Section 5.1 on the CIFAR-10 dataset. Specifically, each row corresponds to CIFAR-10 test example that is misclassified by a ResNet-18 model. The left subplot in each row shows how applying Coar-Edit (by ablating components (9)) increases the correct-class margin (5) of the misclassified example without impacting the average margin over the train or test set. The right subplot reports the drop in overall test accuracy and visualizes examples with correct-class margins that change the most or least due to the edit. Figure 20: Correcting misclassified ImageNet examples via Coar-Edit. We reproduce the Coar-Edit experiment from Section 5.1 on additional ImageNet examples (one per row) misclassified by a ResNet-50 model. The left subplot shows that applying Coar-Edit (by ablating components (9)) increases the correct-class margin (5) of the misclassified example without impacting the average margin over the train or test set. (Right) The right subplot visualizes examples with margins that change the most or least due to the edit. Figure 21: Adversarially inducing misclassifications on ImageNet examples via Coar-Edit. Each row corresponds to an ImageNet test example that is correctly classified by a ResNet-50 model. In the left subplot of each row, we show that applying Coar-Edit (by ablating the top-k components (9)) decreases the correct-class margin (5) of the correctly classified example without impacting the average margin over the train or test set. On the right, we shw that the edit does not impact visually dissimilar examples, but does increase or decrease the correct-class margin of examples containing visually similar objects, e.g., tennis balls in the second row. Figure 22: Which components does Coar-Edit target to fix model errors? We analyze the specific convolution layers from which Coar-Edit ablates components (convolution filters) to correct ImageNet examples misclassified by a ResNet-50 model. On the y-axis, we plot the 30303030 most common convolution layers corresponding to the ablated components. On the x-axis, we plot the percentage of ablated components that belong to each convolution layer. We find that Coar-Edit primarily targets convolution filters from the last few layers (closet to the output) of the ResNet-50 model in order to make fine-grained edits that do not impact overall model performance. For example, more than 25%percent2525\%25 % of the ablated components belong to layer4.0,1,2.conv3—the last convolution layer in the first three residual blocks of the last layer group of the ResNet-50 model. D.2 Forgetting a class Experiment details. In Section 5.2, we use Coar-Edit to selectively forget a class of a ResNet-50 model on ImageNet. In this experiment, we set the “target” examples to be set of 10101010 examples from the class to be forgotten and the “reference” examples to be a set of 50505050 randomly selected ImageNet examples. Using these examples, we use Coar-Edit to ablate components (9) that decrease the average correct-class margin (5) of the target examples without impacting the average margin over the reference examples. Additional experiments We show that Coar-Edit can be used to selectively forget additional ImageNet classes. Specifically, in Figure 23, we reproduce the Coar-Edit experiment in Section 5.2 on three additional ImageNet classes: “folding chair”, “military uniform”, and “revolver”. Like in Figure 4, we again observe that Coar-Edit can specifically degrade the accuracy on the target class without impacting the average accuracy over the train or test set by ablating a few components (convolution filters) in the ResNet-50 model. Figure 23: Forgetting ImageNet classes via Coar-Edit. We reproduce the Coar-Edit experiment from Section 5.2 on additional ImageNet classes (one per subplot). Specifically, in each subplot, we find that ablating 15151515 of 22,7202272022,72022 , 720 convolution filters (identified via Coar-Edit) suffices to significantly degrade the accuracy of a ResNet-50 model on a specific class (in green). This edit is targeted in that it does not impact the average accuracy over the train set (in blue) or test set (in orange). D.3 Improving subpopulation robustness. Experiment details. In Section 5.3, we use Coar-Edit to improve subpopulation robustness of models trained on two benchmark datasets: Waterbirds and CelebA. In both cases, we fine-tune an ResNet-50 model via standard “empirical risk minimization” using SGD hyperparameters taken from 2020distributionally. The resulting fine-tuned models attain 64%percent6464\%64 % and 47%percent4747\%47 % worst-subpopulation accuracy on the Waterbirds and CelebA test sets, respectively. To improve subpopulation robustness on Waterbirds, we set the “target” examples to a set of 10101010 random training examples from the “waterbirds on land” (the worst-performing subpopulation) and the “reference” examples to be 10101010 random examples from other subpopulations. Analogously, for CelebA, we set the “target” examples to the set of 20202020 random examples from the “blond male” worst-performing subpopulation and the “reference” examples to 20202020 random examples from other subpopulations. Then, we use Coar-Edit to identify components that, when ablated, increase the average correct-class margin (5) of the target examples without impacting the average margin over the reference examples. In both cases, the number of components to ablate is a hyperparameter that we select by tracking the worst-subpopulation accuracy on a validation set. D.4 Mitigating backdoor attacks. Experiment details. We now describe the experiment setup in Section 5.4, where we used Coar-Edit to mitigate the effect of a backdoor attack on a ResNet-18 model trained on a backdoored CIFAR-10 dataset. The CIFAR-10 dataset is modified by adding a small blue-squared trigger to the upper left corner of 50%percent5050\%50 % of examples in the “airplane” class. Training a model with standard SGD hyperparameter on this dataset causes the model to spuriously associate the trigger with the “airplane” class, leading to a backdoor attack. That is, while the resulting model attains 89%percent8989\%89 % test accuracy, applying the attack to examples in the test set causes the model to misclassify them as “airplanes”, resulting in 37%percent3737\%37 % accuracy on test examples with the trigger. To mitigate the effect of the backdoor attack, we first sample ten examples from the training set. Then, we set the “target” examples to these two examples with the trigger and the “reference” examples to these two examples without the trigger. Then, we use Coar-Edit to ablate components (9) that increase the correct-class margin (5) of the target examples without impacting the average margin over the reference examples. Additional analysis. Recall that our experiment in Section 5.4 shows that Coar-Edit can significantly mitigate the effect of a backdoor attack on a ResNet-18 model by ablating a few backdoor-specific components. We now qualitatively analyze the components ablated via Coar-Edit to mitigate the effect of a backdoor attack in Figure 24. Specifically, we visualize the ablated components (convolution filters in this case) using the input-times-gradient saliency map method from the Captum library [kokhlikyan2020captum]. As shown in Figure 24, these visualizations suggest that the ablated components are sensitive to the blue-squared trigger. Figure 24: Visualizing components ablated via Coar-Edit to mitigate a backdoor attack. Recall that in Section 5.4, we used Coar-Edit to mitigate the effect of a backdoor attack (a blue-squared spurious trigger) on a ResNet-18 model trained on a backdoored CIFAR-10 dataset. Here, we visualize the components ablated via Coar-Edit to reduce the model’s reliance on this spurious feature. The first row shows a set of random examples from the modified CIFAR-10 test set that contain the trigger. Each subsequent row corresponds to an ablated component—a convolution filter of the ResNet-18 model in this case. In each of these rows, we use the input-times-gradient saliency map method from the Captum library [kokhlikyan2020captum] to (qualitatively) highlight parts of the examples that are most “important” for the ablated component’s output. These maps suggest that all ablated components are sensitive to the blue-squared trigger. D.5 Improving robustness to typographic attacks. Experiment details. In Section 5.5 and Figure 7 in particular, we show that Coar-Edit can be used to improve robustness of zero-shot CLIP classifiers to typographic attacks. In this experiment, we consider a zero-shot CLIP ViT-B/16 classifier [radford2021learning] and specify a computation graph over 82,9448294482,94482 , 944 components, where each component corresponds to a weight vector in the ViT (across all layers). We evaluate the robustness of this model in a zero-shot setting on 180180180180 images and four real-world typographic attacks—“taxi”, “twitter”, “EU”, and “iPad”—taken from the dataset in [materzynska2022disentangling]. We also consider synthetic typographic attacks, where we render a blob of text on a white background and place it in the center of a given image. The zero-shot performance of the CLIP model drops from 95%percent9595\%95 % to 51%percent5151\%51 % and 54%percent5454\%54 % on the real and synthetic typographic attacks, respectively. To improve robustness, we set the “target” examples to be the 25252525 examples with a randomly picked synthetic attack and the “reference” examples to the same set of examples without any attack. Then, we use Coar-Edit to ablate components (9) that increase the average correct-class margin (5) of the target examples without impacting the average margin over the reference examples. We use a validation set comprising examples with and without the synthetic attack to select the number of components to ablate from the model. Appendix E Analyzing design choices in Coar In this section, we analyze three design choices in Coar: (a) the train-time ablation fraction α used to sample a subset of components C′⊂Csuperscript′C ⊂ C′ ⊂ C of size α⁢|C|α|C|α | C |, (b) the ablation method (Remark 2) used to intervene on the sampled components C′, and (c) the specific model output function used to compute component counterfactuals fM⁢(⋅,C′)subscript⋅superscript′f_M(·,C )fitalic_M ( ⋅ , C′ ) (1), i.e., model output fM⁢(⋅)subscript⋅f_M(·)fitalic_M ( ⋅ ) after ablating the component subset C′. E.1 Effect of ablation fraction The first step of Coar—constructing a component dataset (Equation 2)—requires choosing a ablation fraction α∈(0,1)01α∈(0,1)α ∈ ( 0 , 1 ). This hyperparameter determines the size of the random α-fraction subsets C′⊂Csuperscript′C ⊂ C′ ⊂ C used to compute component counterfactuals. A priori, however, it is not clear which ablation fraction α is best suited for learning accurate component attributions. For example, ablating too large a component subset (large α) can induce a significant drop in model performance to a point where the ablated model is no longer representative of the original model. Effect of train-time ablation fraction αtrainsubscripttrain _trainαtrain We use two metrics to quantify the effect of ablation fraction α on model outputs: • Change in model performance. We measure the effect of ablating random α-fraction subsets C′⊂Csuperscript′C ⊂ C′ ⊂ C of components on model performance, e.g., test accuracy. • Correlation between example-level model outputs. We measure the correlation between model outputs before and after ablation, e.g., logits or margins. We use these (heuristic) metrics to ensure that the ablations are not too severe to nullify model performance and that the outputs of the ablated models are still predictive of the outputs of the original model. Effect of train-time ablation fraction αtrainsubscripttrain _trainαtrain. Figure 25 evaluates how varying the train-time ablation fraction αtrainsubscripttrain _trainαtrain changes both metrics—model performance and correlation between model outputs—for all three settings considered in Section 4: CIFAR-10 ResNet-18, ImageNet ResNet-50, and ImageNet ViT-B/16. In all three settings, we find that model accuracy and margin correlation decrease as the ablation fraction α increases. For instance, ablating 15%percent1515\%15 % of components (α=0.150.15α=0.15α = 0.15) results in a significant accuracy drop for ResNets, but not for ViTs. On the other hand, ablating 1%percent11\%1 % of all components (α=0.010.01α=0.01α = 0.01) results in a small drop in accuracy and correlation, e.g., for the ResNet-18 model trained on CIFAR-10 (first row of Figure 25). Therefore, our experiments in Section 4 use α=0.100.10α=0.10α = 0.10 for the CIFAR-10 model and α=0.050.05α=0.05α = 0.05 for both ImageNet models. These findings also suggest that the choice of α depends on the model architecture and the task at hand, e.g., ViTs are more robust to zero ablations than ResNets. E.2 Effect of ablation method As discussed in Remark 2, we use a simple ablation method that sets the weights/activations of a subset of components C′⊂Csuperscript′C ⊂ C′ ⊂ C to zero. However, our method Coar is not dependent on any specific ablation method, and can be used to compute component attributions with other ablation methods as well. Alternative ablation method based on scaling. In this section, we consider an alternative ablation method that scales down the activations of a component by a factor of γ∈[0,1]01γ∈[0,1]γ ∈ [ 0 , 1 ]. Note that setting γ=00γ=0γ = 0 corresponds to the zero ablation method described in Remark 2; we use γ=0.50.5γ=0.5γ = 0.5 in our experiments. Experiment results. We find that the alternative scaling-based ablation maintains high correlation between model outputs before and after ablations, resulting in accurate component attributions. Specifically, we make three key observations. • We first observe that on a ResNet-18 model trained on CIFAR-10, the scaling-based ablation method described above maintains high correlation between model outputs before and after ablation, even at high ablation fractions α∈0.30,…,0.050.30…0.05α∈\0.30,…,0.05\α ∈ 0.30 , … , 0.05 (fourth row of Figure 25). • Then, in Figure 26, we apply Coar with the scaling-based ablation method to a CIFAR-10 ResNet-18 model. The resulting component attributions attain an average correlation of more than 0.750.750.750.75 for most ablation fractions α∈0.40,…,0.010.40…0.01α∈\0.40,…,0.01\α ∈ 0.40 , … , 0.01 . The correlation between Coar attribution estimates and ground-truth counterfactuals is high across a range of ablation fractions α from 0.010.010.010.01 to 0.450.450.450.45. • In Figure 27, we compare Coar attributions computed with the scaling ablations to attributions computed with zero-ablations. We find that (a) these attributions exhibit high cosine similarity (Figure 27a) and that (b) attributions learned with scaling-based ablations are predictive of ground-truth component counterfactuals computed using zero-ablations (Figure 27b). This indicates that both ablations—scaling down the activations of a component by a factor of γ=0.50.5γ=0.5γ = 0.5 and setting the activations of a component to zero—change model outputs in a similar manner. E.3 Effect of model output function Recall that in Section 4, we use the correct-class margin (5) as the model output function to estimate Coar attributions for classification tasks. However, our approach is not tied to a specific model output function. Depending on the task at hand, one can use an alternative model output function to estimate Coar attributions. For example, in a multi-label classification task, we can also use the logit of a fixed class of interest as the model output function to estimate Coar attributions. In Figure 12, we apply Coar to a pre-trained ImageNet ResNet50 model fine-tuned on MIMIC-CXR [johnson2019mimic]—a dataset of labeled chest radiographs—and set the model output function to be the logit of the “Cardiomegaly” class. Our results show that Coar attributions remain predictive with this model output function, and attain a correlation of 0.70.70.70.7 and 0.60.60.60.6 with the ground-truth counterfactuals on “Cardiomegaly” logits when α=αtrain=0.05subscripttrain0.05α= _train=0.05α = αtrain = 0.05 and α=0.100.10α=0.10α = 0.10 respectively. Additionally, in Appendix B, we also apply Coar to the next-token prediction task in language modeling, using average correct-class margin over all tokens in a given sequence as the model output function. Figure 25: Effect of ablation fraction α on model outputs. We evaluate the effect of ablating α-fraction subsets C′⊂Csuperscript′C ⊂ C′ ⊂ C of components (x-axis) on model accuracy (y-axis in the left column) and the correlation between model outputs before and after ablation (y-axis in the right column). In all settings considered in Section 4 (one per row), we find that model accuracy and margin correlation gradually decrease as the ablation fraction α increases. See Section E.1 for more details. Figure 26: Effect of ablation method on Coar attributions. We estimate Coar attributions for a CIFAR-10 ResNet-18 model using an alternative ablation method that scales down the activations of a subset of components C′⊂Csuperscript′C ⊂ C′ ⊂ C by a factor of γ (0.50.50.50.5 in this case) instead of setting them to zero. The resulting attribution-derived estimates (4) exhibit high correlation (y-axis) with ground-truth component counterfactuals. See Section E.2 for more details. Figure 27: Comparing Coar attributions estimated with different ablation methods. We compare Coar attributions on a CIFAR-10 ResNet18 model computed with the zero-ablation method Remark 2 to attributions computed with the alternative ablation method described in Section E.2. The left plot shows that the paired attributions (corresponding to each example) exhibit high cosine similarity. The right plot shows that the counterfactual estimates (4) computed using attributions from the alternative ablation method are predictive of ground-truth component counterfactuals computed using the zero ablation method. See Section E.2 for more details.

Instant research discovery