Editing a classifier by rewriting its prediction rules

1 Introduction . At the core of machine learning is the ability to automatically discover prediction rules from raw data . However , there is mounting evidence that not all of these rules are reliable [ Torralba and Efros , 2011 , Beery et al. , 2018 , Shetty et al. , 2019 , Agarwal et al. , 2020 , Xiao et al. , 2020 , Bissoto et al. , 2020 , Geirhos et al. , 2020 ] . In particular , some rules could be based on biases in the training data : e.g. , learning to associate cows with grass since they are typically depicted on pastures [ Beery et al. , 2018 ] . While such prediction rules may be useful in some scenarios , they will be irrelevant or misleading in others . This raises the question : How can we most effectively modify the way in which a given model makes its predictions ? The canonical approach for performing such post hoc modifications is to intervene at the data level . For example , by gathering additional data that better reflects the real world ( e.g. , images of cows on the beach ) and then using it to further train the model . Unfortunately , collecting such data can be challenging : how do we get cows to pose for us in a variety of environments ? Furthermore , data collection is ultimately a very indirect way of specifying the intended model behavior . After all , even when data has been carefully curated to reflect a given real-world task , models still end up learning unintended prediction rules from it [ Ponce et al. , 2006 , Torralba and Efros , 2011 , Tsipras et al. , 2020 , Beyer et al. , 2020 ] . Our contributions The goal of our work is to develop a toolkit that enables users to directly modify the prediction rules learned by an ( image ) classifier , as opposed to doing so implicitly via the data . Concretely : Editing prediction rules . We build on the recent work of Bau et al . [ 2020a ] to develop a method for modifying a classifier ’ s prediction rules with essentially no additional data collection ( Section 2 ) . At a high level , our method enables the user to modify the weight of a layer so that the latent ∗Equal contribution . 1Our code is available at https : //github.com/MadryLab/EditingClassifiers . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . representations of a specific concept ( e.g. , snow ) map to the representations of another ( e.g. , road ) . Crucially , this allows us to change the behavior of the classifier on all occurrences of that concept , beyond the specific examples ( and the corresponding classes ) used in the editing process . Real-world scenarios . We demonstrate our approach in two scenarios motivated by real-world applications ( Section 3 ) . First , we focus on adapting an ImageNet classifier to a new environment : recognizing vehicles on snowy roads . Second , we consider the recent “ typographic attack ” of Goh et al . [ 2021 ] on a zero-shot CLIP [ Radford et al. , 2021 ] classifier : attaching a piece of paper with “ iPod ” written on it to various household items causes them to be incorrectly classified as “ iPod. ” In both settings , we find that our approach enables us to significantly improve model performance , using only a single , synthetic example to perform the edit—cf . Figure 1 . Large-scale synthetic evaluation . To evaluate our method at scale , we develop an automated pipeline to generate a suite of varied test cases ( Section 4 ) . Our pipeline revolves around identifying specific concepts ( e.g. , “ road '' or “ pasture ” ) in an existing dataset using pre-trained instance segmentation models and then modifying them using style transfer [ Gatys et al. , 2016 ] ( e.g. , to create “ snowy road ” ) . We find that our editing methodology is able to consistently correct a significant fraction of model failures induced by these transformations . In contrast , standard fine-tuning approaches are unable to do so given the same data , often causing more errors than they are fixing . Probing model behavior with counterfactuals . Beyond model editing , our concepttransformation pipeline can also be viewed as a scalable way of generating image counterfactuals . In Section 5 , we demonstrate how such counterfactuals can be useful to gain insights into how a given model makes its predictions and pinpoint certain spurious correlations that it has picked up . 2 A toolkit for editing prediction rules . It has been widely observed that models pick up various context-specific correlations in the data— e.g. , using the presence of “ road ” or a “ wheel ” to predict “ car ” ( cf . Section 5 ) . Such unreliable prediction rules ( dependencies of predictions on specific input concepts ) could hinder models when they encounter novel environments ( e.g. , snow-covered roads ) , and confusing or adversarial test conditions ( e.g. , cars with wooden wheels ) . Thus , a model designer might want to modify these rules before deploying their model . The canonical approach to modify a classifier post hoc is to collect additional data that captures the desired deployment scenario , and use it to retrain the model . However , even setting aside the challenges of data collection , it is not obvious a priori how much of an effect such retraining ( e.g. , via fine-tuning ) will have . For instance , if we fine-tune our model on “ cars ” with wooden wheels , will it now recognize “ scooters ” or “ trucks ” with such wheels ? The goal of this work is to instead develop a more direct way to modify a model ’ s behavior : rewriting its prediction rules in a targeted manner . For instance , in our previous example , we would ideally be able to modify the classifier to correctly recognize all vehicles with wooden wheels by simply teaching it to treat any wooden wheel as it would a standard one . Our approach is able to do exactly this . However , before describing this approach ( Section 2.2 ) , we first provide a brief overview of recent work by Bau et al . [ 2020a ] which forms its basis . 2.1 Background : Rewriting generative models . Bau et al . [ 2020a ] developed an approach for rewriting a deep generative model : specifically , enabling a user to replace all occurrences of one selected object ( say , “ dome ” ) in the generated images with another ( say , “ tree ” ) , without changing the model ’ s behavior in other contexts . Their approach is motivated by the observation that , using a handful of example images , we can identify a vector in the model ’ s representation space that encodes a specific high-level concept [ Kim et al. , 2018 , Bau et al. , 2020a ] . Building on this , Bau et al . [ 2020a ] treat each layer of the model as an associative memory , which maps such a concept vector at each spatial location in its input ( which we will refer to as the key ) to another concept vector in its output ( which we will call the value ) . In the simplest case , one can think of a linear layer with weights W ∈ Rmxn transforming the key k ∈ Rn to the value v ∈ Rm . In this setting , Bau et al . [ 2020a ] formulate the rewrite operation as modifying the layer weights from W to W ′ so that v∗ = W ′k∗ , where k∗ corresponds to the old concept that we want to replace , and v∗ to the new concept . For instance , to replace “ domes ” with “ trees ” in the generated images , we would modify the layer so that the key k∗ for “ dome ” maps to the value v∗ for “ tree ” . Consequently , when this value is fed into the downstream layers of the network it would result in a tree in the final image . Crucially , this update should change the model ’ s behavior for every instance of the concept encoded in k∗—i.e. , all “ domes ” in the images should now be “ trees ” . To extend this approach to typical deep generative models , two challenges remain : ( 1 ) handling non-linear layers , and ( 2 ) ensuring that the edit doesn ’ t significantly hurt model behavior on other concepts . With these considerations in mind , Bau et al . [ 2020a ] propose making the following rankone updates to the parameters W of an arbitrary non-linear layer f : min Λ ∑ ( i , j ) ∈S ∥∥v∗ij − f ( k∗ij ; W ′ ) ∥∥ s.t . W ′ = W + Λ ( C−1d ) > . ( 1 ) Here , S denotes the set of spatial locations in representation space for a single image corresponding to the concept of interest , d is the top eigenvector of the keys k∗ij at locations ( i , j ) ∈ S and C =∑ d kdkd > captures the second-order statistics for other keys kd . Intuitively , the goal of this update is to modify the layer parameters to rewrite the desired key-value mapping in the most minimal way . We refer the reader to Bau et al . [ 2020a ] for further details . 2.2 Editing classifiers . We now shift our attention to the focus of this work : editing classifiers . To describe our approach , we use the task of enabling classifiers to detect vehicles with “ wooden wheels ” as a running example . At a high level , we would like to apply the approach described in Section 2.1 to modify a chosen ( potentially non-linear ) layer L of the network to rewrite the relevant key-value association . But , we need to first determine what these relevant keys and values are . Let us start with a single image x , say , from class “ car ” , that contains the concept “ wheel ” . Let the location of the “ wheel ” in the image be denoted by a binary mask m.2 Then , say we have access to a transformed version of x—namely , x′—where the “ car ” has a “ wooden wheel ” ( cf . Figure 2 ) . This x′ could be created by manually replacing the wheel , or by applying an automated procedure such as that in Section 4 . In the rest of our study , we refer to a single ( x , x′ ) pair as an exemplar . 2Such a mask can either be obtained manually or automatically via instance segmentation ( cf . Section 4 ) . Intuitively , we want the classifier to perceive the “ wooden wheel ” in the transformed image x′ as it does the standard wheel in the original image x . To achieve this , we must map the keys for wooden wheels to the value corresponding to their standard counterparts . In other words , the relevant keys ( k∗ ) correspond to the network ’ s representation of the concept in the transformed image ( x′ ) directly before layer L. Similarly , the relevant values ( v∗ ) that we want to map these keys to correspond to the network ’ s representation of the concept in the original images ( x ) directly after layer L. ( The pertinent spatial regions in the representation space are simply determined by downsampling the mask to the appropriate dimensions . ) Finally , the model edit is performed by feeding the resulting key-value pairs into the optimization problem ( 1 ) to determine the updated layer weights W ′—cf . Figure 2 for an illustration of the overall process . Note that this approach can be easily extended to use multiple exemplars ( xk , x′k ) by simply expanding S to include the union of relevant spatial locations ( corresponding the concept of interest ) across these exemplars .

Summary: The authors propose a model editing method that proceeds in two steps. First, they rely on off-the-shelf segmentation models and style transfer models to learn a data transformation $x \to x'$ that preserves semantic meaning (e.g., in a vehicle recognition task, they might convert a snowy road into a non-snowy road, which preserves the label of the vehicle). Second, they build on recent work by Bau et al. to use these transformed pairs $(x,  x')$ to learn an edit to the model $f$ such that, roughly speaking, $f(x)$ and $f(x')$ agree in their representations at some particular layer and above. They show that their procedure generalizes (astonishingly) well from edits learned on just a single exemplar.