Efficient Automated Circuit Discovery in Transformers using Contextual Decomposition

Previous efforts to interpret neural networks have frequently drawn on ideas from causal inference (Pearl, 2009). Much of this work emphasizes counterfactual reasoning or enforcing causal constraints on computations to explain model outputs (Pearl, 2009; Feder et al., 2020; Geiger et al., 2021; Wu et al., 2021; Kaddour et al., 2022). A related line of research (Vig et al., 2020a; Stolfo et al., 2023; Feng & Steinhardt, 2023) takes a different approach by treating internal components (or nodes, if the model is viewed as a computational graph) as mediators and conducting causal mediation analysis, also called activation patching. In these studies, the contribution of nodes is 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-Dill et al. (2023) and Wang et al. (2023) propose treating input-to-output paths as more expressive mediators. They introduce the technique of path patching to explore how different edge subsets in a model’s computational graph affect its behavior.

Ablation-based methods are essential for identifying key components within models. Through activation 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 of manual 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 et 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 MLPs based on model performance on a 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. (2023) introduced Edge Attribution Patching (EAP, or ACDC++ as later renamed), which improves ACDC’s efficiency by using a linear approximation of activation patching. However, this reliance on linear approximations can lead to overestimation of edge importance and a weaker correlation with true causal effects. EAP also struggles when the gradient of the performance metric is zero. To resolve these issues, Hanna et al. (2024) proposed Edge Attribution Patching with Integrated Gradients (EAP-IG) 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 SAE-learned features to map attention head contributions to identified circuits. However, SAE-based methods either require carefully crafted human-generated examples (O’Neill & Bui, 2024) or require an impractical amount of compute to train an SAE (He et al., 2024; 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 support further attention head splitting at specific sequence positions (Syed et al., 2023; Hanna et al., 2024; O’Neill & Bui, 2024), or it is fundamentally feasible but requires a 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 mathematical 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).

Contextual Decomposition (CD), introduced by Murdoch et al. (2018), attributes local importance to input features in LSTMs by analytically decomposing the output without any changes to the underlying 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 a selected token or phrase (denoted as the relevant) and those that are not (denoted as the irrelevant), using a factorization of the update equations for hidden states and cell states. Jumelet et al. (2019) propose Generalized Contextual Decomposition (GCD) for unidirectional LSTMs, applying it to phenomena like number agreement and pronoun resolution. Singh et al. (2018) extend CD from LSTMs to RNNs and 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 as the feature importance measure. We further extend CD to transformers (CD-T) and we argue that our development of CD-T is a significant advance with two novel contributions. First, we carefully design mathematical decomposition principles tailored for transformers of all types (e.g., decoder-based, encoder-based), which makes it able to 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.

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 to any other position. We start by describing the operation of a general attention mechanism before discussing the modifications needed to specialize it for generation, exemplified by models such as GPT (Brown et al., 2020a), and for sequence-to-sequence transformation models such as BERT (Devlin et al., 2019). Formally, the attention mechanism computes a vector of contributions from a series of vectors $\{x_{i}\}_{i=1}^{l}$ to a target vector $x$ at position $p$ with $x_{i},x\in\mathbb{R}^{d}$. An attention layer typically consists of a small number of heads, with each head parameterized by three functions, the key function $f_{k}$, value function $f_{v}$, and query function $f_{q}$. The output of the attention head to $x$ is computed as follows where $d_{k}$ denotes the output dimension of $f_{k}$:

The functions $f_{k},f_{v},f_{q}$ are commonly simple learnable transformations such as linear transformations 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 an attention layer is a 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 $(t+1)^{th}$ token generated based on the previous $t$ 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 $t+1$ and beyond are not determined. As a consequence, only the representations for tokens before $t$ are used to compute the attention vectors.

We will now describe the general method of contextual decomposition (CD), as well as our extension to the transformer architecture and applications to mechanistic interpretability.

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 a sum of two parts: a $\beta$ part (said to be relevant), which contains the part of this particular state that stems from the input tokens of interest, and a $\gamma$ part (said to be irrelevant), which contains information coming from tokens outside of the list of interest. $\beta$ is often initialized with masked input embeddings ($1$’s specify where token positions of interest are, and $0$’s the opposite), and $\gamma$ the complement of $\beta$.

Given a decomposition of a vector $x\in\mathbb{R}^{d}$ into relevant and irrelevant constituents $\beta+\gamma=x$ (here we use the word constituents instead of “components” to avoid name collision with the “components of a neural network”), and a module $f:\mathbb{R}^{d}\to\mathbb{R}^{k}$, CD consists of a set of mathematical equations to determine the decomposition of the output of a module $f$ when given $x$ as input: $f(x)\in\mathbb{R}^{k}=\beta_{o}+\gamma_{o}$. The general principle is to find a symbolic expression for $f(x)$ in terms of $\beta$ and $\gamma$, and group the terms in the expression according to whether they are solely a function of $\beta$ or not: the sum of terms solely relying on $\beta$ becomes $\beta_{o}$, and the sum of the remaining terms becomes $\gamma_{o}$. We refer interested readers to Murdoch et 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 $f(\beta_{1},\gamma_{1})=\beta_{2},\gamma_{2}$, and $g(\beta_{2},\gamma_{2})=\beta_{3},\gamma_{3}$, then $g(f(\beta_{1},\gamma_{1}))=\beta_{3},\gamma_{3}$) 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 an earlier layer $i-1$) can affect (the relevance term in a later layer $i$) that are unaccounted for by this rule. For example, Singh 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 $\beta$ and $\gamma$:

We succinctly represent the above equations as $\beta_{i},\gamma_{i}=LinearDecomp(\beta_{i-1},\gamma_{i-1})$, where $W$ and $b$ 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 of transformers is the self-attention module. We handle the decomposition of the attention equations as follows:

where above $Mask$ represents a causal mask in appropriate architectures and $Softmax$ is the softmax 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 a 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 1 on a subsequence of the input and 0 elsewhere, but in a mechanistic interpretability context, the CD framework allows us to decompose activations into any choice of $\beta,\gamma$. Thinking of the network 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.

We compare the performance of circuits recovered by CD-T with those obtained by two SOTA baselines: ACDC (Conmy et al., 2023) ³³3We adopt the ACDC optimized for KL diveregence instead of task-specific metrics because it’s shown to perform better in Conmy et al. (2023). and EAP (Syed et al., 2023) ⁴⁴4We consider EAP but not EAP-IG because EAP-IG performs comparably to EAP on the tasks we evaluate (Hanna et al., 2024), but with a runtime increase by a constant factor in theory.. SAE-based methods (O’Neill & Bui, 2024; Marks et al., 2024; He et al., 2024) and classic neural network pruning methods are not considered here because they either require model training or focus more on compressing neural networks for faster inference to reduce storage requirements. Furthermore, pruning for interpretability is also shown to be outperformed by ACDC (Conmy et al., 2023). We evaluate the circuit performance on three tasks: indirect object identification (IOI) (Wang et al., 2023), 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 as prior work (Hanna et al., 2023; Wang et al., 2023; Heimersheim & Janiak, 2023) has manually found circuits in language models explaining for those tasks, which often serve as an imperfect reference standard for circuit discovery work (Conmy et al., 2023; Syed et al., 2023; O’Neill & Bui, 2024) due to the lack of ground-truth circuits. For both baselines, we perform node-level patching to be comparable since CD-T is a node-level patching technique. All experiments are conducted on an NVIDIA A100 GPU.