Sparse Attention Decomposition
Applied to Circuit Tracing

The approach used in this paper to trace circuits starts by decomposing the attention score $A^{\prime}$ in terms of the SVD of $\Omega.$ The matrix $\Omega$ has size $(d+1)\times(d+1)$, but due to its construction it has maximum rank $r$. We therefore work with the SVD of $\Omega=U\Sigma V^{\top}$ in which $U\in\mathbb{R}^{(d+1)\times r}$, $V\in\mathbb{R}^{(d+1)\times r}$ and $\Sigma\in\mathbb{R}^{r\times r}$. $U$ and $V$ are orthonormal matrices with $U^{\top}U=I$ and $V^{\top}V=I$, and $\Sigma=\operatorname{diag}(\sigma_{0},\sigma_{1},\dots,\sigma_{r-1})$ with $\sigma_{0}\geq\sigma_{1}\geq\dots\geq\sigma_{r-1}\geq 0$. Important to our work is that the SVD of $\Omega$ can equivalently be written as

in which $\{\mathbf{u}_{k}\}$ and $\{\mathbf{v}_{k}\}$ are orthonormal sets and each term in the sum is a rank-1 matrix having Frobenius norm $\sigma_{k}$. We refer to each term $D_{k}$ as an orthogonal slice of $\Omega$, since we have $D_{k}^{\top}D_{j}=D_{k}D_{j}^{\top}=0$ whenever $k\neq j$.

When the sparse decomposition hypothesis holds, we can approximate the score computed by the attention head as follows:

where the number of terms in the sum (i.e., $|S_{ij}|$) is small.

We use $S_{ij}$ to denote the subset of values of $k$ for the token pair $(i,j)$. Besides being specific to an attention head and token pair, $S$ also depends on the task. In this paper we do not define ‘task’ precisely; in what follows, we study only a limited set of attention head functions that are performed when generating outputs for IOI prompts in GPT-2 small. We leave the association between $S$ and precisely-defined tasks as a fascinating direction for future work.

Our tracing strategy constructs $A^{\prime}$ using (6), i.e., using $S_{ij}$ in place of all of the singular vectors of $\Omega$. This is akin to denoising in signal processing. When a signal has an approximately-sparse representation in a particular orthonormal basis (e.g., the Fourier basis or a wavelet basis) then removing the signal components that correspond to small coefficients is useful to suppress noise. Likewise, we find that in the orthonormal bases provided by $U$ and $V$, contributions of the inputs $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$ are typically approximately-sparse. Hence removing the dimensions with contributions summing to zero allows us to identify low-dimensional components of the inputs that are responsible for most of the attention head’s output (score). We illustrate the benefits of denoising $A^{\prime}$ in §5.2.

Hence, to fix $S_{ij}$, we consider the individual contributions made by each orthogonal slice to $A^{\prime}_{ij}$: $\{\mathbf{x}_{i}^{\top}D_{k}\mathbf{x}_{j}\}_{k=0}^{r-1}$. Empirically we typically find a few large, positive terms and many others that may be positive or negative. We seek to separate terms into ‘signal’ and ‘noise.’ To do so we adopt a simple heuristic, treating noise terms as a set that, in sum, has little or no effect on the attention head score. Accordingly, we define the noise terms to be the largest set of terms whose sum is less than or equal to zero. The indices of the remaining terms constitute $S_{ij}$. Terms denoted by $S_{ij}$ are strictly positive, and are the largest positive terms. Typically the number of those terms, ie, $|S_{ij}|$, is 20 or less, often just 2 or 3. We refer to this condition where $|S_{ij}|$ is small as the sparse decomposition of attention head scores in terms of the orthogonal slices of $\Omega$.

Given $S_{ij}$, we can decompose model residuals into ‘signal’ and ‘noise’ in terms of their impact on $A^{\prime}_{ij}.$ Define subspaces $\mathcal{U}=\operatorname{Span}\{\mathbf{u}_{k}\,|\,k\in S_{ij}\}$ and $\mathcal{V}=\operatorname{Span}\{\mathbf{v}_{k}\,|\,k\in S_{ij}\}$ and associated projectors $P_{\mathcal{U}}$ and $P_{\mathcal{V}}$. The denoising step separates the inputs $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$ into:

where $P_{\mathcal{U}^{\perp}}=I-P_{\mathcal{U}}$ and $P_{\mathcal{V}^{\perp}}=I-P_{\mathcal{V}}$. Then we have $\mathbf{x}_{i}=\mathbf{s}_{i}+\mathbf{z}_{i},\mathbf{x}_{j}=\mathbf{s}_{j}+\mathbf{z}_{j},$ and

where $|S_{ij}|$ is as small as possible.¹¹1To go from $\mathbf{\tilde{x}}$ to $\mathbf{x}$ we simply drop the last component; this is explained in the Appendix. Intuitively, we interpret a signal $\mathbf{s}$ to approximately represent a feature that is used for communication between attention heads; we present evidence in support of this interpretation below.

We use this framework to trace circuits in GPT-2 small as follows. A prompt corresponding to the IOI task is input to GPT-2. Consider th $a$-th attention head at layer $\ell$ generating attention score $A^{\prime}_{ij}$ for source token $j$ and destination token $i$. Then $S_{ij}^{\ell a}$ defines a set of orthogonal slices that the attention head $(\ell,a)$ is using. Hence we can approximately construct $A^{\prime}_{ij}$ as in (6), where we are using the SVD of $\Omega^{\ell a}$.

We will trace circuits causally with respect to the singular vectors of each attention head. For attention head $(\ell,a)$ generating output on tokens $(i,j)$, we identify the subspaces $\mathcal{U}$ and $\mathcal{V}$. We then look at each attention head ‘upstream’ of $(\ell,a)$, to determine how much each contributes in the subspace $\mathcal{U}$ to destination token $\mathbf{x}_{i}$ and in the subspace $\mathcal{V}$ to source token $\mathbf{x}_{j}.$ Specifically, for attention head $(l,b)$ with $l<\ell,$ denote the output according to (4) that it adds to token $i$ as $\mathbf{o}^{lb}_{i}$ (likewise for $j$). We then compute

Conceptually, $c^{\ell a,lb}_{i,ij}$ is an estimate of how much attention head $(l,b)$, by writing to token $\mathbf{x}_{i}$, has changed the output (score) of attention head $(\ell,a)$ on token pair $(i,j)$. We use $\sqrt{\sigma_{k}}$ to incorporate the magnitude of each singular vector’s contribution to the attention head output, dividing the contribution equally between the source and destination tokens. We refer to $c^{\ell a,lb}_{i,ij}$ as the contribution of attention head $(l,b)$ to attention head $(\ell,a)$ on token pair $(i,j)$ through token $\mathbf{x}_{i}$.

A singular vector trace of a model on a given input is a weighted, directed acyclic graph involving attention head pairs $(\ell,a),(l,b)$ and token pairs $(i,j)$ and having edge weights $c^{\ell a,lb}_{t,ij}$. It includes nodes that correspond to attention heads, and we add dummy nodes to allow for the separate contributions made through each token to be captured. Hence the graph consists of weighted edges of the type $(l,b)\to(\ell,a,\text{token})$ that capture the contributions, and edges of the type $(\ell,a,\text{token})\to(\ell,a)$ that capture which tokens are being attended by the attention head $(\ell,a)$.

To demonstrate the utility of singular vector tracing we apply it to a specific setting: the behavior of GPT-2 small when performing indirect object recognition (IOI). The IOI problem was introduced in [27] and that paper identified circuits that GPT-2 uses to perform the IOI task. As described in [27]: In IOI, sentences such as “When Mary and John went to the store, John gave a drink to” should be completed with “Mary.” To be successful, the model must identify the indirect object (IO, ‘Mary’) and distinguish it from the subject (S, ‘John’) in a prompt that mentions both. Thus, a succinct measure of model performance can be obtained by comparing the output logits of the IO and S tokens. The authors identified a collection of attention heads and the token positions they attend to when performing the IOI task. We refer the reader to [27] for additional details.

The IOI dataset that we use consists of 256 example prompts with 106 different names. We are using the same 15 templates used in [27], with two patterns (‘ABBA’ and ‘BABA’), which refers to the order of the names appearing in the sentence (the IO name is A, and the S name is B). Prompt sizes range from 14 to 20 tokens.

As in [27], we focus on understanding the interaction between attention heads. We believe that extending our framework to include the contributions of MLPs is an important direction for future work. Furthermore, in the analyses in this paper we do not consider cases in which heads attend to the first token – that is, when attention head has a large value of $A_{ij}$ for $j=0$. Because the attention weights for each target token form a probability distribution, when a destination token should not be meaningfully modified, attention heads normally put their weight on the first token. This role for token 0 has been noted in previous work [20].²²2GPT-2 does not use a ‘bos’ token.

Singular vector tracing as described in §4.2 can be applied at every attention head with respect to every token pair. However in our experiments we limit our analysis to cases where the attention head is primarily attending to a single source token for a given destination token. This strategy is similar to analyses in prior work [27, 2, 5]; it reduces complexity in the tracing analysis, but could be relaxed in future work. Specifically we say that an attention head is ‘firing’ if it places more than 50% weight on a particular source token for any given destination. This rule implies that a head can only ‘fire’ on one source token for each destination token. In general, we only trace upstream from attention heads that are firing on specific token pairs.

Note that interpreting (8) as giving the actual magnitude of the change in downstream attention score overlooks processing that may affect the signal in between the upstream head and the downstream head. Previous work has shown that downstream processing can compensate for upstream ablations [15] and that some layers may remove features added by previous layers [24]. We show results in §5.4 that confirm the direct causality of signals on downstream attention head outputs, but also illustrate that downstream processing can at times have a noticeable effect on model performance after signal interventions. Further, (8) ignores the impact of the layer norm. Attention head $(l,b)$ adds its output $\mathbf{o}^{lb}_{i}$ to the residual $\mathbf{x}_{i}$, but $\mathbf{x}_{i}$ will be subjected to the layer norm before being input to attention head $(\ell,a)$. We account for the effect of the layer norm using three techniques: weights and biases are folded into the downstream affine transformations, output matrices are zero centered,³³3These operations are provided by the TransformerLens library; see [20]. and the scaling applied to each token is factored into the contribution calculation. More details on tracing are provided in the Appendix⁴⁴4Code available at https://github.com/gaabrielfranco/sparse-attention-decomposition.

Throughout this and the next section, we use as our examples attention heads that figured prominently in the results of [27]. This aids interpretation and ensures we are paying attention to important components of the model.

We start by demonstrating the approximately-sparse nature of attention scores when decomposed via SVD. Figure 1 shows typical cases. Each plot in the figure shows results for attention head $(8,6)$ and a single source token, destination token, and prompt. The 64 contributions made by each orthogonal slice to the attention score are shown. Red bars correspond to sets $S_{ij}$; the sum of the red bars is approximately equal to the sum of all bars, which is the attention score for this head on these inputs. Figure 1(a) corresponds to the case in which the attention score is negative (and so the attention weight would be nearly zero); Figure 1(b) corresponds to a case in which the attention score is positive (28.1), and the attention weight is large ($\approx$ 0.83). Note that orthogonal slices are shown in order of decreasing singular value; the effect shown is not due to a low-rank property of $\Omega$ itself. Rather, the effect shown corresponds to sparse construction of the attention score when the inputs are encoded in the bases given by the SVD of $\Omega$.

We see considerable evidence in our experiments that the orthogonal slices used by an attention head are similar to each other when the attention head is firing. Figure 2 shows examples of the $S_{ij}$ sets for four attention heads: (3, 0), (4, 11), (8, 6), and (9, 9) across 256 prompt inputs. Furthermore, the nature of these attention head’s functions are evident in the sets of slices that they use. In the case of (8, 6) (an S-inhibition head), there is a set of about 6 slices that are consistently used; these are the same as the red bars in Figure 1(b). Attention head (9, 9) (a name mover head) uses a larger set, but there is still clearly a specific set of slices that frequently appear. Attention head (3, 0) (a duplicate-token head) uses a broad set of slices; this is consistent with its need to look at all the token’s dimensions, since its role is to detect identical tokens. And attention head (4, 11) (a previous token head) primarily uses a single slice; this also is consistent with its role of detecting adjacent tokens, which appears to only require detecting a feature in a one-dimensional subspace. Finally, we show in the Appendix corresponding plots for when these attention heads are not firing; there too, a consistent set of slices is used, but the slices are completely different from those used when the head is firing.

The sparsity of the contributions of the slices of $\Omega$ to attention scores is summarized in Figure 3. In Figure 3(a) we show the distribution of the number of orthogonal slices used, ie, $|S_{ij}|$, across all the prompts in our dataset. The figure shows that the number of slices used is consistently small, much smaller than the number of available dimensions (64). We also ask whether the sparsity observed is an artifact of the input data used in our experiments. To assess this, we run GPT-2 on a dataset consisting of non-specific inputs, having no relationship to the IOI task. We use the first 256 elements of The Pile [21] selecting the first 21 tokens from each element to match the IOI dataset size. The corresponding plot is shown in Figure 3(b). This comparison suggests that sparse decompositions of attention scores is not limited to the IOI setting.

Next we show that it is possible to leverage the sparsity of attention score decomposition. To illustrate this, we ask the following question: given a token that is input to a particular attention head, what similarity does it show to the outputs of upstream attention heads? In other words, could we trace some part of a circuit by simply looking upstream to see who has ‘contributed’ to the token?

We take as our examples heads (9, 9) and (10, 0), which are name mover heads. The authors in [27] find that functionally important contributors to the ‘end’ tokens of these heads are (7, 3), (7, 9), (8, 6), and (8, 10). The heatmaps in Figure 4 show contribution scores (computed via (8)) for two cases: the case where the signal is taken to be the entire residual $\mathbf{x}$, and the case where the signal $\mathbf{s}$ is taken to be just its low-dimensional component as described in §4.1.

The figure shows the strong filtering and correcting effect that results from exploiting the sparsity of attention decomposition. Figures 4(a) and (c) show that if we simply ask how much each upstream head contributes to the attention score of the downstream head, we get a very noisy answer with two problems. First, a large set of attention heads have high scores; and second, the known functional relationships from [27] are not evident. On the other hand, when we focus on Figures 4(b) and (d), we see considerable noise suppression; many heads in middle layers of the model are no longer shown as being significant. Furthermore, we see that one attention head (8, 6) stands out, and it is one of those previously identified as functionally important. And in the case of Figure 4(d), we see that other attention heads with known functional relationship (7, 3) and (7, 9), are also highlighted. In §5.3 we will show many other attention head pairs with functional relationships from [27] that are recovered using the singular vector tracing strategy. In the Appendix we show a hypothetical network trace performed without using singular vectors; the resulting trace is not usable and shows little evidence of known functional relationships in the IOI circuit.

Next we construct a singular vector trace using the concepts from §4.2. We start at a particular attention head and token pair. If the head is firing on the token pair (as discussed in §4.3) we obtain the subspaces $\mathcal{U}$ and $\mathcal{V}$ and associated projectors $P_{\mathcal{U}}$ and $P_{\mathcal{V}}$. We then look at each upstream attention head’s output, and separate from it the signal it contains for the downstream head. The properly adjusted magnitude of this signal constitutes the upstream head’s contribution to the downstream head’s attention score via the corresponding token, as computed via (8).

Empirically we find that the contributions from most upstream heads are small, with only a few upstream heads making large contributions to the downstream attention score. We filter out the small contributions, which are unlikely to have significant impact on model performance. To do this for a given downstream firing, we adopt the simple rule of choosing the smallest set of upstream heads whose contributions sum to at least 70% of the sum of all contributions.

Using this rule, for each token we identify the upstream heads with significant contribution through that token. An edge in the resulting trace graph is defined by the upstream head, the downstream head, the two tokens on which the downstream head is firing, and the choice of which token is being written into by the upstream head. For each upstream head we then ask whether it is firing on that token as a destination. If so, the process repeats from that head and token pair. The process is presented in detail as Algorithm 1 in the Appendix.

We ran singular vector tracing on GPT-2 small using 256 prompts from the IOI dataset described in §4.3. We started the trace at the three name mover heads (9, 6), (9, 9) and (10, 0), as they were identified as having direct effect on model performance in [27]. The resulting trace is shown in Figure 5; we refer to the graph in the figure as $G$. There are two kinds of edges in $G$: communication edges from heads to tokens, and attention edges from tokens to heads. In the figure, blue edges are toward tokens that are source tokens downstream, and red edges likewise are destinations. The width of the edge is the accumulated contribution of the edge over the 256 prompts. Darker nodes fired more often in our traces, and we have placed green borders around nodes that appeared in [27]. Edges that appear very few times (less than 65 times over the 256 prompts) are omitted.

We make a number of observations. First, the trace shows broad agreement with results in [27], although not all nodes in that paper appear in our graph; this may reflect differences of effect size, or differences between the input datasets. The trace shows agreement with previous results on head-to-head connections and also on the tokens through which the communication is effected. Note that in the appendix we show that further filtering this graph to just its most frequent edges yields a set of attention heads in agreement with [27] to a precision of 0.52 and recall of 0.69.

The trace also goes beyond previous results in a number of ways. Whereas in previous work, the upstream contributors to name mover source tokens were not identified, this trace shows where important features for the that (IO) token are added and that this happens very early in the model’s processing. Previous work also proposed that redundant paths were present in GPT-2’s processing for the IOI task; this trace confirms their existence and elucidates the nature of their interconnection pattern. For example, there is distinct lattice structure among nodes at layers 7, 8, and 9. (In the Appendix we isolate this lattice structure for better inspection.) We also identify highly active heads that were not discussed in previous work, including (2, 8) which attends to the ditransitive verb of the prompt and feeds into (4, 11) (a previous token head) as well as (8, 6) (a S-inhibition head); and (4, 3) which attends to the last token (which is always a preposition in our prompts) and its predecessor.

We note as well that singular value tracing itself requires only a single pass over the data, and as used here, doesn’t make use of counterfactual inputs. Most importantly, it identifies specific features that are present in the communication from one head to another, and which can be further investigated as to meaning and role.

We adopt a variety of strategies to validate the graph in Figure 5. To demonstrate the causal effect of each edge’s communication on model performance, we intervene on individual edges; and to demonstrate that the structure of the graph itself is functionally significant, we intervene on various collections of edges simultaneously.

The results in §§5.1 and 5.2 support the sparse decomposition hypothesis. In this section we discuss some possible mechanisms behind this phenomenon. First, a motivation for decomposing attention matrices using SVD comes from the following fact:

The proof is straightforward and provided in the Appendix. This suggests that model training could have the following effect. If an attention head needs to attend to particular vectors $\mathbf{x}$ and $\mathbf{y}$, it needs to output a large value for $\mathbf{\tilde{x}}^{\top}\Omega\mathbf{\tilde{y}}$. In that case, model training could result in construction of $\Omega$ in which one term of the SVD, say $\mathbf{u}_{k}\sigma_{k}\mathbf{v}_{k}^{\top}$, has $\mathbf{u}_{k}\approx\mathbf{\tilde{x}}/\|\mathbf{\tilde{x}}\|$, $\mathbf{v}_{k}\approx\mathbf{\tilde{y}}/\|\mathbf{\tilde{y}}\|$, and with $\sigma_{k}$ reflecting the importance of this term in the overall computation of the attention score. As an illustration, we note that the results in §5.2 and the Appendix show that in GPT-2 the singular vectors of some $\Omega$ matrices are correlated with word features that are relevant for the IOI task.

Consider a hypothetical case in which the sets of vectors to which the attention head needs to attend, say $\{\mathbf{\tilde{x}}_{i}\}$ and $\{\mathbf{y}_{i}\}$, happen to each form orthogonal sets. Then to achieve maximum discrimination power in distinguishing corresponding pairs, the singular vectors of $\Omega$, that is $\{\mathbf{u}_{i}\}$ and $\{\mathbf{v}_{i}\}$ respectively, should be aligned with the corresponding vectors in $\{\mathbf{\tilde{x}}_{i}\}$ and $\{\mathbf{\tilde{y}}_{i}\}$. To move from vectors to features, we refer to the linear representation hypothesis [17, 10, 23] which suggests that concepts, including high-level concepts, are represented linearly in the model.

To move closer to realistic cases, we note that the authors in [3] make observations, based on experimental evidence and geometric considerations, about features constructed by neural models. They argue that models will tend to represent correlated feature sets in a manner such that, considered in isolation, the sets are nearly orthogonal. They term this the use of “local, almost-orthogonal bases.” In our case, if an attention head is attending to vector sets $\{\mathbf{\tilde{x}}_{i}\}$ and $\{\mathbf{\tilde{y}}_{i}\}$ that are important when performing a specific task, then we may hypothesize that training will construct the sets to be “nearly-orthogonal,” meaning that cosine similarities among the vectors in each set would typically be small.⁵⁵5In the language of [3] we are making a “local non-superposition assumption.” In this case, the resulting sets of singular vectors $\{\mathbf{u}_{k}\}$, $\{\mathbf{v}_{k}\}$ are more likely to sparsely encode the $\{\mathbf{\tilde{x}}_{i}\}$ and $\{\mathbf{\tilde{y}}_{i}\}$ than exist in one-to-one correspondence.

Given the above argument, for cases where the linear representation hypothesis holds, we expect that an attention head is testing for a pair of low-dimensional subspaces in the inputs $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$. In that case, we expect that subsets of the singular vectors of $\Omega$ will be constructed during training so as to ‘match’ those subspaces. In this context, a sparse encoding allows the attention head to attend to more than $r$ different subspaces, expanding the number of concepts that the attention head can recognize.

Note however that attention heads have been shown to have a variety of functions, not all of which correspond to testing low-dimensional subspaces. For example, some attention heads have the role of detecting when tokens $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$ are identical [4, 27]. In that case, we expect that $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$ will have non-negligible inner products with most or all of the singular vectors of $\Omega$. In fact, we see exactly this phenomenon in the case of the duplicate-token head (3, 0) as shown in Figure 2(a).

There are a number of limitations of our study and directions for future work. First, the method as used here does not explore alternative pathways that may affect model output [13]. It also does not directly assess adaptive computations in the model [14, 24] nor does it construct minimal circuits in the sense of [27]. However, we believe that it offers an alternative toolbox that can be extended to help investigate those issues, in part by exposing the signals passing between specific pairs of attention heads.

Next, the noise separation strategy we use is not perfect, and is perhaps too permissive. We believe that an approach that determines sets $S_{ij}^{\ell a}$ differently may yield more precise results. For example, Figure 2 shows that certain orthogonal slices are frequently reused by certain attention heads, and a better $S_{ij}^{\ell a}$ identification strategy might use this observation.

Next, the tracing strategy used herein only considers edges that make positive contributions to attention scores. However, as Figure 1 shows, certain orthogonal slices contribute strongly negative terms to the attention computation. We have observed that these negative terms are important for determining when an attention head does not fire. The output of an attention head in general seems to depend on the balance between the few strongly positive terms and the few strongly negative terms. This is a fertile area for for investigation that could lead to deeper understanding of control signals in the model.

In this paper, we have restricted attention to heads that are firing (putting 50% or more attention weight on one source token). Not all attention heads work by firing on a pair of inputs. Conceptually there is no difficulty applying singular vector tracing to non-firing situations, but identifying likely causal paths and functions becomes more challenging.

Finally, we have not explored in depth the nature of the signals themselves. Indirect evidence, such as the difference in effect between local and global ablations, suggests that there are some similarities between signals used on different edges of the network. Exploration of the nature of signals and their relationships is an intriguing direction for future work.