# CytoImageNet: A large-scale pretraining dataset for bioimage transfer
learning

Stanley Bryan Z. Hua
Computer Science
University of Toronto
<EMAIL_ADDRESS>
Alex X. Lu
Computer Science
University of Toronto
(Current Affiliation: Microsoft Research)
<EMAIL_ADDRESS>
Alan M. Moses
Cell and Systems Biology
Computer Science
Center of Genome Evolution and Function
University of Toronto
<EMAIL_ADDRESS>

###### Abstract

Motivation: In recent years, image-based biological assays have steadily
become high-throughput, sparking a need for fast automated methods to extract
biologically-meaningful information from hundreds of thousands of images.
Taking inspiration from the success of ImageNet, we curate CytoImageNet, a
large-scale dataset of openly-sourced and weakly-labeled microscopy images
(890K images, 894 classes). Pretraining on CytoImageNet yields features that
are competitive to ImageNet features on downstream microscopy classification
tasks. We show evidence that CytoImageNet features capture information not
available in ImageNet-trained features. The dataset is made available at
https://www.kaggle.com/stanleyhua/cytoimagenet.

##  1 Introduction

Automated microscopy has given biologists a tool for studying various cell
types under varying experimental conditions with minimal assistance. In recent
years, automated image analysis has proven useful for drug discovery and
functional genomics research; identifying the function of novel genes and
cellular response to various drugs [1]. With the capacity to perform large-
scale screens rapidly, tens of thousands of images are being collected daily,
demanding computational methods to extract biology from these images [2].

Methods for automated image analysis remains an active area of research. One
strategy for automating the analysis of microscopy images is to use low-level
features from classic computer vision strategies; extracting thousands of
general features related to area, shape, texture and intensity of single-cells
[3, 4]. However, these features usually require extensive engineering,
parameter tuning, and effective single-cell segmentation, all of which must be
customized to the images of interest. Furthermore, these features may not
capture relevant biological signal and instead capture experimental noise [5].
To reduce labor and improve the sensitivity of extracted features, some
studies rely on deep learning to automatically learn features. These
strategies range from supervised training with labels, autoencoders to more
recently, self-supervised learning [5, 6, 7]. However, these models require
users to train on their own data, which is computationally expensive, time-
consuming, and requires specialized hardware like GPUs.

As an alternative to training such models, some studies instead rely on
transfer learning. Most commonly, studies re-use features from neural networks
trained to classify images from ImageNet, a large-scale natural image
classification dataset [8, 9, 10]. By training models to perform well on large
datasets, they can learn useful features that may be transferred to other
datasets.

ImageNet is a large-scale dataset of more than 14 million diverse and
annotated natural images [11]. A defining factor in the success of ImageNet is
their diversity of image labels and image sources. Models that perform the
best on ImageNet-1K become the standard for transfer learning [12]. These
models learn a diverse set of features that are able to capture complex and
generalizable patterns inherent in their dataset. Surprisingly, ImageNet
features have even shown good transferability on microscopy classification
tasks [13, 14].

However, a recent study suggests that “ImageNet features are less general than
previously suggested,” and that pretraining on datasets with a closer domain
match to the downstream task benefits transfer learning [12]. In fact, Caicedo
et al. showed that pretraining convolutional networks (CNN) on weakly labeled
microscopy images yields features that outperform ImageNet features on
downstream classification tasks [5]. We hypothesize that features from
pretraining on a diverse large-scale microscopy dataset will yield more
domain-relevant features than ImageNet.

We present CytoImageNet, a large-scale image dataset of weakly labeled
microscopy images (890K images, 894 classes). Incorporating image data and
labels from 40 distinct microscopy datasets, CytoImageNet mimics the diverse
and complex nature of ImageNet. We provide a framework for curating similar
large-scale pretraining datasets using openly available microscopy datasets.

We believe that pretraining on CytoImageNet may yield biologically-meaningful
image representations useful in any downstream biological task. Since
concatenation of CytoImageNet-trained features and ImageNet-trained features
performs significantly better than either one alone, we believe that
CytoImageNet pretraining provides features that are useful, and different from
ImageNet features.

##  2 Methods

CytoImageNet is a dataset of 890,737 microscopy images sourced from 40 openly
available datasets. The selected datasets originated from 5 databases: (1)
Recursion, (2) Image Data Resource (IDR) [15], (3) Broad Bioimage Benchmark
Collection (BBBC) [16], (4) Kaggle and (5) Cell Image Library (CIL). Electron
microscopy and histopathology images were excluded. Code for the methods below
is available at https://github.com/stan-hua/CytoImageNet. The dataset is
available at https://www.kaggle.com/stanleyhua/cytoimagenet.

Figure 1: Examples of CytoImageNet classes. Two representative images from two representative classes are shown for each metadata category (names are in vertical text) Table 1: Number of labels contributed by each image database Database | Number of Labels Contributed
---|---
Recursion | 651
Image Data Resource | 450
Broad Bioimage Benchmark Collection | 202
Kaggle | 27
Cell Image Library | 1

###  2.1 Weak label assignment & stratified downsampling

As a training target for CytoImageNet, we sought to associate each microscopy
image with a class label. As expert-assigned annotations are not readily
available or standardized across publicly available image datasets, we instead
chose to assign images to classes based upon metadata associated with each of
the image datasets we used. We considered 7 kinds of metadata: (organism,
cell_type, cell_visible, phenotype, compound, gene, and sirna). Each column
was searched for unique values and their counts to find potential weak labels.
To create 894 labels, unique values with at least 287 images served as
potential labels. Beginning with the least counts, each potential label was
iterated through, searching for images that can be assigned to the label.
Images that were already assigned labels are excluded, and their unique image
index is stored in a hash table pointing to their label. To improve the
diversity of images selected for a label, stratified sampling was done by
sampling from images that were grouped by values in the remaining potential
label columns.

###  2.2 Image standardization & upsampling

To preprocess images, we converted all images to PNG. RGB images are converted
to grayscale by averaging over channels. We normalized each channel between
[0, 255] using the 0.1 and 99.9th percentile pixel intensity as the lower and
upper ranges respectively [6]. Finally, for fluorescent images, we merged the
channels.

We took 1-4 crops from each image of varying spatial scales where possible.
Unless the image was smaller than 70x70 pixels, we cropped images into four
quadrants. To introduce variations in spatial scale, we further cropped these
quadrants by sampling crops of either full, 0.5, 0.25, or 0.125 of the
quadrant (excluding any crop scales that would produce images smaller than
70x70 pixels). Finally, we filtered any crops that were all black, all white,
or where the 75th percentile pixel intensity is equal to 0. Crop pixel
intensities were re-normalized as we previously did for full images.

###  2.3 Training

CytoImageNet is split into a training and validation set with 10% used for
validation. This yields roughly 900 training samples for each label. Images
are fed in batches of 64 with random 0 to 360 degrees rotations. We train
convolutional networks (EfficientNetB0) to classify one of the 894 labels, by
minimizing the categorical cross-entropy loss of predictions to ground truth
labels. Randomly initialized models were trained for 24 epochs (2 weeks) on an
NVIDIA Tesla K40C. The loss was optimized via the Adam optimizer with learning
rate of 0.001.

###  2.4 Evaluation

We compare features extracted from a CytoImageNet-pretrained model, an
ImageNet-pretrained model, and a randomly-initialized model on three distinct
downstream tasks: 1) BBBC021 mechanism-of-action classification, 2) Cells-Out-
Of-Sample (COOS-7) protein localization classification, and 3) a CYCLoPs
dataset for protein localization classification [17, 18, 19]. Additionally, we
would like to know if CytoImageNet and ImageNet pretraining lead to different
learnt features. If so, fusing the extracted features from both may increase
performance on downstream tasks. To test this, we concatenate CytoImageNet and
ImageNet extracted features for each image in the transfer tasks.

BBBC021 is a dataset of fully imaged human cells. Cells are treated with one
of 113 small molecules at 8 concentrations, and fluorescent images are
captured staining for nucleus, actin and microtubules. The phenotypic
profiling problem is presented, where the goal is to extract features
containing meaningful information about the cellular phenotype exhibited. Each
of 103 unique compound-concentration treatment is labeled with a mechanism-of-
action (MOA). The MOA is predicted for each unique treatment (averaging
features over all treatment examples) by matching the MOA of the closest point
excluding points of the same compound. We use a 1-nearest neighbors using
cosine distance as the metric and report the not-same-compound (NSC) accuracy.

On the other hand, the WT2 dataset from the CYCLoPs database consists of
27,058 single-cell images of yeast cells. The task is to classify the
subcellular localization of a fluoresced protein, given two channels staining
for the protein of interest and the cytosol. An 11-nearest neighbors is used
to perform leave-one-out classification, and the accuracy is reported.

Lastly, COOS-7 contains 132,209 single-cell images of mouse cells. Images are
spread over 1 training set and 4 testing sets, where each single-cell image
contains a protein and nucleus fluorescent channels. COOS-7 was curated, such
that each succeeding testing set has a greater degree of covariate shift from
the training set. Similar to the CYCLoPs dataset, the COOS-7 dataset presents
a protein subcellular localization classification task. We use a 11-nearest
neighbors fit on the training set to predict on the testing sets and report
the accuracy.

We chose these datasets (1) to assess the generalizability of CytoImageNet-
pretrained models on distinct biological image datasets and problems, and (2)
to compare against ImageNet features that have already been shown to perform
relatively well on these datasets [13, 7, 19].

Features are extracted from the penultimate layer of EfficientNetB0 models.
ImageNet-pretrained EfficietNetB0 weights are loaded in using Keras, but the
weights themselves belong to the source paper [20]. Predictions are generated
using k-Nearest Neighbors (kNN), and accuracy is reported with a 95%
confidence interval. BBBC021 kNN is implemented with scikit-learn, while kNN
is implemented using the faiss library for COOS-7 and CYCLoPs datasets.

We compare the performance of features extracted following two optional
preprocessing methods: normalization and channel merging. Normalization refers
to channel normalization with the lower bound and upper bound of the 0.1 and
99.9th percentile pixel intensity respectively [6]. Channel merging refers to
the merging of fluorescent channels into a single grayscale image before
extracting features. If not, features are extracted for each fluorescent
channel then concatenated.

Table 2: kNN classification accuracy (with 95% confidence intervals) on downstream tasks, comparing EfficientNetB0 penultimate layer features. ’Random’ refers to the model with randomly-initialized weights. ’Fusion’ refers to the concatenation of ImageNet features and CytoImageNet features. Concatenation of randomly-initialized features with ImageNet features gave no increase in performance (data not shown). Downstream Task | Random | ImageNet | CytoImageNet | Fusion
---|---|---|---|---
BBC021 | 27.18 $\pm$ 8.59% |  83.5 $\pm$ 7.17% |  83.5 $\pm$ 7.17% |  86.41 $\pm$ 6.62%
CYCLoPs | 53.06 $\pm$ 0.59% |  68.47 $\pm$ 0.55% | 65.19 $\pm$ 0.57% |  77.97 $\pm$ 0.49%
COOS7 Test Set 1 | 65.68 $\pm$ 0.91% |  88.87 $\pm$ 0.61% | 88.58 $\pm$ 0.61% |  94.80 $\pm$ 0.43%
COOS7 Test Set 2 | 67.81 $\pm$ 0.7% | 88.93 $\pm$ 0.47% |  89.37 $\pm$ 0.46% |  95.07 $\pm$ 0.33%
COOS7 Test Set 3 | 48.77 $\pm$ 0.54% |  75.91 $\pm$ 0.46% | 65.97 $\pm$ 0.51% |  78.97 $\pm$ 0.44%
COOS7 Test Set 4 | 51.88 $\pm$ 0.56% |  82.19 $\pm$ 0.43% | 78.52 $\pm$ 0.46% |  87.47 $\pm$ 0.37%

##  3 Discussion

CytoImageNet pretraining shows comparable performance to ImageNet pretraining
on all downstream tasks. Normalizing pixel intensities for each channel and
extracting features for each channel led to the best performance for all
features, implying that automated methods for feature extraction still require
some form of image preprocessing.

Given that CytoImageNet has a closer domain match to the downstream microscopy
tasks, it is surprising that CytoImageNet features alone do not outperform
ImageNet features. A recent study on ImageNet pretraining reported a strong
correlation between ImageNet validation accuracy and transfer accuracy [12].
Our trained EfficientNetB0 model only achieved 13.42% accuracy on the training
set and more importantly, 11.32% on the CytoImageNet validation set, where 402
of the 894 labels had no images correctly predicted. Of which, 375 of the 402
were compound labels. Meanwhile, EfficientNetB0 was shown to attain a top-1
accuracy of 77.1% on ImageNet-1K [20]. Given our low validation accuracy, we
speculate that hyperparameter optimization will lead to further increases in
benchmark classification accuracy. We believe this is an important area for
future research.

Encouragingly, fusing extracted CytoImageNet and ImageNet features (via
concatenation) performs the best on downstream tasks by a significant margin.
This implies that CytoImageNet and ImageNet pretraining result in the learning
of different albeit meaningful image representations. We believe this is due
to CytoImageNet-trained models learning some feature of microscopy images that
are not available in natural images. These results suggest that future
applications of bioimage transfer learning may benefit the most from the
fusion of CytoImageNet and ImageNet features.

In summary, we present CytoImageNet; a pretraining dataset of 890,737
microscopy images and 894 classes. We contribute a simple framework for
combining openly available microscopy datasets, in the hopes that this sparks
interest in the vast amount of rich microscopy data available online.
Furthermore, we hope that this may inspire others to experiment on their own;
both to create new large-scale datasets and to develop image classification
architectures and training procedures that generalize to other image domains.

## Acknowledgements

Stanley Z. Hua was funded by the University of Toronto Cell & Systems Biology
undergraduate research award. Alan M. Moses holds a Tier II Canada Research
Chair. The authors would like to acknowledge Recursion, Image Data Resource,
Broad Bioimage Benchmark Collection and Kaggle for allowing open access to
their microscopy datasets, which were used to construct the CytoImageNet
dataset. We would like to thank the authors responsible for uploading their
data on these databases. We would also like to thank Juan Caicedo and Anne
Carpenter for their valuable input regarding the Broad Institute datasets.
Finally, the authors would like to acknowledge the Canadian Foundation for
Innovation that sponsored the NVIDIA GPU with which the models were trained
on.

## References

[1] Caicedo JC, Singh S, Carpenter AE. Applications in image-based profiling
of perturbations. Current Opinion in Biotechnology. 2016;39:134–142.
doi:10.1016/j.copbio.2016.04.003

[2] Wollman R, Stuurman N. High throughput microscopy: From raw images to
discoveries. Journal of Cell Science. 2007;120(21):3715–3722.
doi:10.1242/jcs.013623

[3] Danuser G. Computer Vision in cell biology. Cell. 2011;147(5):973–978.
doi:10.1016/j.cell.2011.11.001

[4] Carpenter AE, Jones TR, Lamprecht MR, Clarke C, Kang I, Friman O, Guertin
DA, Chang J, Lindquist RA, Moffat J, et al. CellProfiler: image analysis
software for identifying and quantifying cell phenotypes. Genome Biology.
2006;7(10). doi:10.1186/gb-2006-7-10-r100

[5] Caicedo JC, McQuin C, Goodman A, Singh S, Carpenter AE. Weakly supervised
learning of single-cell feature embeddings. 2018 IEEE/CVF Conference on
Computer Vision and Pattern Recognition. 2018. doi:10.1109/cvpr.2018.00970

[6] Kraus OZ, Grys BT, Ba J, Chong Y, Frey BJ, Boone C, Andrews BJ. Automated
Analysis of high-content microscopy data with Deep Learning. Molecular Systems
Biology. 2017;13(4):924. doi:10.15252/msb.20177551

[7] Lu AX, Kraus OZ, Cooper S, Moses AM. Learning unsupervised feature
representations for single cell microscopy images with paired cell inpainting.
PLOS Computational Biology. 2019;15(9). doi:10.1371/journal.pcbi.1007348

[8] Huh M, Agrawal P, Efros AA. What makes ImageNet good for transfer
learning? arXiv. 2016;1608.08614.

[9] Razavian AS, Azizpour H, Sullivan J, and Carlsson S. CNN features off-the-
shelf: an astounding baseline for recognition. In IEEE Conference on Computer
Vision and Pattern Recognition Workshops (CVPRW),pages 512–519. IEEE, 2014.

[10] Donahue J, Jia Y, Vinyals O, Hoffman J, Zhang N, Tzeng E, and Darrell T.
Decaf: A deep convolutional activation feature for generic visual recognition.
In International Conference on Machine Learning, pages 647–655, 2014.

[11] Deng J, Dong W, Socher R, Li L-J, Kai Li, Li Fei-Fei. ImageNet: A large-
scale hierarchical image database. 2009 IEEE Conference on Computer Vision and
Pattern Recognition. 2009. doi:10.1109/cvpr.2009.5206848

[12] Kornblith S, Shlens J, Le QV. Do Better Imagenet Models Transfer Better?
2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).
2019. doi:10.1109/cvpr.2019.00277

[13] Pawlowski N, Caicedo JC, Singh S, Carpenter AE, Storkey A. Automating
morphological profiling with generic deep convolutional networks. bioRxiv.
2016. doi:10.1101/085118

[14] Kensert A, Harrison PJ, Spjuth O. Transfer learning with deep
convolutional neural networks for classifying cellular morphological changes.
SLAS DISCOVERY: Advancing the Science of Drug Discovery. 2019;24(4):466–475.
doi:10.1177/2472555218818756

[15] Williams E, Moore J, Li SW, Rustici G, Tarkowska A, Chessel A, Leo S,
Antal B, Ferguson RK, Sarkans U, et al. Image Data Resource: A Bioimage Data
Integration and publication platform. Nature Methods. 2017;14(8):775–781.
doi:10.1038/nmeth.4326

[16] Ljosa V, Sokolnicki KL, Carpenter AE. Annotated high-throughput
microscopy image sets for validation. Nature Methods. 2012;9(7):637–637.
doi:10.1038/nmeth.2083

[17] Caie PD, Walls RE, Ingleston-Orme A, Daya S, Houslay T, Eagle R, Roberts
ME, Carragher NO. High-content phenotypic profiling of drug response
signatures across distinct cancer cells. Molecular Cancer Therapeutics.
2010;9(6):1913–1926. doi:10.1158/1535-7163.mct-09-1148

[18] Koh JL, Chong YT, Friesen H, Moses A, Boone C, Andrews BJ, Moffat J.
Cyclops: A comprehensive database constructed from automated analysis of
protein abundance and subcellular localization patterns in saccharomyces
cerevisiae. G3 Genes|Genomes|Genetics. 2015;5(6):1223–1232.
doi:10.1534/g3.115.017830

[19] Lu AX, Lu AX, Schormann W, Andrews DW, Moses AM. The Cells Out of Sample
(COOS) dataset and benchmarks for measuring out-of-sample generalization of
image classifiers. ArXiv. 2019;abs/1906.07282.

[20] Tan M, Le Q. EfficientNet: Rethinking Model Scaling for Convolutional
Neural Networks. In: 36th International Conference on Machine Learning (ICML
2019) Long Beach, California, USA, 9-15 June 2019. Red Hook, NY: Curran
Associates, Inc; 2019.

[21] Zhang H, Cao X, Tang M, Zhong G, Si Y, Li H, Zhu F, Liao Q, Li L, Zhao J,
et al. A subcellular map of the human kinome. eLife. 2021;10.
doi:10.7554/elife.64943

## Supplementary Material

Supplementary Table 1: List of Datasets Used with External Links

Database | Shorthand | Dataset Name | Link
---|---|---|---
BBBC | bbbc012 | C. elegans infection marker | https://bbbc.broadinstitute.org/BBBC012
BBBC | bbbc010 | C. elegans live/dead assay | https://bbbc.broadinstitute.org/BBBC010
BBBC | bbbc011 | C. elegans metabolism assay | https://bbbc.broadinstitute.org/BBBC011
BBBC | bbbc048 | Cell Cycle Jurkat Cells | https://bbbc.broadinstitute.org/BBBC048
BBBC | bbbc017 | Human HT29 colon-cancer cells shRNAi screen | https://bbbc.broadinstitute.org/BBBC017
BBBC | bbbc026 | Human Hepatocyte and Murine Fibroblast cells - Co-culture experiment | https://bbbc.broadinstitute.org/BBBC026
BBBC | bbbc006 | Human U2OS cells (out of focus) | https://bbbc.broadinstitute.org/BBBC006
BBBC | bbbc022 | Human U2OS cells - compound cell-painting experiment | https://bbbc.broadinstitute.org/BBBC022
BBBC | bbbc013 | Human U2OS cells cytoplasm-nucleus translocation | https://bbbc.broadinstitute.org/BBBC013
BBBC | bbbc014 | Human U2OS cells cytoplasm-nucleus translocation (2) | https://bbbc.broadinstitute.org/BBBC014
BBBC | bbbc015 | Human U2OS cells transfluor | https://bbbc.broadinstitute.org/BBBC015
BBBC | bbbc025 | Human U2OS cells - RNAi Cell Painting experiment | https://bbbc.broadinstitute.org/BBBC025
BBBC | bbbc045 | Human White Blood Cells | https://bbbc.broadinstitute.org/BBBC045
BBBC | bbbc051 | Human kidney cortex cells | https://bbbc.broadinstitute.org/BBBC051
BBBC | bbbc038 | Kaggle 2018 Data Science Bowl | https://bbbc.broadinstitute.org/BBBC038
BBBC | bbbc020 | Murine bone-marrow derived macrophages | https://bbbc.broadinstitute.org/BBBC020
BBBC | bbbc041 | P. vivax (malaria) infected human blood smears | https://bbbc.broadinstitute.org/BBBC041
BBBC | bbbc005 | Synthetic cells | https://bbbc.broadinstitute.org/BBBC005
CIL | cil_kinome | Kinome Atlas | http://cellimagelibrary.org/pages/kinome_atlas
IDR | idr0081 | Adenovirus | https://idr.openmicroscopy.org/webclient/?show=screen-2406
IDR | idr0017 | Chemical-Genetic Interaction Map | https://idr.openmicroscopy.org/webclient/?show=screen-1151
IDR | idr0016 | Compound Profiling | https://idr.openmicroscopy.org/webclient/?show=screen-1251
IDR | idr0009 | Early Secretory Pathway | https://idr.openmicroscopy.org/webclient/?show=screen-803
IDR | idr0041 | Mitotic Atlas | https://idr.openmicroscopy.org/webclient/?show=project-404
IDR | idr0021 | Pericentriolar Material | https://idr.openmicroscopy.org/webclient/?show=project-51
IDR | idr0080 | Perturbation | https://idr.openmicroscopy.org/webclient/?show=screen-2701
IDR | idr0088 | Phenomic Profiling | https://idr.openmicroscopy.org/webclient/?show=screen-2651
IDR | idr0003 | Plasticity | https://idr.openmicroscopy.org/webclient/?show=screen-51
IDR | idr0072 | Subcellular Localization | https://idr.openmicroscopy.org/webclient/?show=screen-2952
IDR | idr0037 | Variation in Human iPSC lines | https://idr.openmicroscopy.org/webclient/?show=screen-2051
IDR | idr0067 | Yeast Meiosis | https://idr.openmicroscopy.org/webclient/?show=project-904
IDR | idr0093 | siRNA screen for cell size and RNA production | https://idr.openmicroscopy.org/webclient/?show=screen-2751
Kaggle | kag_cell_cycle | Cell Cycle Experiments | https://www.kaggle.com/paultimothymooney/cell-cycle-experiments
Kaggle | kag_hpa_single | Human Protein Atlas - Single Cell Classification | https://www.kaggle.com/c/hpa-single-cell-image-classification/data
Kaggle | kag_hpa | Human Protein Atlas Image Classification | https://www.kaggle.com/c/human-protein-atlas-image-classification/data
Kaggle | kag_leukemia | Leukemia Classification | https://www.kaggle.com/andrewmvd/leukemia-classification
Recursion | rec_rxrx1 | RxRx1 | https://www.rxrx.ai/rxrx1#the-data
Recursion | rec_rxrx19a | RxRx19a | https://www.rxrx.ai/rxrx19a
Recursion | rec_rxrx19b | RxRx19b | https://www.rxrx.ai/rxrx19b
Recursion | rec_rxrx2 | RxRx2 | https://www.rxrx.ai/rxrx2


# Circuit Breaking: Removing Model Behaviors with Targeted Ablation

Maximilian Li  Xander Davies  Max Nadeau

###### Abstract

Language models often exhibit behaviors that improve performance on a pre-
training objective but harm performance on downstream tasks. We propose a
novel approach to removing undesirable behaviors by ablating a small number of
causal pathways between model components, with the intention of disabling the
computational circuit responsible for the bad behavior. Given a small dataset
of inputs where the model behaves poorly, we learn to ablate a small number of
important causal pathways. In the setting of reducing GPT-2 toxic language
generation, we find ablating just 12 of the 11.6K causal edges mitigates toxic
generation with minimal degradation of performance on other inputs.

Figure 1: In targeted ablation, we (1) rewrite our model as a computation
graph of a desired granularity, (2) learn a binary mask over edges while
regularizing to penalize ablations, and (3) ablate edges at inference time to
avoid the target bad behavior.

##  1 Introduction

Language models (LMs) often exhibit undesirable behaviors useful during pre-
training that prove hard to remove during fine-tuning. This has resulted in
capable LMs which competently hallucinate, lie, manipulate, and exhibit
undesirable biases (OpenAI, 2023; Brown et al., 2020).

In this work, we propose a new method for removing undesirable behaviors:
targeted edge ablation. In targeted edge ablation, we target a bad behavior by
removing a small number of causal pathways through the model at inference time
(Figure 1). Targeted edge ablation follows recent work in using causal
mediation to discover computational circuits responsible for particular model
behaviors (Wang et al., 2022; Goldowsky-Dill et al., 2023; Geiger et al.,
2023a). Rather than discovering circuits, targeted edge ablation discovers
causal cuts through circuits, disabling circuits responsible for bad
behaviors.

#### Main Contributions.

We formulate the problem of behavior removal and propose targeted edge
ablation as a possible solution (Section 3). We then present preliminary
results in performing targeted edge ablation to harm performance in toxic
language generation (Section 4).

##  2 Background

#### Circuit analysis.

We can write any model as a connected directed acyclic graph (DAG) with source
nodes representing the model’s (typically vector-valued) input, sink nodes
representing the model’s output, and intermediate nodes representing units of
computation (e.g. Figure 1, left; see Appendix B). Circuit analysis attempts
to mechanistically understand model computation by identifying a subgraph of
this DAG that is responsible for a given behavior, and assigning semantic
meaning to (groups of) nodes (Wang et al., 2022; Räukur et al., 2022; Chan et
al., 2022). Circuits have also been discussed in the context of treating nodes
as “features,” usually defined as directions in the latent space (Olah, 2022;
Cammarata et al., 2020).

#### Ablating edges in a computational graph.

Since edges in the model’s computational graph represent dependencies between
nodes, we can simulate what the model would have computed without a certain
node-to-node dependency by performing ablation on an edge in the graph. While
previous work has largely focused on ablation of nodes (Ghorbani & Zou, 2020),
an advantage of our strategy of ablating edges rather than nodes is the
mitigation of polysemantic behavior of model components (Olah et al., 2020),
since we investigate the causal importance of each causal path into and out of
the component. In our experiments, we use zero ablation, in which we compute
the destination node as if the source node’s value were zero, and mean
ablation (Wang et al., 2022), in which we compute the destination node as if
the source node’s value were set to its mean value over the training set. See
Appendix C for more.

##  3 Targeted Ablation for Behavior Removal

Let $\mathcal{L}(M,\mathcal{D})$ indicate the loss of model $M$ on a
distribution $\mathcal{D}$ over input-label pairs. We specify a behavior as
some distribution $\mathcal{D}$ on which the model achieves low loss
$\mathcal{L}(M,\mathcal{D})<K$ for some appropriate hyperparameter $K$. We can
define the disjointness $\delta(\mathcal{D},\mathcal{D}^{\prime})$ for
behaviors $\mathcal{D}$ and $\mathcal{D}^{\prime}$ to be the total variation
distance between $\mathcal{D}$ and $\mathcal{D}^{\prime}$. In particular, the
total variation distance is 1 if $\mathcal{D}$ assigns probability 0 to all
regions that $D^{\prime}$ assigns positive probability and vice versa.

######  Definition 3.1 (Behavior Removal).

Given a model $\mathcal{M}$ and unlimited access to training samples, produce
a model $\mathcal{M}^{*}$ which achieves high loss
$\mathcal{L}(\mathcal{M}^{*},\mathcal{D})>K$, without harming distinct
behaviors. In particular, for all behaviors $\mathcal{D}^{\prime}$ completely
disjoint from $\mathcal{D}$, i.e.
$\delta(\mathcal{D},\mathcal{D}^{\prime})=1$, we wish to preserve
$\mathcal{L}(\mathcal{M}^{*},\mathcal{D}^{\prime})\leq\mathcal{L}(\mathcal{M},\mathcal{D}^{\prime})$.

Thus, behavior removal has two goals: efficacy – the edited model should
achieve high loss on $\mathcal{D}$; and specificity – the edited model should
achieve low loss on all disjoint behaviors $\mathcal{D}^{\prime}$ for which
the original model achieves low loss.

Let $D_{\text{train}}$ be our train set, and $D_{\text{behavior}}$ be samples
from $\mathcal{D}$. One reason the model might exhibit a behavior is if
$\mathcal{D}$ overlaps with the training distribution, which would incentivize
the model to produce low loss on $\mathcal{D}$. Thus, it is reasonable to
assume $D_{\text{train}}$ and $\mathcal{D}$ may not be completely disjoint.

###  3.1 Baseline: Finetuning

We form an approximate objective function by encouraging preserving
performance on the training set, while increasing loss on the bad behavior
set:

$\mathcal{L}(\mathcal{M},D_{\text{train}})-\alpha\cdot\mathcal{L}(\mathcal{M},D_{\text{behavior}})$
(1)

where $\alpha$ is a hyperparameter. We can now finetune using Equation 1.
Since $D_{\text{behavior}}$ is often small, we use early stopping to avoid
overfitting.

###  3.2 Baseline: Task Arithmetic

In task arithmetic (Ilharco et al., 2023), we finetune $\mathcal{M}$ on
$\mathcal{L}(\mathcal{M},D_{\text{behavior}})$ towards the bad behaviors, and
find the “task vector”, or difference in weights between the finetuned model
and $\mathcal{M}$. We then form $\mathcal{M}^{*}$ by adding the negated task
vector to $\mathcal{M}$.

###  3.3 Targeted Edge Ablation

Following Figure 1, we describe targeted edge ablation as three steps.

1\. Rewrite the model. We first choose at what level of granularity to
represent the model’s computation. Since we learn a mask over edges in the
resulting graph, increasing the granularity results in a more expressive
ablation process. We call the specified graph $G$, and call its set of edges
$E_{G}$.

2\. Learn an ablation mask. Let $G_{-E}$ be our graph $G$ with the edges in
$E$ ablated. Then we wish to select $E\subset E_{G}$ that minimizes

$\displaystyle\mathcal{L}(G_{-E},D_{\text{train}})-\alpha\cdot\mathcal{L}(G_{-E},D_{\text{behavior}})+\lambda\cdot
R(E)$ (2)

for hyperparameters $\alpha,\lambda$ and some regularization function
$R$.111The regularization term penalizes large sizes of $E$ to apply pressure
to find a minimal subset of edges that disables the behavior. To compute an
optimal edge subset $E$, we optimize an edge mask $W_{\text{mask}}$ on a
continuous relaxation of Equation 2. Every edge $e=(A,B)$ is given a learnable
weight $w_{e}\in[0,1]$, where $w_{e}=0$ corresponds to ablating $e$, $w_{e}=1$
corresponds to preserving $e$, and $0<w_{e}<1$ corresponds to node $B$
observing the following convex combination of the preserved value ($v_{A}$)
and the ablated value ($\mu_{A}$) for node $A$:

$w_{e}\cdot v_{A}+(1-w_{e})\cdot\mu_{A}$ (3)

When $w_{e}=0$, node $B$’s observation of node $A$ is replaced by its ablated
value, and when $w_{e}=1$, node $B$ fully observes the value of node $A$. We
initialize the mask parameters $W_{\text{mask}}$ to a vector of 1s (indicating
fully faithful model computation) and train $W_{\text{mask}}$ on the loss
function

$\displaystyle\mathcal{L}(W_{\text{mask}};\alpha,\lambda,R)=\ $
$\displaystyle\mathcal{L}(W_{\text{mask}},D_{\text{train}})$
$\displaystyle-\alpha\cdot\mathcal{L}(W_{\text{mask}},D_{\text{bad
behavior}})$ $\displaystyle+\lambda(t)\cdot R(W_{\text{mask}})$ (4)

We train with a regularization weight $\lambda(t)$ that increases over time,
since we find that this training dynamic encourages the edge mask to find a
set of ablations that removes the bad behavior and then revise it to minimize
the number of ablations. When training is finished, we then round all the mask
weights to either 0 or 1 by selecting the set of ablated edges to be
$\hat{E}^{*}=\\{e\ |\ w_{e}\leq\tau\\}$ for some threshold $\tau\in(0,1)$.

3\. Ablate during inference. We form $\mathcal{M}^{*}$ by ablating the edges
learned in step (2) at inference time.

###  3.4 Conceptual Advantages over Fine-Tuning

#### Limited Expressivity.

LMs and other large models may have millions or billions of parameters and
thus may be vastly overparameterized for the task of performing poorly on the
bad-behavior examples, especially if generating bad-behavior examples is
expensive and the set of examples is small.222For example, collecting
jailbreaks to remove jailbreaking behavior is challenging and expensive.

A particular advantage of limiting the expressivity of our solution class is
avoiding the negative effects of training on a mis-specified objective
function like Equation 1, which encourages low loss on samples in
$D_{\text{train}}$ which exhibit the behavior but are not included in
$D_{\text{behavior}}$. Allowing the model to overfit to this loss function may
result in memorization of the points in $D_{\text{behavior}}$ to maintain low
loss on all of $D_{\text{train}}$, including those points which have high
likelihood in $\mathcal{D}$. On the other hand, edge ablation limits the
expressivity of the solution space and relies on the model’s previously
learned specialization of causal pathways.

#### Preserving Structure.

Since edge ablation edits the model at a high level, it preserves most of the
model’s mechanistic calculus. Even subtle fine-tuning has the potential to
entirely reorganize the model’s reasoning process, disrupting any mechanistic
interpretability work that has already been performed. Targeted edge ablation
is unlikely to induce the model to change its reasoning structure or increase
its knowledge because it strictly decreases the amount of information
available to the model’s computation.

##  4 Removing Toxicity in GPT-2

Figure 2: Ablating GPT-2 Small to remove toxicity. Left: Grey nodes are
attention heads, and purple nodes are MLPs. Computation proceeds upwards, with
horizontal alignment corresponding to layers. The computational graph has
11,611 edges; red edges are the 12 ablations learned to remove toxicity.
Right: Examples of improved non-toxic generation.

We apply our model editing methodology to preventing the generation of toxic
(e.g. offensive, swear-filled) sequences in a pre-trained GPT-2 Small (Radford
et al., 2019). Our goal is to edit GPT-2 so that it achieves high loss on
toxic sequences, so our $\mathcal{D}$ is a distribution over toxic sequences
for which the model achieves low loss.333All code is available at
https://anonymous.4open.science/r/circuit-breaking-5DE5/.

As an approximation of our train set $D_{\text{train}}$, we use 10,000 samples
from OpenWebText (OWT) (Gokaslan & Cohen, 2019). See Appendix E for results in
removing a sub-class in an image classification model.

#### Constructing a bad behavior dataset.

We sample excerpts from highly toxic comments posted to the Politically
Incorrect board of 4chan imageboard forum (Papasavva et al., 2020). We sample
from posts assigned a toxicity score of greater than 0.9, as calculated by
Google’s Perspective API Toxicity V6 (Google, 2023).

###  4.1 Learning Edge Mask Details

Similar to (Goldowsky-Dill et al., 2023; Wang et al., 2022), we write GPT-2 as
a graph consisting of the input, the output, attention heads, and MLPs (158
nodes total) by considering a “residual rewrite” of the model’s computational
structure. The canonical description of a transformer model expresses the
attention head $A_{i,j}$ (the $j$th attention head in layer $i$) as taking an
argument $R_{i-1}$, the residual from the previous layer. However, since
$R_{0}=I$ (where $I$ represents the input embeddings) and
$R_{i}=R_{i-1}+\sum_{j}A_{i,j}+M_{i}$ (where $M_{i}$ is the output of the MLP
node in layer $i$), we can instead consider attention head $A_{i,j}$ as
operating on the sum
$S^{A}_{i}=I+\sum_{i^{\prime}<i}\left(M_{i^{\prime}}+\sum_{j^{\prime}}A_{i^{\prime},j^{\prime}}\right)$,
and taking all nodes in previous layers as separate input arguments.
Similarly, we can consider MLP node $M_{i}$ as operating on the sum
$S^{M}_{i}=I+\sum_{i^{\prime}<i}M_{i^{\prime}}+\sum_{i^{\prime}\leq
i}\sum_{j^{\prime}}A_{i^{\prime},j^{\prime}}$, and the output node as
operating on the sum of the input embeddings and all attention head and MLP
outputs. In total, this residual rewrite gives us a nearly-dense graph
containing 11,611 edges: one between every pair of (attention head, MLP,
input, and output) nodes, except for attention heads in the same layer, which
do not communicate with each other. Concretely, ablating an edge from
$A_{i^{\prime},j^{\prime}}$ to $A_{i,j}$ entails replacing the
$A_{i^{\prime},j^{\prime}}$ term in $S^{A}_{i}$ for the input to attention
head $A_{i,j}$ with zero (for zero ablation) or the mean value of head
$A_{i^{\prime},j^{\prime}}$ (for mean ablation).

We train two ablated models using a continuous edge mask. First, we train a
zero-ablation mask against $\mathcal{L}(W_{\text{mask}};\alpha,\lambda,R)$
described in equation 4, with $\alpha=0.2$, $\lambda(t)=(t-20)/10000$, and
$R(W_{\text{mask}})=\sum_{e\in E_{G}}w_{e}$. This search process finds a mask
that ablates 12 edges (Figure 2) and mitigates toxicity while preserving
coherence. Second, we train a mean-ablation mask with $\alpha=0.15$ and using
the same hyperparameters otherwise, which finds a mask that ablates 84 edges
and produces a similar effect.

As a baseline, we fine-tune on the loss given by Equation 1 directly, with
$\alpha=0.2$. We use early stopping with a validation set to prevent
overfitting.444We note this is a stronger baseline than naively training for
high loss on our bad behavior set as done in (Ilharco et al., 2023), which we
call “gradient ascent” in Table 1. We also compare to task arithmetic (Ilharco
et al., 2023) (Section 3.2).

|  Toxic-loss |  Toxic-loss (filtered) |  Toxic generation |  Toxic generation (filtered) |  Incoherence |  TPP Incoherence
---|---|---|---|---|---|---
Original | 4.954 | 4.435 | 0.453 | 0.944 | 4.264 | 4.617
Gradient Ascent | 21.339 | 20.980 | 0.015 | 0.013 | 15.287 | 18.415
Task Arithmetic | 5.357 | 4.827 | 0.351 | 0.631 | 4.427 | 4.731
Joint Fine-Tuned | 11.817 | 13.020 | 0.009 | 0.008 | 4.240 | 7.402
Ablated (12 edges) | 5.027 | 4.486 | 0.328 | 0.567 | 4.280 | 4.623
Ablated (84 edges) | 4.895 | 4.470 | 0.280 | 0.441 | 4.180 | 4.515
Table 1: Toxic-loss measures the model’s loss on toxic prompts. Toxic
generation measures the average toxicity score of model generations on toxic
prompts, according to the Detoxify classifier. The filtered columns denote the
loss or generation toxicity on test samples filtered by the original model
achieving low loss ($<5$) or highly toxic generation ($>.9$). Incoherence
measures the model’s loss on OWT. Toxic Pre-Pended (TPP) incoherence measures
the model’s loss after on OWT sequences that have been preceded by toxic text.

###  4.2 Evaluation Metrics

Following Definition 3.1, we evaluate both the model’s avoidance of toxic
generation (efficacy) and the detriment to other behaviors (specificity).
Since our goal is for the ablated model to achieve high loss on all toxic
sequences (i.e. minimizing its probability of predicting subsequent tokens
that would cause the sequence to be toxic), we evaluate efficacy in a few
ways. First, we consider the ablated model’s loss on with-held toxic text and
in particular its loss on sequences for which the original model achieves low
($<5$) loss. Second, we consider the toxicity of the model’s completions when
prompted with toxic text, as measured by the score in $[0,1]$, 0 being the
least toxic, given by the toxic-comment classifier Detoxify. We emphasize the
toxicity of model completions on the specific prompts for which the original
model produces highly toxic ($>0.9$) output.

We evaluate specificity by using the perplexity on withheld sequences from
OWT, along with the perplexity on withheld OWT sequences prepended with toxic
content. The original model produces low loss (4.617) on these sequences, and
we choose to highlight the behavior of retaining coherence when prompted with
toxic text as one that is particularly likely to be inadvertently removed when
editing the model to produce high loss on toxic text.

###  4.3 Results

Results are shown in Table 1. We train a model with 12 edges zero-ablated that
substantially mitigates toxic generation, decreasing the average toxicity
score on model generations for toxic prompts from 0.458 to 0.328 and in
particular for the most toxic-inducing prompts from 0.944 to 0.567. This
minimal edge ablation outperforms task arithmetic on every efficacy and
specificity metric, and causes a lower increase in incoherence following toxic
prompts than joint fine-tuning, though it does not eradicate the model’s
toxicity. Our mean-ablation mask with 84 edges achieves a similar result,
greatly mitigating toxic generations without detracting from the model’s other
behaviors.

##  5 Related Work

#### Causal mediation for circuit analysis.

Causal mediation (Pearl, 2009; Iwasaki & Simon, 1994) has been proposed as a
framework for evaluating mechanistic causal explanations for model outputs
(Goldowsky-Dill et al., 2023; Geiger et al., 2023a; Vig et al., 2020).
Experimental evaluation for causal explanations involves performing a set of
ablation experiments to check whether they match hypothesized effects. For
example, ablating allegedly unimportant paths should have little impact on the
target behavior. Previous work has used the causal mediation framework to
discover circuits, including in transformers (Chan et al., 2022; Wang et al.,
2022; Nanda et al., 2023).

Existing causal mediation tests and circuit discovery methods built upon these
tests evaluate whether a given set of edges are sufficient for a given model
behavior (i.e. if they contain a vertical path along the circuit), while our
circuit breaking technique finds a set of edges that are necessary for the
behavior (i.e. a horizontal “cut” through the circuit).

#### Automated circuit discovery.

Recent work has explored automated approaches to discovering circuits,
including greedy algorithms which crawl the computational graph and remove
edges which preserve behavior above a fixed threshold (Conmy et al., 2023),
and gradient descent-based methods which use interchange intervention training
(Geiger et al., 2022) to learn alignments between a source model and a
proposed high-level causal model (Geiger et al., 2023b). Our work differs in
attempting to find neither single features (Vig et al., 2020; Gurnee et al.,
2023) nor full computational circuits (Geiger et al., 2023b; Goldowsky-Dill et
al., 2023; Wang et al., 2022); instead we discover edges where removing their
causal effect breaks a given behavior.

#### Weight-masking and model pruning.

Much prior work has sought to compress models by masking parameters (LeCun et
al., 1989; Hassibi & Stork, 1992). Most relevant to our work are approaches
which learn masks from data by encouraging sparsity and preserving performance
(Louizos et al., 2017; Wang et al., 2019; Cao et al., 2021). In our work, we
disincentivize sparsity (since we want fewer ablations), and use an objective
function tailored to removing a specific behavior instead of preserving
general performance. Additionally, our edge-masking technique is more general
than weight-masking, since we can ablate internal connections between high-
level model components that do not correspond directly to particular weights,
such as communication channels between pairs of attention heads. Finally, we
prune using mean ablation instead of zero ablation.

#### Model editing to change or remove behaviors.

Recent work has made changes to model behavior by making targeted edits to
model weights (Meng et al., 2022) or activations (Hernandez et al., 2023),
which differ from our goal of removing behaviors. (Gandikota et al., 2023)
propose a fine-tuning approach to erasing concepts from diffusion models.
(Elazar et al., 2021) remove information from a language model’s
representation by iteratively learning linear probes to extract the
information and projecting onto the null space. Compared to such work, we
consider coarser ablations, allow editing around multiple components, and seek
to break behaviors as opposed to erasing information. Like us, (Ilharco et
al., 2023) attempt to remove the toxic generation behavior in GPT-2, but do so
by fine-tuning on bad behavior and subtracting the weight-difference from the
original model.

##  6 Conclusion

Using a small dataset of examples of inputs on which a neural network exhibits
a “bad behavior,” we find that our method can make high-level modifications to
the network that mitigate the bad behavior on the provided examples,
generalize to removing the bad behavior across other inputs that trigger it,
and cause only small amounts of damage to the model’s performance on all other
inputs (see D for limitations). We conjecture that model editing may be an
alternate tool for targeted behavioral modification to fine-tuning, and
encourage future work further investigating our approach.

## References

  * Bourtoule et al. (2021) Bourtoule, L., Chandrasekaran, V., Choquette-Choo, C. A., Jia, H., Travers, A., Zhang, B., Lie, D., and Papernot, N.  Machine unlearning.  In _2021 IEEE Symposium on Security and Privacy (SP)_ , pp. 141–159. IEEE, 2021.
  * Brown et al. (2020) Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D.  Language models are few-shot learners, 2020.
  * Cammarata et al. (2020) Cammarata, N., Carter, S., Goh, G., Olah, C., Petrov, M., Schubert, L., Voss, C., Egan, B., and Lim, S. K.  Thread: circuits.  _Distill_ , 5(3):e24, 2020.
  * Cao et al. (2021) Cao, S., Sanh, V., and Rush, A. M.  Low-complexity probing via finding subnetworks.  _arXiv preprint arXiv:2104.03514_ , 2021.
  * Chan et al. (2022) Chan, L., Garriga-Alonso, A., Goldowsky-Dill, N., Greenblatt, R., Nitishinskaya, J., Radhakrishnan, A., Shlegeris, B., and Thomas, N.  Causal Scrubbing: a method for rigorously testing interpretability hypotheses [Redwood Research], December 2022.  URL https://www.alignmentforum.org/posts/JvZhhzycHu2Yd57RN/causal-scrubbing-a-method-for-rigorously-testing.
  * Chen et al. (2018) Chen, J., Song, L., Wainwright, M. J., and Jordan, M. I.  Learning to explain: An information-theoretic perspective on model interpretation, 2018.
  * Conmy et al. (2023) Conmy, A., Mavor-Parker, A. N., Lynch, A., Heimersheim, S., and Garriga-Alonso, A.  Towards automated circuit discovery for mechanistic interpretability.  _arXiv preprint arXiv:2304.14997_ , 2023.
  * Covert et al. (2020) Covert, I., Lundberg, S., and Lee, S.-I.  Understanding global feature contributions with additive importance measures, 2020.
  * Covert et al. (2022) Covert, I., Lundberg, S., and Lee, S.-I.  Explaining by removing: A unified framework for model explanation, 2022.
  * Dabkowski & Gal (2017) Dabkowski, P. and Gal, Y.  Real time image saliency for black box classifiers.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/0060ef47b12160b9198302ebdb144dcf-Paper.pdf.
  * Deng (2012) Deng, L.  The mnist database of handwritten digit images for machine learning research.  _IEEE Signal Processing Magazine_ , 29(6):141–142, 2012.
  * Devlin et al. (2019) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
  * Elazar et al. (2021) Elazar, Y., Ravfogel, S., Jacovi, A., and Goldberg, Y.  Amnesic probing: Behavioral explanation with amnesic counterfactuals.  _Transactions of the Association for Computational Linguistics_ , 9:160–175, 2021.
  * Fong & Vedaldi (2017) Fong, R. C. and Vedaldi, A.  Interpretable explanations of black boxes by meaningful perturbation.  In _2017 IEEE International Conference on Computer Vision (ICCV)_. IEEE, oct 2017.  doi: 10.1109/iccv.2017.371.  URL https://doi.org/10.1109%2Ficcv.2017.371.
  * Gandikota et al. (2023) Gandikota, R., Materzynska, J., Fiotto-Kaufman, J., and Bau, D.  Erasing concepts from diffusion models.  _arXiv preprint arXiv:2303.07345_ , 2023.
  * Geiger et al. (2022) Geiger, A., Wu, Z., Lu, H., Rozner, J., Kreiss, E., Icard, T., Goodman, N., and Potts, C.  Inducing causal structure for interpretable neural networks.  In _International Conference on Machine Learning_ , pp. 7324–7338. PMLR, 2022.
  * Geiger et al. (2023a) Geiger, A., Potts, C., and Icard, T.  Causal abstraction for faithful model interpretation, 2023a.
  * Geiger et al. (2023b) Geiger, A., Wu, Z., Potts, C., Icard, T., and Goodman, N. D.  Finding alignments between interpretable causal variables and distributed neural representations.  _arXiv preprint arXiv:2303.02536_ , 2023b.
  * Ghorbani & Zou (2020) Ghorbani, A. and Zou, J.  Neuron shapley: Discovering the responsible neurons, 2020.
  * Gokaslan & Cohen (2019) Gokaslan, A. and Cohen, V.  Openwebtext corpus.  http://Skylion007.github.io/OpenWebTextCorpus, 2019.
  * Golatkar et al. (2020) Golatkar, A., Achille, A., and Soatto, S.  Eternal sunshine of the spotless net: Selective forgetting in deep networks.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 9304–9312, 2020.
  * Goldowsky-Dill et al. (2023) Goldowsky-Dill, N., MacLeod, C., Sato, L., and Arora, A.  Localizing model behavior with path patching.  _arXiv preprint arXiv:2304.05969_ , 2023.
  * Google (2023) Google.  Perspective API, 2023.  URL https://perspectiveapi.com/.
  * Guan et al. (2022) Guan, J., Tu, Z., He, R., and Tao, D.  Few-shot backdoor defense using shapley estimation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 13358–13367, 2022.
  * Gurnee et al. (2023) Gurnee, W., Nanda, N., Pauly, M., Harvey, K., Troitskii, D., and Bertsimas, D.  Finding neurons in a haystack: Case studies with sparse probing.  _arXiv preprint arXiv:2305.01610_ , 2023.
  * Hanu & Unitary team (2020) Hanu, L. and Unitary team.  Detoxify.  Github. https://github.com/unitaryai/detoxify, 2020.
  * Hassibi & Stork (1992) Hassibi, B. and Stork, D.  Second order derivatives for network pruning: Optimal brain surgeon.  _Advances in neural information processing systems_ , 5, 1992.
  * Hernandez et al. (2023) Hernandez, E., Li, B. Z., and Andreas, J.  Measuring and manipulating knowledge representations in language models.  _arXiv preprint arXiv:2304.00740_ , 2023.
  * Ilharco et al. (2023) Ilharco, G., Ribeiro, M. T., Wortsman, M., Gururangan, S., Schmidt, L., Hajishirzi, H., and Farhadi, A.  Editing Models with Task Arithmetic, March 2023.  URL http://arxiv.org/abs/2212.04089.  arXiv:2212.04089 [cs].
  * Iwasaki & Simon (1994) Iwasaki, Y. and Simon, H. A.  Causality and model abstraction.  _Artificial intelligence_ , 67(1):143–194, 1994.
  * Lakkaraju et al. (2016) Lakkaraju, H., Bach, S. H., and Leskovec, J.  Interpretable decision sets: A joint framework for description and prediction.  In _Proceedings of the 22Nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , KDD ’16, pp. 1675–1684, New York, NY, USA, 2016. ACM.  ISBN 978-1-4503-4232-2.  doi: 10.1145/2939672.2939874.  URL http://doi.acm.org/10.1145/2939672.2939874.
  * LeCun & Cortes (2010) LeCun, Y. and Cortes, C.  MNIST handwritten digit database.  http://yann.lecun.com/exdb/mnist/, 2010.  URL http://yann.lecun.com/exdb/mnist/.
  * LeCun et al. (1989) LeCun, Y., Denker, J., and Solla, S.  Optimal brain damage.  _Advances in neural information processing systems_ , 2, 1989.
  * Louizos et al. (2017) Louizos, C., Welling, M., and Kingma, D. P.  Learning sparse neural networks through $l\\_0$ regularization.  _arXiv preprint arXiv:1712.01312_ , 2017.
  * Lundberg & Lee (2017) Lundberg, S. M. and Lee, S.-I.  A unified approach to interpreting model predictions.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper/2017/file/8a20a8621978632d76c43dfd28b67767-Paper.pdf.
  * Meng et al. (2022) Meng, K., Bau, D., Andonian, A., and Belinkov, Y.  Locating and editing factual associations in gpt.  _Advances in Neural Information Processing Systems_ , 35:17359–17372, 2022.
  * Nanda (2023) Nanda, N.  Attribution patching: Activation patching at industrial scale, 2023.  URL https://www.neelnanda.io/mechanistic-interpretability/attribution-patching.
  * Nanda et al. (2023) Nanda, N., Chan, L., Liberum, T., Smith, J., and Steinhardt, J.  Progress measures for grokking via mechanistic interpretability.  _arXiv preprint arXiv:2301.05217_ , 2023.
  * Olah (2022) Olah, C.  Mechanistic interpretability, variables, and the importance of interpretable bases.  _Transformer Circuits Thread(June 27). http://www. transformer-circuits. pub/2022/mech-interp-essay/index. html_ , 2022.
  * Olah et al. (2020) Olah, C., Cammarata, N., Schubert, L., Goh, G., Petrov, M., and Carter, S.  Zoom in: An introduction to circuits.  _Distill_ , 2020.  doi: 10.23915/distill.00024.001.  https://distill.pub/2020/circuits/zoom-in.
  * OpenAI (2023) OpenAI.  Gpt-4 technical report, 2023.  URL https://arxiv.org/abs/2303.08774.
  * Ouyang et al. (2022) Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al.  Training language models to follow instructions with human feedback.  _Advances in Neural Information Processing Systems_ , 35:27730–27744, 2022.
  * Papasavva et al. (2020) Papasavva, A., Zannettou, S., Cristofaro, E. D., Stringhini, G., and Blackburn, J.  Raiders of the lost kek: 3.5 years of augmented 4chan posts from the politically incorrect board, 2020.
  * Pearl (2009) Pearl, J.  _Causality_.  Cambridge university press, 2009.
  * Poursabzi-Sangdeh et al. (2018) Poursabzi-Sangdeh, F., Goldstein, D. G., Hofman, J. M., Vaughan, J. W., and Wallach, H. M.  Manipulating and measuring model interpretability.  _arXiv_ , abs/1802.07810, 2018.  URL http://arxiv.org/abs/1802.07810.
  * Radford et al. (2019) Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Räukur et al. (2022) Räukur, T., Ho, A., Casper, S., and Hadfield-Menell, D.  Toward transparent ai: A survey on interpreting the inner structures of deep neural networks.  _arXiv preprint arXiv:2207.13243_ , 2022.
  * Sekhari et al. (2021) Sekhari, A., Acharya, J., Kamath, G., and Suresh, A. T.  Remember what you want to forget: Algorithms for machine unlearning.  _Advances in Neural Information Processing Systems_ , 34:18075–18086, 2021.
  * Shapley (1997) Shapley, L. S.  A value for n-person games.  _Classics in game theory_ , 69, 1997.
  * Stiennon et al. (2022) Stiennon, N., Ouyang, L., Wu, J., Ziegler, D. M., Lowe, R., Voss, C., Radford, A., Amodei, D., and Christiano, P.  Learning to summarize from human feedback, 2022.
  * Touvron et al. (2023) Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G.  Llama: Open and efficient foundation language models, 2023.
  * Vig et al. (2020) Vig, J., Gehrmann, S., Belinkov, Y., Qian, S., Nevo, D., Singer, Y., and Shieber, S.  Investigating gender bias in language models using causal mediation analysis.  _Advances in neural information processing systems_ , 33:12388–12401, 2020.
  * Wang et al. (2022) Wang, K., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J.  Interpretability in the wild: a circuit for indirect object identification in gpt-2 small, 2022.
  * Wang et al. (2019) Wang, Z., Wohlwend, J., and Lei, T.  Structured pruning of large language models.  _arXiv preprint arXiv:1910.04732_ , 2019.
  * Wu & Wang (2021) Wu, D. and Wang, Y.  Adversarial neuron pruning purifies backdoored deep models.  _Advances in Neural Information Processing Systems_ , 34:16913–16925, 2021.
  * Yoon et al. (2019) Yoon, J., Jordon, J., and van der Schaar, M.  INVASE: Instance-wise variable selection using neural networks.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=BJg_roAcK7.
  * Zhou et al. (2015) Zhou, B., Khosla, A., Lapedriza, A., Oliva, A., and Torralba, A.  Object detectors emerge in deep scene cnns, 2015.
  * Ziegler et al. (2022) Ziegler, D. M., Nix, S., Chan, L., Bauman, T., Schmidt-Nielsen, P., Lin, T., Scherlis, A., Nabeshima, N., Weinstein-Raun, B., de Haas, D., Shlegeris, B., and Thomas, N.  Adversarial training for high-stakes reliability, 2022.

##  Appendix A Additional Related Work

#### Unlearning.

Machine unlearning aims to modify a model to match the behavior of a model
which had not seen certain data points (Sekhari et al., 2021; Bourtoule et
al., 2021; Golatkar et al., 2020). However, a key difference in our setting is
that we are not able to enumerate the full set of undesirable data points in
our training set.

#### Backdoor removal.

(Wu & Wang, 2021) learn a binary mask to zero ablate neurons sensitive to
adversarial perturbations, and finds that doing so removes injected backdoors.
(Guan et al., 2022) target backdoors by estimating Shapley importance values
(Shapley, 1997) for every edge and then zero ablating neurons which have a
high attack success rate attribution score, finding they are able to remove
backdoors with very limited (and sometimes no) data. We believe our technique
could be effective for disabling the activation of backdoor mechanisms and
find this application a promising direction for future work.

##  Appendix B Writing models as computational graphs.

Figure 3: We can subdivide an attention head into its own computational graph.

We can write any model as a connected directed acyclic graph (DAG) with source
nodes representing the model’s (typically vector-valued) input, sink nodes
representing the model’s output, and intermediate nodes representing units of
computation. Each intermediate node represents a function of the values of its
parent nodes. On a forward pass, given values for its input nodes, the model
computes the value of each node in any topologically sorted order until it has
computed the value of the output nodes.

For any model, there are many equivalent graphs that faithfully represent its
computation. In particular, a computational graph can represent a model at
varying levels of detail. At one extreme, intermediate nodes can designate
individual additions, multiplications, and nonlinearities – such a graph would
have at least as many nodes as model parameters. On the other hand, many model
architectures have self-contained computational modules, which allows them to
be represented by graphs that convey a high level of abstraction. For example,
in convolutional networks, intermediate nodes can represent convolutional
filters and pooling layers, while in transformer models (Devlin et al., 2019),
the natural high-level computational units are attention heads and multi-layer
perceptron (MLP) modules. To be more granular, we can subdivide each attention
head node into nodes that compute queries, keys, and values and combine them
into attention patterns (Figure 3).

##  Appendix C Ablation Types

One mode of ablating an edge is zero ablation, in which we compute the value
of its destination node as if the value of its source node were zero. However,
a value of zero on an intermediate node can sometimes be highly unusual, and
can thus in some cases convey a strong idiosyncratic signal to the destination
node.

One other technique is mean ablation, in which we compute the destination node
as if the source node’s value were set to its mean value over the training
set. Mean ablation arguably better captures a lack of information flow: if the
specific value of the source node were withheld from the destination node, the
source node’s mean value would be the most general estimate of its true value.

##  Appendix D Limitations

We ablate edges by setting their input values to zero and the train-set mean.
However, recent work has argued that ablating model components with random
samples from their marginal distributions may be preferable and that mean
ablation may lead to out of distribution resampling (Goldowsky-Dill et al.,
2023). Additionally, circuit-breaking interventions on the model could be made
even more surgical by using more granular model nodes and edges (for example,
splitting attention heads into query, key, and value nodes). Finally, our
results could be strengthened by considering stronger baselines and additional
approaches to learning binary masks.

##  Appendix E Additional Experiments: Breaking Digit Classification in an MLP

Figure 4: The learned mask for MNIST classification over the course of
training. Note that versions of this mask in the middle of training are
allowed to partially ablate each edge, so “Edges Ablated” is calculated by
summing the coefficients assigned to the ablation value. The “train” points
are those that the MLP was trained on, and the “test” points are those it was
not. The “bad behaviors” line indicates its accuracy on the 30 exemplar
digits.

We train a one-hidden-layer MLP with a 50 hidden neurons to classify the
handwritten digits of MNIST, then use a small (30 example) dataset of a
particular digit (say, 3) to remove the model’s ability to correctly classify
that digit. We consider the most granular computational graph for the MLP with
one node for each pixel of the input, one for each hidden neuron, and one for
each output neuron. The graph contains an edge corresponding to each weight in
the network. To prevent our learned mask from simply ablating the edges
feeding into the output neuron corresponding to 3, we arbitrarily pair digits
and merge their labels so that the MLP has only 5 output neurons rather than
10. This pairing forces the network to retain edges to the output neuron that
aid in correctly classifying the digit that is paired with 3, while not using
the neuron for the 3 inputs.

We search for a binary mask over edges by training a continuous edge mask
against $\mathcal{L}(W_{\text{mask}};\alpha,\lambda,R)$ described in Equation
4. Specifically, we use $\alpha=0.3$, $\lambda(t)=t$, and
$R(W_{\text{mask}})=\sum_{e\in E_{G}}\sqrt{1-w_{e}}$. The sublinear cost
imposed by $R$ incentivizes masks that are binary and ablate few edges;
conceptually, if the mask were half-ablating two edges, it would receive a
lower penalty for instead ablating one edge completely. We set the rounding
threshold $\tau=0.5$.

Using this technique, we find a binary mask that ablates 400 of the model’s
38K edges, bringing its accuracy on held-out “3”s from near-perfect to 21%
(20% is random classification on this task), while accuracy on other (held-
out) inputs stays high (dropping from 99% to 97%). We consider this a modest
success for both the efficacy of the edit (i.e. its ability to transfer to
other inputs on which the model exhibits the bad-behavior of classifying a “3”
correctly) and also its specificity (i.e. the model’s continued ability to
classify non-“3”s correctly) – see Figure 4.