Efficient Automated Circuit Discovery in Transformers using Contextual Decomposition
Authors: Aliyah R. Hsu
Georgia Zhou
Yeshwanth Cherapanamjeri
Yaxuan Huang
Anobel Odisho
Peter R. Carroll
Bin Yu
Venue: ICLR 2025 Poster (2025)
Keywords: Automated Circuit Discovery
Explainable AI
Interpretation
Machine Learning
Language Models
Transformers
Primary area: interpretability and explainable AI

Abstract: Automated mechanistic interpretation research has attracted great interest due to its potential to scale explanations of neural network internals to large models. Existing automated circuit discovery work relies on activation patching or its approximations to identify subgraphs in models for specific tasks (circuits). They often suffer from slow runtime, approximation errors, and specific requirements of metrics, such as non-zero gradients.
In this work, we introduce contextual decomposition for transformers (CD-T) to build interpretable circuits in large language models. CD-T can produce circuits at any level of abstraction and is the first to efficiently produce circuits as fine-grained as attention heads at specific sequence positions.
CD-T is compatible to all transformer types, and requires no training or manually-crafted examples.
CD-T consists of a set of mathematical equations to isolate contribution of model features. Through recursively computing contribution of all nodes in a computational graph of a model using CD-T followed by pruning, we are able to reduce circuit discovery runtime from hours to seconds compared to state-of-the-art baselines.
On three standard circuit evaluation datasets (indirect object identification, greater-than comparisons, and docstring completion),
we demonstrate that CD-T outperforms ACDC and EAP by better recovering the manual circuits with an average of 97% ROC AUC under low runtimes.
In addition, we provide evidence that faithfulness of CD-T circuits is not due to random chance by showing our circuits are 80% more faithful than random circuits of up to 60% of the original model size. 
Finally, we show CD-T circuits are able to perfectly replicate original models' behavior(faithfulness  = 1) using fewer nodes than the baselines for all tasks.
Our results underscore the great promise of CD-T for efficient automated mechanistic interpretability, paving the way for new insights into the workings of large language models.

Published as a conference paper at ICLR 2025
EFFICIENT AUTOMATED CIRCUIT DISCOVERY IN TRANSFORMERS USING CONTEXTUAL DECOMPOSI- TION
Aliyah R. Hsu
Georgia Zhou Department of EECS UC Berkeley
Yeshwanth Cherapanamjeri CSAIL, MIT
Department of EECS
UC Berkeley al lyyhhauu@eereele.yyccoo
Yaxuan Huang
Anobel Y. Odisho & Peter R. Carroll
Department of Statistics
Department of Urology, Epidemiology and Biostatistics UC San Francisco
UC Berkeley
Bin Yu
Department of Statistics, EECS Center for Computational Biology UC Berkeley
ABSTRACT
Automated mechanistic interpretation research has attracted great interest due to its potential to scale explanations of neural network internals large models. Existing automated circuit discovery work relies on activation patching or its ap- proximations identify subgraphs in models for specific tasks (circuits). They often suffer from slow runtime, approximation errors, and specific requirements of metrics, such as non-zero gradients. In this work, we introduce contextual decomposition for transformers (CD-T) to build interpretable circuits in large lan- guage models. CD-T can produce circuits at any level abstraction and is the first to efficiently produce circuits fine-grained as attention heads at specific sequence positions. CD-T is compatible to all transformer types, and requires no training manually-crafted examples. CD-T consists of set of mathematical equations to isolate contribution of model features. Through recursively comput- ing contribution of all nodes in computational graph of model using CD-T followed by pruning, we are able to reduce circuit discovery runtime from hours seconds compared state-of-the-art baselines. On three standard circuit eval- uation datasets (indirect object identification, greater-than comparisons, and doc- string completion), we demonstrate that CD-T outperforms ACDC and EAP by better recovering the manual circuits with an average of 97% ROC AUC under low runtimes. In addition, we provide evidence that faithfulness of CD-T eircuits is not due to random chance by showing our circuits are 80% more faithful than random circuits up to 60% the original model size. Finally, we show CD-T circuits are able to perfectly replicate original models' behavior (faithfulness = 1) using fewer nodes than the baselines for all tasks. Our results underscore the great promise of CD-T for efficient automated mechanistic interpretability, paving the way for new insights into the workings of large language models. All code for using CD-T and reproducing results is made available Github.
INTRODUCTION
Transformers (Vaswani et al. 2017) have recently demonstrated impressive predictive capabili- ties (Brown et al. 2020b) by learning intricate nonlinear relationships between features. How-
https://github.ce

---

Published as a conference paper at ICLR 2025
ever, the challenge of comprehending these relationships has resulted in transformers largely being thought of as black boxes. Despite this, transformers are increasingly utilized in high-stakes do- mains such medicine (e.g. medical image analysis (He et al. 2023) and science (c.g. protein structure prediction aiding drug discovery (Jumper et al. (021)) D. This highlights the necessity of understanding model behaviors. Mechanistic interpretability. ity, study to explain behaviors of neural networks in terms of their internal components, is at the frontier of interpretability research (Geiger ct al. 2021: Geva et al. 2020 Räuker al. 2023), as it uniquely provides an avenue for guiding manual modifications (Elhage et al. 2021| Vig et al. 2020b) and reverse- reverse-engineering solutions (El- hage et al., 2022 Meng et al. 20233.
One main direction in mechanistic interpretability research is circuit discovery, where researchers aim to identify a computational subgraph (circuit) in neural networks responsible for solving spe- cific task (Olah et al.) 2020). Prior circuit discovery work in language models has found subgraphs of attention heads and multi-layer perceptrons (MLPs) that partially or fully explain model behav- iors on tasks such as indirect object identification, doestring completion, and greater-than compar- isons (Wagg et al. 2023; Hanna et al.] 2023; Nanda et al. 2023; Heimersheim & Janiak 2023). However, most of them require tedious manual inspection, and thus limited smaller mod- els (Lindner et al., (023). Automated Circuit DisCovery (ACDC) (Conmy et al.] (023) has been proposed scalable end-to-end technique to identify circuits in models of arbitrary sizes by au- tomating the pruning of unimportant edges via activation patching. Although effective, ACDC suf- fers from slow runtimes. More recent work has focused on improving the efficiency, for example, through linear approximating activation patching with attribution patching (Syed et al. 2023 Hanna al. 2024), or training sparse autoencoders (SAEs) learn attention head output patterns (0 O'Neill & Bui 2024). While faster variants have been leveraged with some success, they still suffer from several drawbacks, such as approximation errors (Syed et al., 2023 Hanna et al. (2024) and the
need of model training and carefully designed examples by human (O O'Neill& & Bui 2024).
Our main contribution in this work the introduction of contextual decomposition for transform- ers (CD-T), mathematical or analytical method to efficiently compute the contribution of model features to other model features within a transformer. This work builds on CD methods for CNNs and RNNs to examine local importance proposed in prior work (Murdoch et al. 2018 Singh al., 2018 Jumelet et al.; 2019). We significantly extend CD by defining novel mathematical decom- position principles which can be applied to transformers of all types (e.g., decoder- based, encoder- based), because these principles are indispensable for CD to scale to modern state-of-thh-ar (SOTA) deep learning models as they are mostly based on transformers.
We are also the first, to our knowledge, to show the utility of CD methods for mechanistic inter- pretability by proposing an automated circuit discovery algorithm utilizing CD-T. That is, CD-T is able to discover circuits with finer granularity by further splitting into sequence positions while at the same time reduce runtime from hours to seconds, when it compared with prior work (Syed al., 2023 Conmy et al. 2023). CD-T is automatic or does not require any training or manually-cratir examples by humans. CD-T is fundamentally compatible with all common transformer architec- tures: in particular, CD-T supports circuit discovery in both BERT-like (Devlin et al. 2019) and GPT-like (Brown et al. 2020a) architectures.
Specifically, for evaluation, we use three standard circuit evaluation datasets: indirect object identi- fication (IOI), greater-than comparisons threa, and docstring completion (Docstring). We compare recovered circuits by CD-T on these datasets with two SOTA baselines, ACDC (Conmy) ct al. 2023) and EAP (Syed et al. 2023). We demonstrate that CD-T outperforms the baselines by better identifying attention heads in the manual circuits with average of 97% ROC AUC due to its outstanding ability in recovering negative and supporting heads efficiently. EAP obtains similar runtime as CD-T but performs slightly worse in ROC AUC for Greater-than and IOL. ACDC the least efficient as it often takes hours to identify a circuit, yet with no better ROC AUC on all tasks compared to other methods. In addition, we demonstrate faithfulness of CD-T circuits not due to random chance by showing our circuits are 80% more faithful than random circuits of up to 60% of the original model size. Finally, we show CD-T circuits are able perfectly replicate original models' behavior (faithfulness = 1) using fewer nodes than the baselines for all tasks.
In our experiments, "nodes" refers to attention heads attention heads' output at specific sequence. See the general definition in Section 3
2

---

Published as a conference paper at ICLR 2025
2
RELATED WORK
2.1 CAUSAL INTERPRETATION
Previous efforts to interpret neural networks have frequently drawn ideas from causal infer- ence (Pearl 2009). Much of this work emphasizes counterfactual reasoning or enforcing causal constraints on computations
explain model outputs (Pearl 2009 Feder et al. [2020] Geiger al.,
2021
Wu et al. 202T; Kaddour et al. 2022). A related line of research (Vig et al. 2020a] Stollo
ct al. 2023 Feng & Steinhardi 2023) takes different approach by treating internal components (or
nodes, if the model is viewed as a computational graph) as mediators and conducting causal media- tion analysis, also called activation patching. In these studies, the contribution of nodes assessed by ablating subsets of them and measuring direct and indirect effects, with the direct effects serving as proxies for their overall contribution. In contrast to viewing nodes as mediators, Goldowsky- paths more expressive me-
ett al. (2023) and Wang et al. (2023) propose treating diators. They introduce the technique of path patching to explore how different edge subsets in a model's computational graph affect its behavior.
2.2 CIRCUIT DISCOVERY METHODS
Ablation-based methods are essential for identifying key components within models. Through ac- tivation patching, prior research on circuit discovery in language models has uncovered subgraphs of attention heads and MLPs that explain model behavior (Wang et al. 2023 Hanna et al. 2023 Heimersheim & Janiak 2023). However, most of these methods require multiple iterations man- ual inspection and ad-hoc analysis which is specific to the model and task, limiting their application to small numbers of circuits in smaller models (Lindner al., 2023). To address this, Conmy et al. (2023) introduced Automated Circuit Discovery (ACDC), which scales to models of any size by automatically calculating edge importance between attention heads and MI.ss based on model per- formance on chosen metric via activation patching. While ACDC is effective, it suffers from slow runtimes and fails to recover certain attention head subsets within the manual circuits. Syed et al. (023) introduced Edge Attribution Patching (EAP, ACDC++ later renamed), which improves ACDC's efficiency by using a linear approximation of activation patching. However, this reliance on linear approximations with true causal effects. EAP also struggles when the gradient of the performance metric is zero. To resolve these issues, [Hanna al. (2024) proposed Edge Attribution Patching with Integrated Gradients (EAP-II)) for more faithful circuit recovery. although it increases runtime by a constant factor. Another approach trains sparse autoencoders (SAEs) on attention head outputs and uses the SAB-learned features to map attention head contributions to identified circuits. However, SAE-
lead to overestimation of edge importance and weaker correlation
based methods either require carefully crafted uuman-gennrated examples (O'Neill & Buil (024) or require impractical amount of compute to train SAE (He et al. Marks Marks et al. 2024).
Another common limitation of prior methods is the limited level of granularity of the circuits they discover, because most work either doesn't 't support further attention head splitting at specific se- quence positions (Syed et al. 2023] Hanna et al. 2024: O Neill & Bui 2024), or it is fundamentally feasible but requires prohibitive amount of time to do so (Conmy et al., 2023), which drastically hinders more fine-grained mechanistic interpretability analysis. Our method CD-T addresses the limitations of previous circuit. discovery approaches by significantly reducing runtime, using mathe- matical decomposition equations with no approximations, allowing for arbitrary level of abstraction of circuits, and avoiding both model training and manually crafted examples (i.e. CD-T achieves all the above in an automated manner).
2.3 CONTEXTUAL DECOMPOSITION IN NEURAL NETWORKS
Contextual Decomposition (CD), introduced by Murdoch et al. (018)) 3 attributes local importance to input features in LSTMs by analytically decomposing the output without any changes to the under- lying model. For each word in a sentence, a forward pass computes activations for all cells and gates while partitioning each neuron's activation into parts influenced by selected token or phrase (de- noted as the relevant) and those that are not (denoted as the irrelevant). using factorization of the update equations for hidden states and cell states. Jumelet al.) (2019) propose Generalized Contextual Decomposition (GCD) for unidirectional LSTMs, applying it to phenomena like num- ber agreement and pronoun resolution. Singh et al. 2018) extend CD from LSTMs RNNs and
3

---

Published as a conference paper at ICLR 2025
CNNs and introduce hierarchical interpretation with feature clustering for improved input feature attributions. A distantly related work, contextual explainability (Bertolini et al. 2022). , focuses on feature attributions for each convolutional layer of CNNs by utilizing gradients the feature impor- tance measure. We further extend CD to transformers (CD-T) and we argue that our development of CD-T is significant advance with two novel contributions. First, we carefully design mathematical decomposition principles tailored for transformers of all types (e.g., ricooer-basee. encoder-based). which makes it able scale to larger models. Second, given prior work on CD only involves input feature attribution, to our knowledge, we are the first to demonstrate the utility of CD for mechanistic interpretability by proposing an automated circuit discovery algorithm that leverages CD-T.
3
OUR METHOD: CD-T AND MECHANISTIC INTERPRETABILITY
In this section, we first recall the basic operations of transformers (Section 3.1) We then provide technical details on contextual decomposition and its application to transformers (section 3.2) Fi- nally, we conclude with circuit discovery algorithm based on CD-T (section 3.3
3.1 TRANSFORMERS
The main technical element unifying transformer-based architectures is the attention mechanism. This mechanism allows contextually useful information to be transmitted from one position in the sequence any other position. We start by describing the operation of general attention mech- anism before discussing the modifications needed to specialize it for generation, exemplified by models such as GPT (Brown et al.) 2020a), and for sequence-to-sepue se-sueence transformation models such as BERT (Devlin et al. 2019). Formally, the attention mechanism computes vector of con- target vector
tributions from series of vectors {=r))== attention layer typically consists of a small number of heads, with each head parameterized by three functions, the key function fx. value function fur and query function fg. The output of the attention head to x is computed as follows where de denotes the output dimension of fa:
position p with 2i, E Rd. An
Vi E [4] : k Sm(x;), Oz Lg(x)) 0s So(x;)
Vi [1] : a = exp(a k2/ /d///d))
(2)
Ejetg exp(11 k/ /da)
y2 = Σ a₂u,
(3)
The functions fur fg are commonly simple learnable transformations such as linear transfor- mations or one hidden-layer MLPs. An attention layer is composed of a series of attention heads applied to every position in the sequence such that the concatenation of their outputs equals the input dimension d. As the output of attention layer is series of vectors of the same length and dimensionality as its input, this allows for stacking of these layers to build increasingly complex representations of the input so far.
The two classes of transformer models we consider in this work, sequence-to-sequence models such as BERT and generative models such as GPT mainly differ in the set of positions that the attention mechanism operates on. Concretely. BERT maps a fully specified input sequence to an output sequence of the sample length while generative models produce increasingly larger sequences autoregressively with the (f + 1) 1)th token generated based on the previous tokens. Since all the input tokens are specified for BERT-based models, the computation of the attention vector at position t utilizes representations at every position of the sequence. On the other hand, when generating the (t + 1)th, the positions at + 1 and beyond are not determined. As a consequence, only the representations for tokens before t are used to compute the attention vectors.
3.2 CONTEXTUAL DECOMPOSITION FOR TRANSFORMERS (CD-T)
We will now describe the general method of contextual decomposition (CD), as well as our extension to the transformer architecture and applications mechanistic interpretability.
4

---

Published as a conference paper at ICLR 2025
CD (Murdoch et al. 2018) was first proposed to compute contributions of input tokens to the output of LSTMs as a local interpretation method. CD divides each cell and hidden state into sum two parts: | part (said to be relevant), which contains the part of this particular state that stems from the input tokens of interest, and a part (said to be irrelevant), which contains information coming from tokens outside of the list of interest. ,3 is often initialized with masked input embeddings (I's specify where token positions of interest are, and O's the opposite), and the complement of 3.
Given a decomposition of vector x E Rd into relevant and irrelevant constituents / + Y - (here we use the word constituents instead of "components" to avoid name collision with the "com- ponents of a neural network"), and module f I + R CD consists of set of mathemati- cal equations to determine the decomposition of the output of a module f when given z as input: f(x) E T = 3, + %oo The general principle is to find symbolic expression for [() in terms of B and 7, and group the terms in the expression according to whether they are solely func- tion of 3 or not: the sum of terms solely relying on becomes Bo, and the sum of the remaining terms becomes %o. We refer interested readers to Murdoch al. (2018) and Singh et al. (2019) for further examples of decomposition formulas of various modules. These decomposition rules can be composed through multiple modules (i.e. if (1) 82, 72. and 9 %/2) B3, Y3+ then of ((f.))))) m)) = B3, 7a) in order to define decomposition rules for larger computational blocks, including entire neural networks.
Other authors have heuristically adjusted the decomposition rules for a module to reflect specific mechanisms by which (the relevance term in earlier layer in later layer i) that are unaccounted for by this rule. For example, |Sinhh et al. (2019) adjust the decomposition rule for affine transformations f(x) Wx + b. so that the bias term does not affect the relative magnitudes of B and 7/:
can affect (the relevance term
W B
B. = W3-1
(4)
313-11 + IW 31
|W
7/ Wh-1- |W3-11+ + [W
b.
(5)
We succinetly represent the above equations as 33, Linear DeriaareeDoppppppp n-1). where W and are understood to be tunable parameters of a neural network module.
Given the query, key, and value functions mentioned in Section 3.1 are linear transformations, the only module not accounted for in the context transformers the self-attention module. We handle the decomposition of the attention equations follows:
Bquery: "query = Linear Decomp(B:n. )im)
(6)
3key, Ikey Linear Decoorpecmpp..p (Bm )in)
(7)
Buatues Vvalue Linear Decomp(33n. Tim)
(8)
Buttention 9eayyquery /de
(9)
Nattention
= ((8key + +key) T (Aquery + /query)/ / dx) Battention
(10)
Sprohs Softmax M (atttenttottion) ))
(11)
Pprobs Softmax M Mask(Battenttt (atteention + Mattention)) Bprobs
(12)
32 Bprots * Beatue
(13)
7z (Sprors Tprohs) (Bratue + Vootue) - Bz
(14)
Boutpitts Toutpuu
Linear LinearDecome Decomp((2x....
(15)
where above Mask represents causal mask in appropriate architectures and Soflmaz is the soft- max along the same dimension as in a standard implementation of self-attention.
Although in the original setting of CD, "input" and "output" refer to an input sequence x and model's output logits, the method generalizes to arbitrary source and target activations in the model. Furthermore, in the original context, the "relevant" portion of the decomposition was equal to on a subsequence of the input and 0 elsewhere, but in a mechanistic interpretability context, the CD framework allows again as a computational graph, this gives us a method of calculating the effect of an arbitrary constituent of the activations of an arbitrary node on the activations of another node.
to decompose activations into any choice of 8, Thinking of the network
5

---

Published as a conference paper at ICLR 2025
3.2.1 CHOOSING DECOMPOSITION OF SOURCE NODES FOR MECHANISTIC INTERPRETABILITY
In a manner analogous to the choice of ablation method elsewhere in mechanistic interpretability. the correct choice of starting decomposition can be subtle, and sometimes specific to the task. A general principle is that the relevant constituent of the activations at some node is the constituent that we hypothesize has a counterfactual effect on the task; often, this means the deviation from the mean activations over some distribution. (It follows that the irrelevant constituent of the activation at this node is the mean activation, which matches intuitions about mean-ablation.) We provide details on the evaluation tasks and name the distributions used for ablation in Appendix A
3.3
AUTOMATING CIRCUIT DISCOVERY
3.3.1 DEFINITIONS OF CIRCUIT DISCOVERY
Before delving into our method, first formally define a circuit. To maintain consistency in no- tations with prior work (Shi et al., 2024 Elhage et al., 2021), we view model as a computational graph M, where nodes are activations of the model components in its forward pass (e.g. attention heads and MLPs) and edges are the interactions between those components, and a circuit C' is sub- graph of M responsible for certain behavior of interest. Naturally, the first step of circuit discovery is to define a task of interest to investigate. This requires specifying a dataset and a task-specifie metric to measure the performance of M and C. typically, by maximization of the task metric. Given an input x, similarly as to how the entire model defines function (x) from inputs to logits, we also associate each circuit C' with a function C(x)), defined by ablating away the effect of all components in M (i.e. the components not included in C). Circuit discovery may be performed on a single input example or, likelier, the result may be averaged over multiple input examples to reduce variance. Typically studied, can yield stable performance. Detailed information about which models are studied, in- put sample counts, and how the task metrics are defined is found in Appendix A Finally, note that our method finds the circuit of interest but does not automatically generate an interpretation of it, however, known interpretability techniques (Wang et al., 2023 can be used do so.
find averaging across 10- 100 samples, depending
the tasks being
3.3.2 METHOD
Here we present our algorithm for discovering circuits with CD-T. Although it is possible to perform CD at arbitrary granularity and on arbitrary components of transformer (e.g.. MLPs, query/key/value vectors), in this paper we focus solely on finding circuits consisting of attention heads in transformer to be comparable to prior work, and our unit of analysis is either the output of an attention head or the output of attention head at specific sequence position. For clarity, our method focuses on including or excluding specific nodes from a computational graph, rather than edges; CD-T at once implicitly models both direct and indirect effects of one node another, and our algorithm does not make a distinction between the two.
In each iteration, search a set of target nodes (initialized to the task objective). Once we find the set of nodes which are most relevant the task, we prune them by some simple heuristic (in our experiments, simply greedily removing the nodes with the least magnitude of impact was often sufficient, or sometimes pruning was not necessary), and designate these most important nodes the target nodes for the next iteration of the algorithm. We halt when the collection of nodes achieves adequate performance (e.g, comparable the whole network), or has not improved from the previous iteration.
potential source nodes find the ones with highest relevance to
We need to specifyy how "highest relevance to set of target nodes" is determined. As described above, we start with decomposition 3., % of the output of a source attention head 8, and we are interested in its relevance R(s, 7') to a set of target nodes T. Starting from we propagate the decomposition forward through the network using the CD-T equations for each module in the computational graph between and T; again, if there is some series of functions fn,... so that and (1 +7,) = 81 f1 +7) = 82 +22 ..., fnn
T= fn (n 3r + Tr, can straightforwardly apply the modules in order to obtain 81 Tr. In the case that T is the output of the network, Br will have matching dimensions, and it is usually appropriate to let R(s, T) be the task-specific objective as evaluated on 3r. (One intuitive instance of this is the case
=

---

Published as a conference paper at ICLR 2025
Algorithm Building circuit using CD-T for specified task
Input:
datapoint æ. precomputed mean activations μ
Denote Og (s(s) as the activation of on imput
Denote yu(s) as the precomputed mean activations of s on an appropriate distribution
Initialize C' to store the circuit model
Initialize 7 to be the output
repeat
for each attention head 8 upstream of
Run the model from to s, and compute 3g ag p(sss 7/- u(s)
Propagate decomposition to target nodes with CD-T, obtaining Pr: Use Br, Fr to compute the relevance of to T, R(s, T) using equation 16
end for
Heuristically pick the source nodes with the highest values of R(s, T); define this set S
C + CUS
(# Refinement: Greedily prune away the unnecessary nodes)
repeat
for each node in C::
C' +- C'
If Cf C()):C+C
end for until C has not changed since the last iteration
If C(x) - [x)] < return C {# Circuit performs close to the full model)
If (ar) not greater than last iteration, return C
T + S
until upstream attention heads to T are available
return C
where T is single "score" whose value we care about, so that Br is the contribution of to that score.) However, when T an arbitrary set nodes the network's internals, we must define proxy metric for the relevance of to T; in this case, we define the relevance of source node to set of target nodes to be the sum of the relevances to the target nodes T, and in this paper for R(s, t) we choose quick-to-compute measure of the size of B5:
|| Ielll
R(s,T)= {(stt)=
Σ
(16)
Careful readers may have noticed that this particular metric measures the magnitude of node's contribution to another, but not the "direction", so that nodes which maximize this metric may actu- ally contribute negatively to the task under investigation. We found in our experiments that this does not preclude us from finding functioning circuits. On the contrary, this allows us to straightforwardly find nodes which have significant negative contribution to task metric, such as the "Negative Name Mover Heads" in the IOI task. In principle, it is possible that different choices of target relevance metric can be developed to adjust the behavior of the circuit discovery algorithm to avoid (or ac- centuate) this behavior; we do not explore this in the paper but identify it possible avenue for further investigation. Our complete algorithm is described in Algorithm [ presented in the specific case where have chosen to decompose source nodes that 8g the activation 's deviation from the mean over some distribution. More details on the heuristics used and complexity analysis of the algorithm is found in Appendix B

---

Published as a conference paper at ICLR 2025
Coarse Nodel-level Runtime vs ROC AUC
CD-T
14
X
9.352
0.8
ROC AUC
0.6
0.903
P. P.Chee feitems
Methid
0.4
CD-
9.853-
EAP
ACDC of EL Div Hotte
9.823
Taak
Docttin in
Greater than
Berstring
101
x
(Runtime
Size (C') as fraction [M]
0.4
0.6
0.5
10
Figure 1: Left: Log algorithm runtime and ROC AUC comparison. Each dot represents an aver- age measurement faithfulness CD-T circuits compared to random circuits from the reference distribution of varying sizes (x-axis). Dotted vertical lines indicate the actual size of the eircuits. C denotes CD-T circuit. C denotes a random circuit. M denotes the full model.
task with specific methods differentiated by colors.; Right: The relative
EVALUATING CIRCUIT DISCOVERY ALGORITHMS
4.1 EXPERIMENTAL SETUP
We compare the performance of circuits recovered by CD-T with those obtained by two SOTA base-
lines: ACDC (Conmy et al. 2 3 and EAP (Syed et al. 2023 SAE-based methods O O'Neill & Buil 2024, Marks al. [2024] He al. 2024 and classic neural network pruning methods are not considered here because they either require model training or focus more compressing neural net- works for faster inference to reduce storage requirements. Furthermore, pruning for interpretability is also shown to be outperformed by ACDC (Conmy et al., (023)). We evaluate the circuit perfor- mance on three tasks: indirect object identification (IOI) [Wang et al. (023), greater-than (Greater- than) (Hanna et al. 2023), and docstring completion (Docstring) (Heimersheim & Janiak 2023 (see Appendix A for details). They are three standard evaluation tasks to benchmark circuit discovery methods prior work (Hanna et al.) 2023 Wang et al., [023) Heimersheim & Janiak 2023) has
manually found circuits in language models explaining for those tasks, which often serve as an im-
perfect reference standard for circuit discovery work (Conmy et al. [2033] [yed et al. [2023] O Neill
& Bui 2024) due to the lack of ground -tuth circuits. For both baselines, we perform node-level patching to be comparable since CD-T node-level patching technique. All experiments are conducted on an NVIDIA A100 GPU.
4.2 CD-T EXCELS IN IDENTIFYING THE CIRCUIT RESPONSIBLE FOR THE UNDERLYING ALGORITIIM EFFICIENTLY
4.2.1 RUNTIME AND THE RECOVERY OF MANUAL CIRCUITS
Leveraging the manually discovered circuits as reference standard, we measure how much over- lap there between the recovered circuits and the reference circuits using ROC AUC. Specifically, we measure the ROC AUC by sweeping through a range of thresholds that determine cutoff scores keep nodes in a circuit. For CD-T, we test by varying the percentile of top nodes to extract in every iteration, in the range of [90,99). At the same time, we also compute the average algorithm runtime across runs when sweeping through the thresholds for each method. Although CD-T sup- ports arbitrary granular representations of the circuit, the runtime is measured by building "coarse node-level" circuits by not splitting into sequence positions to be comparable with EAP and ACDC
We adopt the ACDC optimized for KI. diveregence instead of task-specific metrics hecause it's shown perform better in Conmy et al. (2023).
We consider EAP but not EAP-IG because EAP-IG performs comparably to EAP on the tasks evalu-
(Hanna al. 2024 but with runtime increase by constant factor in theory.
8

---

Published as a conference paper at ICLR 2025
because EAP doesn't support measuring node importance at sequence positions, and ACDC takes an impractical amount of time to achieve sequence-positeene splitting. To demonstrate CD-T's capa- bility to identify circuits responsible for underlying algorithm for tasks with great efficiency, we plot the algorithm runtime against the ROC AUC and report the result in Figure Lett. From Figure D Left we can CD-T outperforms the haselines by achieving high ROC AUC and low runtime at the same time. EAP obtains similar runtime as CD-T but performs slightly worse in recovering manual circuits for Greater-than and IOI. Even under this coarse nodel-level abstraction, ACDC is the least efficient as it often takes hours to identify circuit, yet with no better ROC AUC on all tasks compared to other methods. To explain the better ROC AUC performance of CD-T, we find that CD-T is especially good at identifying negative heads and supporting heads, such as the neg- ative name-mover heads and backup name-mover heads in IOI, which other algorithms struggle to do (more discussion in Appendix
4.2.2 PROBABILITY OF IDENTIFIED CIRCUITS TO BE MORE FAITHFUL THAN RANDOM CIRCUITS
Given a circuit and task, one can evaluate how well the circuit performs on the task by measuring the task task-specific metric, such as the logit difference on IOI. However, in mechanistic interpretability, instead of aiming to identify the best-performing circuits on the task-specific metric, an ideal circuit should be able to faithfully replicate the full model behavior. Following prior work (O' OO'Neil& & Bui] 2024 Marks et al.) 2024), faithfulness is defined as the proportion of the full model's performance that a circuit explains, relative to the baseline performance when no specific input information is provided. Given a task, faithfulness is computed as mi m(M))m((f)> where m(C), m00). and m(M) are the average of the task-specific metric performance over the dataset for the circuit, the corrupted model with all heads ablated, and the full model, respectively. To complement the imperfect nature of the manual circuits, leveraging faithfulness, here we provide another evidence to validate CD-T's mechanism perservation of the original model. We adopt a circuit hypothesis testing framework proposed in Shi al. (2024) to test whether circuit preserves the original model's performance by measuring the probability of a circuit C' to be more faithful to the original model M, compared to random circuits O" of the same size sampling from a reference distribution. This test verifies that the circuit is not simple lucky draw from the distribution of random circuits, and ensures that it is better than at least a fraction qa of random circuits. In our experiment, we sample random circuits from all possible nodes in the full model, and we repeat the process for 10 times. Figure | results of the best CD-T circuits (as measured by faithfulness and among different granularity) on the 3 tasks. Our best circuit for IOI is from attention heads at specific sequence positions, while for Greater-than and Docstring, they are from attention heads at a coarse level without the splitting. The results show CD-T circuits are 80% (q- 0.8) more faithful than random circuits of up around 60% of the original model's ' size for 101 and Greater-thar and even those of up to 80% of the original model's size for Docstring, with just a small percentage of nodes in the CD-T circuits (0.4% of full model size for 101, 5.5% for Greater-than than, and 31.3% for Docstring). This finding suggests the faithfulness of CD-T circuits is not due random chance, and that more fine-gre -ne-ggriieed abstraction of circuits sometimes is necessary to obtain better faithfulness with an extremely small and specific set of nodes, as in the 1OI case.
m(C))m(()
Right shows the
4.3 IDENTIFIED CIRCUITS MATCII FULL MODEL PERFORMANCE
4.3.1 FAITHFULNESS OF CIRCUITS OF VARYING SIZES
To understand how faithfulness changes as the size of the circuits grows, we compute node impor- tance before pruning for each method and track faithfulness performance as we add circuit nodes in order of importance, while ablating all other nodes with corrupted activations (see Appendix (A). We sample attention heads in the original model one by one in random order to construct the random baseline, and report the results in Figure 2 Since the goal of mechanistic interpretability is to repli- cate the original model's 's behavior, namely achieving faithfulness o with the smallest amount
$ ACDC is omitted in this experiment hecause of the prohibitive amount of runtime in hours to obtain individual node importance.
Faithfulness can go heyond 1 because there are both positively and negatively contributing heads in the model. The original model performance relies on a certain dynamic of interactions hetween them. When the

---

Published as a conference paper at ICLR 2025
Faltttiness (CI)
Falth/uiness (Groatortthan)
Fathhuhesss (Decsting)
Pathitlrees
0.5-
Falth Palthfulmmss
Falte Saleran
699-
64
CDT lean indrle ren MMdiee
12
CD mesh sk-bb CTT Desittate
(BI nocan anirte
Fanden
53
Natal
Manal
28
88
100
14
20
46
83
300
130
34
12 25 No. Atton, Hoads
30
No. Attinn Heads
No. Atton, Heads
Figure 2: Faithfulness of CD-T circuits, EAP circuits, and randomly selected circuits of equiva- lent size for IOI, freater-than and Docstring tasks. CD-T circuits obtains the full model's perfor- mance (faithfulness of 1) faster than EAP as attention heads are added in order of importance.
of nodes possible, this is naturally the first question we ask. First, we find CD-T circuits are able to perfectly replicate model behavior (faithfulness 1) using fewer nodes than EAP for all the three tasks. Second, manual circuits outperform both CD-T and EAP of same size, indicating room of improvement of advanced circuit discovery methods. Finally, we compare the effect of ablation methods on CD-T circuits' performance by including CD-T circuits obtained from zero-ablation in the plot, and find that mean-ablation tend to yield higher faithfulness compared to zero-ablation Greater-than and Docstring, except for IOI where they are quite comparable. Mean-ablation of- ten viewed preferred way of ablation because it preserves general concept of the task, hence is less destructive than zero-ablation. Similar finding is also discussed in other circuit discovery
work (Conmy al. 2023 Wang al. 2023
DISCUSSION AND LIMITATIONS
Despite having shown both qualitative and quantitative evidence of the strengths of CD-T auto- mated circuit discovery, our results is limited in scale and to the datasets and prior interpretation methods evaluated. More work is needed to generalize these findings to a broader set of models, datasets, and interpretation methods. Although comprised of any of a neural network's s internal components, we only discussed circuits built with purely attention heads. Automating the discovery of circuits with different/heterogeneuur compo- nents, particularly the MLPs of a transformer, is a promising direction for further investigation. Finally, our method is itself limited in two ways: first, as mentioned in 3.3 the metric which mea- sures the relevance of a source node to target node does not differentiate signs of the contribution, which might be useful in an interpretability context. Second, while the algorithm is capable of finding circuits under reasonable default settings, some amount of effort is still applied to tune the heuristic elements, such as the number of heads to select in each iteration.
methods can be applied to construct circuits
CONCLUSIONS
We have proposed contextual decomposition for transformers (CD-T), a set of mathematical decom- position principles efficiently compute the contribution of model features to other model features within a transformer. CD-T can produce circuits of arbitrary granularity, and is the first able to produce circuits fine-grained
attention heads specific sequence positions while at the same time reduce runtime from hours to seconds, when it is compared with prior work (Syed et al. 2023 Conmy et al. 2023). CD-T improves upon limitations of prior work, as it is automatic and does not require any training with all common transformer architectures. On three standard circuit evaluation datasets (indirect object identification, thea comparisons, and docstring completion), we demonstrate that CD-
craffedd examples by humans. CD-T is fundamentally compatible
dynamic is not perfectly replicated and there's extra positively contributing effect, the circuit might give higher prediction score than the original model would give, which results in faithfulness >
10

---

Published as a conference paper at ICLR 2025
T outperforms ACDC and EAP by better recovering the manual circuits with average of 97% ROC AUC under low runtimes. In addition, we provide evidence that faithfulness of CD-T circuits is not due to random chance by showing our circuits are 80% more faithful than random circuits of up to 60% the original model size. Finally, we show CD-T circuits are able to perfectly replicate origi- nal models' behavior (faithfulness that our future open-source implementation of CD-T benefits mechanistic interpretability research community by opening up a window from enabling more fine-grained analysis within practical runtime.
using fewer nodes than the baselines for all tasks. We hope
ACKNOWLEDGMENTS
GZ would like to thank Shawn Hu for valuable discussions that contributed greatly to mathematical formalization, experiment design, and important heuristic components of our method. We gratefully acknowledge partial support from NSF grant Research: Foundations of Data Science Institute (FODSI), the NSF and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-203 1833 and 81 $14639, NSF grant MC2378 to the Institute for Artificial CyberThreat Intelligence and Opera- tioN (ACTION), and Berkeley Deep Drive (BDD) grant from BAIR and Dean's fund from CoE, both at UC Berkeley.
2-122413265 NSF grant 2023505 on Collaborative
REFERENCES
Marco Bertolini, Linlin Zhao, Ajork-Amé Clevert, and Floriane Montanari. Beyond atoms and bonds: contextual explainability via molecular graphical depictions. 2022.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Berbert-Voss,. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Ma teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam Mc- Candlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. ArXiv, abs/2005. 2020a. URL https: /api. www.sticsschoollateeoeoeoemmaaee org/
CorpusID:21897:1
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020b.
Arthur Conmy, Augustine N. Mavor-Parker, Aengus Lynch. Stefan Heimersheim, and Adrià Garriga-Alonso. Towards automated circuit discovery for mechanistic interpretability. ArXiv, abs/2304,1499 14997, 2023. URI. https: semant escholar wwascholar..rg/ooopp.....
258418244]
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In North American Chapter the Associ- ation Computational Linguistics, 2019. URL https:/ api. semanticscholee org/
Corpus ID: 52967399
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https:/transfommr-ccrriiii wttp:issssuto00///aaaooooooo
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac HaffieldDDodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse,

---

Published as a conference paper at ICLR 2025
Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah.
Toy models of superposition. Transformer Circuits Thread, 2022. https//transffrmmr
cireuits.uu//22222oy model/index.html.
Amir Feder, Nadav Oved, Uri Shalit, and Roi Reichart. Causalm: Causal model explanation through
counterfactual language models. Computational Linguistics, 47:333-386, 2020. URL https:
7/api. sommantcsscool.. www.ar..rg/cCooe Sarg//CrpusID:: 218901061
Jiahai Feng and Jacob Steinhardt. How do language models bind entities in context? ArXiv, 17.113, 2023. URL http:://pp...mmmnccccclllll www.nsscccoooho..oooooo..... 264490454
Atticus Geiger, Hanson Lu, Thomas F. Ieard, and Christopher Potts. Causal abstractions of neural networks. In Neural Information Processing Systems, 2021. URL httpps://ap)
semanticschol.e. orgg/oorpssi::233555
Mor Geva, R. Schuster, Jonathan Berant, and Omer Levy, Transformer feod-ooorrr layers are key- value memories. ArXiv, abs/2012. 149913, 2020. URL https://api. htt.:appp..emmaticccchoooooooooo
brc/CozpusID: :2: 229923720
Nicholas W. Goldowsky-Dill, Chris Macleod, Lucas Jun Koba Sato, and Aryaman Arora. Lo-
calizing model behavior with path patching. ArXiv, abs/2304.05969, 2023. URL https:
org. IC:
Michael Hanna, Ollie Liu, and Alexandre Variengien. How does gpt-2 compute treater-than?:
Interpreting mathematical abilities in a pre-trained language model, 2023. URL https:
//arxxiv.rrr//bb//235555 20050
Michael Hanna, Sandro Pezzelle, and Yonatan Belinkov. Have faith in faithfulness: Going beyond circuit overlap when finding model mechanisms, 2024. URL https:://rxiv..rggg
2403. 17806
Kelei He, Chen Gan, Zhuoyuan Li, Islem Rekik, Zihao Yin, Wen Ji, Yang Gao, Qian Wang. Junfeng Zhang, and Dinggang Shen. Transformers in medical image analysis. Intelligent Medicine, 3(1): 59-78, 2023. ISSN 2667-1026. doi: https:/doi..or///
0000... URL https:
7/ //www.
rect com
Zhengfu He, Xuyang Ge, Qiong Tang. Tianxiang Sun, Qinyuan Cheng, and Xipeng Qiu. Dictionary learning improves fatch-free circuit discovery mechanistic interpretability: A case study on othello-gpt, 2024. URL/httss: /ttttps://axxx..oottt abs 2402 12201
Stefan Heimersheim and Jett Janiak. A circuit for python docstrings in a 4-layer attention-only transformer. URL: https://www. alignmentfrrum wrsposssseEKkkktkrxyaaaaaaxdd 2023.
Jaap Jumelet, Willem H. Zuidema, and Dicuwke Hupkes. Analysing neural language mod- els: Contextual decomposition reveals default reasoning in number and gender assignment. In Conference Computational Natural Language Learning, 2019. URL httpps://api, semanticscholen orq Ctrgg//rrussDDp. 20:
John M. Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridg- land, Clemens Meyer, Simon A A Kohl, Andy Ballard, Andrew Cowie, Bernardino Romera- Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David A. Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with alphafold. Nature, 596:583 589, 2021. URL httpps:/ stttpp::/api..emmtiiioo semant wrw.mannnnccchola..
Crg/CorpusID:: 22L.L:::3359575
Jean Kaddour, Aengus Lynch, Qi Liu, Matt J. Kusner, and Ricardo Silva. Causal machine leamn ing: A survey and open problems. ArXiv, 154775 2022. URL https: : //api.
semanti cschol ar. wwt.:org//oocssa 2501 44475
12

---

Published as a conference paper at ICLR 2025
David Lindner, J'anos Kram't ar, Matthew Rahtz, Tom McGrath, and Vladimir Mikulik. Tracr: Compiled transformers as a laboratory for interpretability. ArXiv, 0as/2331..50002 2023. URL
htttps:/appl wwwantisccohholar.....ooo www.:aa.rgroooopuoooeeeeo 25554499999
Samuel Marks, Can Rager, Eric J. Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. Sparse feature circuits: Discovering and editing interpretable causal graphs in language mod-
cls. ArXiv, ab4/2403,19647, 2024. URL https:/ /ttp//app... semant
semattt.. wcscholarc.cccocccccoococ org/
CorpusSD:2287327732
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. associations in gpt, 2023,
Locating and editing factual
W. James Murdoch, Peter J. Liu, and Bin Yu. Beyond word importance: Contextual decomposition to extract interactions from LSTMs. In International Conference on Learning Representations, 2018. URL
net/ wttt:/forrmmidddkkk....o htt.pp/eerr.....ceeoooooooooooor OZ
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures
for grokking via mechanistic interpretability. ArXiv, 0tt///311.021111 2023. URL https://
api sew.nticccccc wwr.chhlaarrrgo.... wtt.aoorr//orr....25553.....
Christopher Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and
Shan Carter. Zoom in: An introduction circuits. 2020. URL https://api.
semant: 1 cscho l ar.
Charles O' O'Neill and Thang Bui. Sparse autoencoders enable scalable and reliable circuit identifica- tion in language models, 2024. URL
http:/g/ab//44000 12522
Judea Pearl. Causality: Models, Reasoning and Inference. Cambridge University Press, USA, 2nd edition, 2009. ISBN 052189560X.
Tilman Räuker, Anson Ho, Stephen Casper, and Dylan Hadfield-Menell, Toward transparent ai: A survey interpreting the inner structures of deep neural networks, 2023.
Claudia Shi, Nicolas Beltran-Velez, Achille Nazaret, Carolina Zheng, Adrià Garriga-Alonso, An- drew Jesson, Maggie Makar, and David Blei. Hypothesis testing the circuit hypothesis in LLMs. In ICML 2024 Workshop on Mechanistic Interpretabilitty 2024. URL https:// openreview net/ tttp:///oouummnddd-soooooooo-o
Chandan Singh, W. James Murdoch, and Bin Yu. Hierarchical interpretations for neural network predictions. ArXiv, Abs/1806.05337, 2018. URL http://aapp...eannccccccooo..a org/ Corpus 19214673
Chandan Singh. W. James Murdoch, and Bin Yu. Hierarchical interpretations for neural network predictions. In International Conference on Learning Representations, 2019. URL https:
Alessandro Stolfo, Yonatan Belinkov, and Mrinmaya Sachan. mechanistic interpretation arith- metic reasoning in language models using causal mediation analysis. In Conference Empirical Methods in Natural Language Processing, 2023. URL http::::s//api..eeaaattcccoooaao
ort:/CorpusDD: 258865177
Aaquib Syed, Can Rager, and Arthur Conmy, Attribution patching outperforms automated circuit discovery, 2023.
10348
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Autention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998 0008,6008,000 2017. URL //t /ttpa/////oooaoor.eoeoo.e//////a
neuriss httpp//pp...paapee//////a//aee
13

---

Published as a conference paper at ICLR 2025
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart M. Shieber. Causal mediation analysis for interpreting neural nlp: The case of gen- der bias. ArXiv, abs/2004.12265 2020a. URL stw.p//cssseaattttcccooooo......
Corpus : 21 6553696
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart M. Shieber. Investigating gender bias in language models using causal media- tion analysis. In Neural Information Processing Systems, 2020b. URL https://api. wschoo ort/CorpusID:
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: circuit for indirect object identification in GPT-2 small. In
The Eleventh International Conference on Learning Representations, 2023. URL https:
Zhengxuan Wu, Atticus Geiger, Josh Rozner, Elisa Kreiss, Hanson Lu, Thomas F. Icard, Christopher Potts, and Noah D. Goodman. Causal distillation for language models. ArXiv, abs/2112.025005. 2021. URL bottps:/aapp wtt.aeettttcccooooo...oooo/ooooo o.
A EXPERIMENT DETAILS
A.1 TASK DESCRIPTION
Here we describe each task and how the experiments are performed in more detail; we also specify the exact model components which comprise the circuits we analyzed in the main paper. In this section, we use the convention of specifying attention heads in model by the tuple (layer index, head index). For the IOI task, we add a third entry for output of the head at a specific sequence position, and follow their nomenclature for semantically labeling the sequence positions.
A.1.1 INDIRECT OBJECT IDENTIFICATION (IOI) (WANG ET AL. 2023)
The indirect object identification (IOI) task is to predict the indirect object in a sentence with two entities, such as identifying "Mary" in the sequence "When Mary and John went for a walk, John gave an apple to ...". The objective originally defined for this task, which we also use, is the differ- ence between the predicted logit of the indirect object (IO) token and the subject (S) token. Wang et al. (2023) carefully design their experiment and dataset code so that mean ablation occurs using the mean activations over the corrupted "ABC dataset", which replaces the three name tokens in the task with random names. Likewise, when setting the decomposition of node, we set the "relevant" component to the deviation from the mean activations on the ABC dataset. This task was performed on GPT2-small, GPT-2 correctly performs this task about 99 percent of the time, so we did not take special measures to account for those samples where it doesn't in our circuit analysis. We identify circuits using 25 IOI samples drawn from mixed templates, and mean ablation is conducted using the corrupted ABC dataset. Another set of 100 IOI samples are used in evaluation. Manual circuit of IOI compared against in the ROC AUC experiment is: [(2, 2), (4, 11), (0. 1), (3, 0), (0. 10), (5, 5), (6. 9). (5, 8), (5, 9), (7, 3), (7, 9), (8, 6). (8, 10), (10, 7), (11, 0), (9, 9), (9, 6), (10, 0), (9, 0), (9, 7), (10, 1), (10, 2). (10, 6). (10, 10). (11, 2). (11, 9)], which is from Figure 2 in Waan et al. (0023).
Details of Circuit Analysis For this task, in order to provide some intuitions about what CD-T calculates, we provide number of heatmaps of the relevance scores at specific positions during various iterations of the circuit analysis. In this section, we don't follow our automated circuit discovery algorithm exactly, but instead partially follow Wang et al. (2023)'s 's analysis decide what sequence positions to search over and visualize.
The first iteration of the algorithm finds the Name Mover Heads: (9, 9, end), (10, 0. end), and (9, 6, end); the Negative Name Mover Heads: (10, 7, end). (11, 10, end); and some Backup Name Mover Heads: (10, 2, end), (10, 6, end), (10, 10, end), described by Wang et al. (2023)
Following Wang et al. (203)) ' analysis further, and deviating from the normal course of our circuit- finding algorithm, we compute the relevance of nodes to the Name Mover Heads just the end
14

---

Published as a conference paper at ICLR 2025
Relevance nodes to IOl task metric at last sequence position
0.4
0.3
0.2
0.1
Layer
0.0
-0.1
-0.2
-0.3
10
-0.4
0
1 2
3
4
5
6
7
8
9
10
11
Head
position, and find what they named the "S-Inhibition Heads", at (8, 10, end), (7, 9, end), and (7, 3, end), though there are some other relevant-looking contenders:
Relevance of nodes to output residuals of Name Mover Heads at last sequence position, normalized by mean rel per layer
0.200
0.175
0.125
Layer
5
0.100
0.075
0.050
0.025
Head
Continuing to follow their analysis, Wang et al. (2023) 's analysis further, we compute the relevance
of nodes to the S-Inhibition Heads the S2 position, and find our first minor disagreement with their process: though we find two of what they named the Induction Heads (5, 5, S2), (5, 8, S2), (5, 9, S2) we don't 't find the at (6, 9, S2) but instead find one at (5, 10, S2), and find a head their later analysis called a Duplicate Token Head, at (3, 0, S2).
Next compute the relevance of nodes to the Induction Heads at the S2 position, and find that (3, 0, S2) mostly drowns out the signal of the other Duplicate Token Heads:
Finally, compute the relevance of nodes to the Induction Heads the S1+1 position, and find (4, 11, S1+1), with the other Previous Token Head they found at (2, 2, S1+1) top contender, though there other heads not accounted for in their analysis:
Overall, there is significant but not perfect agreement with the results and analysis the IOI paper. We also attempted analysis by computing relevance intermediate matrices (i.e., the key, query,
15

---

Published as a conference paper at ICLR 2025
Relevance of nodes to output residuals of S-inhibition heads at S2 sequence position, normalized
0.8
Layer
0.3
0.2
10
0.1
0.0
3
4
5
6
7
8
10
11
Head
(Duplicate Token Heads) Relevance of nodes to output residuals Induction heads at S2 position, normalized for mean relevance among heads in layer
0.35
0.30
0.25
0.15
0.10
11 10
11
Head
value vectors) in the attention calculation, but found the plots to be qualitatively similar, possibly due the fact that CD-T by design propagates relevances from these vectors to the attention head outputs as well.
Another fact of note is that the scales on these plots vary substantially. It would be desirable find an interpretable normalization method to put these the same scale with some intrinsic meaning, or otherwise explain the causes of this phenomenon.
A.1.2 GREATER- THAN (HANNA AL. 2023)
The Greater than task to predict the last two digits in an incomplete sentence following the tem- plate "The [noun] lasted from the year XXYY to the year XX". And expect the model to assign higher probability to years greater than YY. The objective originally defined for this task, which we also use, is the sum of probabilities assigned to tokens corresponding sum of probabilities assigned to tokens corresponding to lesser years. (Some probability assigned tokens which don't correspond to numbers all.) For "mean-ablation" simply take the
greater years, minus the
16

---

Published as a conference paper at ICLR 2025
(previous token heads) relevance of nodes to output of Induction heads at S1+1 position, normalized for mean relevance among heads layer
0.18
0.16
0.14
Layer
0.12
0.10
0.06
0.04
1
2
3
4
5
6
7
8
10
11
Head
mean over the activations over 100 negative datapoints (impossible completions, with the ending year preceding the starting century), and as above, when setting the decomposition at a source node, define the relevant component to be the deviation from the mean activation over this distribution. This task was performed on GPT2-small. GPT-2 correctly performs this task about 99 percent the time, we did not take special measures to account for those samples where it doesn't 't in circuit analysis. We identify circuits using a random sample of 100 datapoints provided by [Hanna et al.) 2023), and mean ablation is conducted using the negative impossible completion samples. Another set of 100 samples are used in evaluation.
It should be noted that result and the results in Hanna et al. (2033) are not directly comparable, since they also attempt to investigate the influence of MLPs and their resulting circuit removes most of the MLPs in GPT-2. For the comparisons found in the main paper, we have compared our result (with all MLPs) to the circuit which is obtained by taking all the attention heads named in Hanna al. (2023) (also with all MLPs), which is circuit distinct from the one found in their work.
Manual circuit of theater-than compared against in the ROC AUC experiment is: [(5, 1), (5, 5), (6, 1), (6, 9), (7, 10), (8. 8). (8, 11), (9, 1)].
Independently of the comparison, it remains true that keeping the relatively small proportion of attention heads in our circuit results in recovering almost all of GPT-2's capability on this task; see below.
Task specifse metric
Correct guess rate
Full model
0.817
0.992
All attention heads ablated
-2.095
0
Their circuit
0.768
0.989
Their circuit (most MLPs ablated)
0.727
N/A
Our circuit
0.761
0.981
A.1.3 DOCSTRING (HEIMERSHEEM JANIAK 2023)
The goal of the Docstring task is to predict the next variable name in Python docstring. For example, given function with variable names LOAD, SIZE, FILES, and LAST, the task is to predict the word after one :PARAM. According to docstring conventions, this should be the variable name in the function definition which follows the most recent variable name appear after :PARAM. The objective originally defined for this task is the logit assigned to the correct variable name, minus the max logit assigned to all other variable name tokens found in the function signature. For
17

---

Published as a conference paper at ICLR 2025
their raadom_randdo dataset which randomize both the variable names in the function definition and in the docstring of prompts, and correspondingly, when setting the decomposition at source node, define the relevant component to be the deviation from the mean activation over this distribution. We identify circuits using a 100 datapoints sampling for the dataset provided by |eemmrrshim & Janiak (2023), and mean ablation is conducted using the corrupted aantoommraddon dataset. Another set of I00 samples are used in evaluation. This task was performed on a 4-layer attention-only transformer trained on natural language and Python code (attn-onlyy 4-4 released with the TransformerLens library for the express purpose of facilitating mechanistic interpretability research. Another complication is that the toy model only guesses the correct token between 60 and 65 percent of the time. To account for this, we perform our circuit analysis and evaluation on the subset of input examples for which the model performs the task correctly.
"mean-ablation"
Our circuit consists of the nodes (3, 0), (3, 6) (1, 4), (0, 5). 0). (1, 0) (1, 4), (0. (2,
Heimersheim & Janiak (2023) find the circuit (0, 2), (0, 4), (0, 5), (1, 2). (1, 4), (2, 0), (3, (3. with their initial analysis, and heuristically observe that three heads help to obtain the "augmented circuit" (0, 2), (0, 4), (0, 5), (1, 2), (1, 4), (2, 0), (3, 0), (3, 6), (1, 0), (0, 1), (2.
Manual circuit Docstring compared against in the ROC AUC experiment is: [(0, 2). (0, 4), (0, 5), (1, 2), (L, 4), (2, 0), (3, 0), (3. 6)].
Task sees--peccfiic metric
Correct guess rate
Full model
3.661
1.0
All attention heads ablated -2.095
0
Their circuit Augmented circuit
2.884
0.54
3.524
0.66
Our circuit
3,235
0.57
B ALGORITHM DETAILS
B. HEURISTICS
In addition to the greedy pruning presented in Algorithm 1 as a refinement step, here we describe other heuristic elements to the algorithm.
The threshold for determining which set S of highest-contributing nodes may be varied: is possible to pick the top N nodes or fraction of nodes, or to automatically detect outliers. Empirically, we find that the distribution of node contributions varies quite severely, so finding a heuristic which works in all cases is actually object of future work. In this paper, we set the threshold by varying the percentile of top nodes to extract, in the range of [90, 99] to obtain the ROC AUC in section 4.2.1
Empirically we find, for fixed target node, the relevance scores of source nodes in different layers to this target node may have different expected magnitudes, due to the numerical effects of propagation through the network. To account for this, we normalize the relevance scores by dividing by the average magnitude across scores found in given layer.
To avoid the risk of propagation equations becoming numerically unstable after a large number of iterations, especially the relevant and irrelevant constituents differ in sign at a specific index, set the value of one of rel/irrel at this position to 0 and the other to the sum of the two terms.
B.2 COMPLEXITY ANALYSIS
To provide more clarity, the computational complexity of the algorithm satisfy the following prop- erties:
Each decomposition (of a set target nodes with respect to a set source nodes) requires cost in FLOPs similar one forward pass of the model (and often less, since values prior
18

---

Published as a conference paper at ICLR 2025
the source nodes be cached and values after the target nodes do not need to be calculated).
The core algorithm a loop, where each iteration, we search over all nodes which can potentially have high relevance to the target nodes. (This means that heuristically excluding some nodes from the search can potentially significantly decrease cost FLOPs, and cost in FLOPs increases linearly with respect to the granularity with which we separate the nodes in the model.)
The memory footprint of a single decomposition, including the forward pass, is small (less than 3) constant multiple of the cost of forward pass; the only added costs are to keep track of the relevant and irrelevant constituent tensors separately, as well as bookkeeping of components of the target decomposition metric.
The cost (in FLOPs or memory) of performing the analysis on set of input examples is linear in the number of input examples, since the same set of computations needs to be done with respect to each example.
19


Reviews:
[{"review_id": "QbHYvDlNFw", "reviewer": "Reviewer_ZeEK", "summary": "This paper introduces a novel method CD-T that leverages contextual decomposition for mechanistic interpretability in transformers. CD-T's granularity is fine-grained up to attention heads (does not include MLP layers). They report the results of their method (classification metrics and faithfulness of the circuit as compared to randomly sampled circuits) for three tasks, namely indirect object identification, greater-than comparisons, and docstring completion.", "questions": "Line 20: \"CD-T is the first ---?\" there seems to be a missing word here", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Motivation, relevant scientific terms and prior works are well-written, making the paper accessible to researchers who are non-experts in this area. \n\n2. The proposed method can be easily used for all transformer architectures (encoder-decoder and decoder-only, as per line 72). Further, the proposed method significantly reduces the runtime required for inference.", "weaknesses": "1. Manual circuits are not fully explained (definition, how they are computed, cost of computation, etc.) - manual circuits are used as reference during evaluation (line 405), and it is necessary to provide these details regarding them. \n\n2. It is unclear how CD-T works for inference - is there a distinct circuit discovered for each inference datapoint, or is the circuit found on a broader scale, so as to tradeoff size vs. performance for a larger test set?\n\n3. The authors are encouraged to discuss the practical utility of CD-T. Once their method yields a circuit, could additional analyses, such as on attention heads, be conducted? For instance, are there heads that consistently produce positive or negative outputs across all data points, or do the importance of attention heads vary based on input data characteristics?\n\n4. What transformer architecture is used for the experimental results?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gFI0fjtwce", "reviewer": "Reviewer_tnfo", "summary": "The paper introduces Contextual Decomposition for Transformers (CD-T) for building interpretable circuits in llms for better mechanistic interpretability in transformers efficiently. In contrast to other methods, CD-T employs a mathematical decomposition that isolates the contributions of specific features, allowing CD-T to discover circuits at various levels of granularity. Results show significant improvements in runtime and interpretability over existing methods such as ACDC and EAP on several tasks.", "questions": "1. See W1; In addition, what’s the computational complexity of the algorithm in terms of the input parameter?\n2. How does pruning affect the mechanistic interpretability of the model?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe paper is relatively well presented. It provides a clear description of the method, experiment details, and contribution. \n2.\tUsing the decomposition approach for designing/scaling LLM models is relatively novel and effective. Results show promising improvements in computational efficiency while maintaining circuit quality, and the method seems to be agnostic to the specific types of transformer models.", "weaknesses": "1.\tThe description of the algorithm can be improved. In particular, the algorithm description is relatively informal and could benefit from more details and formalism. For “Prune nodes from S for which doing so increases the task metric”, how is increasing the task metric evaluated?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pbVSUYX7CO", "reviewer": "Reviewer_Lhid", "summary": "The paper proposes an automatic, mechanistic interpretability approach for finding circuits in transformer models. Their approach propagates backwards from the model output, finding a circuit of nodes most relevant to the given output, and from there the nodes most relevant to those nodes, and so on. These circuits are task-specific subgraphs in the larger model compute graph that are most stimulated (and most relevant) to the input type or task under consideration. \n\nThe authors consider circuit discovery for the attention heads of transformer models, choosing to leave the MLP layers in an attention block for later work. Their linear decomposition model is mechanistic in that it solely depends on the weights and biases of each layer of the network and the relevance metric calculated for the previous layer. It does appear to be data-dependent in that it still relies on a dataset for calculating these relevance metrics. It also relies on precomputed mean activations to do so.", "questions": "1) I would suggest clearly presenting model architectures under evaluation and assumptions being made on input data for the approach - the mechanistic relevance metrics are ultimately calculated from the outputs derived from input data.\n2) How would this scale to large training datasets in LLMs? Is there a minimum dataset required? Can the authors comment on this?\n3) What is the effect of autoregressive models and sequence to sequence decoding on this approach? Is there a difference? I don't see a clear analysis in Section 5, and while that is not necessary it would be nice to have the authors comment on this in Section 6.\nQuestions addressed in comments.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1) The paper's originality is in applying a novel formulation of an existing method - Context Decomposition - to transformer models, albeit to the attention heads of the transformer models. While novel, I am not in the space of circuit discovery in neural networks and unable to judge how novel this is. \n2) The presentation is clear and the limitations in data dependence and precomputation are made clear to the reader in algorithm pseudocode, which is also concise and clear. Experimentation is presented to clearly establish the work's superior performance over the baselines. Notation remains consistent through the paper to my eyes.\n3) Significance: The work presents the an automatic circuit discovery approach using CD (extending it as CD-T) in transformer models, and clearly outperforms its chosen baselines. The approach is moreover mechanistic and iterative rather than probabilistic and therefore also comes with a guarantee of circuit discovery given certain assumptions on available data and time.", "weaknesses": "1) I am unsure of the significance of the work due to the iterative nature of the process and the need for data to adequately stimulate different 'tasks' (or behavior modes, perhaps?) and therefore discover circuits in the model. This is to my eyes not clearly explained. [Addressed]\n2) Prior work is noted to have conducted circuit discovery on large language models, and I can't easily find the models used for circuit discovery here. That lack of clarity is concerning, especially when using the manually discovered circuits as a target and prior art as baselines. [Addressed]\n3) I am also uncertain of scalability of this approach, considering that the network architectures under evaluation are not clear and that it has been restricted to attention heads rather than MLP layers. I am also uncertain of the effects of sequence length, autoregressive or non-autoregressive decoders. Section 5 does not make these things clear. [Addressed in part; prior art remains deficient as well here - future work?]", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)