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 paper presents a method to improve robustness of object classifiers to style-changes in components, termed “concepts” in the paper, which can include both background elements as well as parts of the object of interest. This is done by first localizing confounding concepts using existing instance segmentation methods, performing style-transfer on the localized concept using existing methods, and then using an existing method which performs a rank-1 update to (a layer of) network parameters to map the transformed-concept-region to the same value as the original concept-region was mapped to by the layer. Experiments illustrate that this pipeline is effective at making the network robust to style-changes for concepts identified to be potentially misleading.