Locating and Editing Factual Associations in Mamba
Authors: Arnab Sen Sharma
David Atkinson
David Bau
Venue: COLM (2024)
Keywords: language models, state space models, factual associations, model editing, interpretability
TL;DR: Many of the tools developed to interpret and edit factual associations in transformer LMs can be adapted to Mamba. Additionally, when it comes to factual recall, autoregressive LMs show similar locality despite differences in their architectures.

Abstract: We investigate the mechanisms of factual recall in the Mamba state space model.  Our work is inspired by previous findings in autoregressive transformer language models suggesting that their knowledge recall is localized to particular modules at specific token locations; we therefore ask whether factual recall in Mamba can be similarly localized. To investigate this, we conduct four lines of experiments on Mamba. First, we apply causal tracing or interchange interventions to localize key components inside Mamba that are responsible for recalling facts, revealing that specific components within middle layers show strong causal effects at the last token of the subject, while the causal effect of intervening on later layers is most pronounced at the last token of the prompt, matching previous findings on autoregressive transformers. Second, we show that rank-one model editing methods can successfully insert facts at specific locations, again resembling findings on transformer LMs. Finally we adapt attention-knockout techniques to Mamba in order to dissect information flow during factual recall. We compare Mamba directly to a similar-sized autoregressive transformer LM and conclude that despite significant differences in architectures, when it comes to factual recall, the two architectures share many similarities.

Published as a conference paper at COLM 2024
Locating and Editing Factual Associations in Mamba
Arnab Sen Sharma, David Atkinson, and David Bau Khoury College of Computer Sciences, Northeastern University
Abstract
We investigate the mechanisms of factual recall in the Mamba state space model. Our work is inspired by previous findings in autoregressive trans- former language models suggesting that their knowledge recall is localized to particular modules at specific token locations; we therefore ask whether factual recall Mamba can be similarly localized. To investigate this, we conduct four lines of experiments on Mamba. First, we apply causal tracing or interchange interventions to localize key components inside Mamba that are responsible for recalling facts, revealing that specific components within middle layers show strong causal effects at the last token of the subject, while the causal effect of intervening on later layers is most pronounced at the last token of the prompt, matching previous findings on autoregressive transformers. Second, we show that rank-one model editing methods can successfully insert facts at specific locations, again resembling findings on transformer LMs. Finally we adapt autentioonkkoocouu techniques Mamba in order to dissect information flow during factual recall. We com- pare Mamba directly conclude that despite significant differences in architectures, when it comes factual recall, the two architectures share many similarities.
similar-sized autoregressive transformer LM and
1 Introduction
Studies of autoregressive transformer language models' (LMs) processing of factual state- ments such as The Eiffel Tower is located in Paris have identified a localized pattern of internal computations when recalling facts (Meng et al., 2022a;b; Geva et al., 2023; Hernandez et al., 2023; Nanda et al., 2023), and have further found that those LMs can be edited by making single-layer rank-one changes in model parameters to alter specific fact. Although these localized phenomena appear to generalize across autoregressive transformer LMs, the extent to which similar locality might appear in very different architectures-uuce networks (RNNs)-has not yet been investigated.
recurrent
In this paper we investigate the internal mechanisms of Mamba (Gu & Dao, 2023), recently-proposed state-space language model, type of RNN that achieves per-parameter performance that is competitive with transformers. Specifically, we ask whether factual recall within Mamba exhibits locality similar to the patterns observed in autoregressive transformer language models.
Our paper is case study confronting a key methodological challenge that broadly faces interpretability researchers: as state-of-the-art neural network architectures evolve, we must ask, can the detailed analytical methods and tools developed for one neural architecture, such as transformer LMs, be generalized and applied to a different neural architecture, such as Mamba? In this paper we are able to answer the question with qualified "yes": we find that many of the methods used to analyze transformers can also provide insights on Mamba. We also discuss mismatches that is, interpretation methods (such as attention-knockouty that do not transfer to Mamba as easily due to architectural constraints.
We begin by studying whether activation patching (Wang et al., 2022) can be successfully applied to Mamba. Known variously as causal mediation analysis (Vig et al., 2020), causal
*Correspondence semsharma.@@horthesteem.. Code available at romba. comba.baulabb info

---

Published as a conference paper at COLM 2024
tracing (Meng et al., 2022a), and interchange interventions (Geiger et al., 2021), activation patching techniques can successfully identify specific model components in transformer LMs that play crucial roles in performing task. We ask whether Mamba can be productively studied the same way, even though the architectural components of Mamba are very different: for example, instead of attention heads and MLP modules, Mamba is composed of convolutions, gates, and state-space modules. To answer, we adapt activation patching to Mamba, and ask if any sparsity patterns emerge which provide insights into the respective roles of its components.
We also study whether rank-one model editing can be applied Mamba. While studies of transformers (Meng et al., 2022a;b; Hase et al., 2024) have found that there are range of MLP modules within which factual knowledge can be inserted by making a single rank-one change in parameters, Mamba does not have MLP modules, so we ask if there are any other modules that can be similarly edited insert knowledge. As with previous studies of transformers, the key question is whether factual associations can be edited with both specificity (without interfering with unrelated facts) and generalization (while remaining robust to rewordings of the edited fact).
Finally, we apply methods for understanding the overall information flows in Mamba. Inspired by Geva et al. (2023), we adapt attention-blocking methods to the attention-free Mamba LMs.
In this work we conduct our experiments Mamba-2.8b, the largest available LM in Mamba family, and for comparison we conduct the same experiments on the similarly sized Pythia-2.8b (Biderman al., 2023) autoregressive transformer LM.
Background on Mamba
Mamba, introduced in Gu & Dao (2023), is a recent family of language models based on state space models (SSMs). SSMs are designed to model the evolution of a hidden state across time with first-order differential equation (Koopman et al., 1999; Durbin & Koopman, 2012); when they are used as the recurrent state of an RNN, they can enable highly efficient parallelized training (Gu et al., 2021). To achieve good performance in language modeling, the Mamba SSM introduces input-dependent parameterization or selective-SSM instead of the traditional time-invariant SSMs. Mamba uses a special architecture called MambaBlock' which is stacked homogeneously, replacing both attention and ML. blocks used in transformer layers. Here, we focus on the different operations performed inside MambaBlock.
Formally, Mamba is an autoregressive language model: M : X -> y over a vocabulary V that maps sequence of tokens x = [x1, X2,.,,,
E X, x; E V to y E y C R which is a proba- bility distribution over the next token continuations of x. Similar to other deep LMs, in Mamba, a token x; is first embedded to hidden state of size d as ([(0) emb(x;). Then n!(0) is transformed sequentially by series of
w
q
Conv SSM
W
w
g
Figure 1: Architecture of a MambaBlock. Projection matrices W and W have the shape 2d x d, while w has the shape x 2d. a,,,,, and are intermediate states of token representation. o is SiLUU activation and @ is elementwise multiplication. Cono + SSM operation abstracts the ConvlD and selectize-SSM operations.
MambaBlocks. The hidden state h: after the tt (1-indexed) MambaBlock is computed as follows:
( - =n (- + +o²
(E)
where of is the output of th MambaBlock for the 7th token
In their paper, Gu & Dao (2023) call this component Mamba- the same name as the LM family.
2

---

Published as a conference paper at COLM 2024
Average IE restoring h," over 400
Clean
Patched Run, Kun,Gil
[[=k]
[PAD]
subject
0.6
subject 2nd last"
carly sit
Jordan
subject last"
professionally
professionally
further tokens
late site
0.2
played
played
last token
0.0
besketball
plans)
5 10 15 25 30 40 50 55
single restored layer within Mambh-2.8b
(a) Activation patching
(b) Tracing the residual states, It
Figure 2: (a) Activation patching. state from the clean run is patched into its corresponding position in the corrupted run G*. This has a downstream effect of potentially changing all the
states that depend the patched state in G*(s ni()). (b) Average IE of applying causal tracing residual stream states (hç in Figure 1) across 400 different facts from the RELATIONS dataset (See Appendix A.2).
MambaBlock (11 (1(5-1) 11h""),,1-1 ((.,,),,.,,((111))) w!!
(s/"
(0)
(2)
Here, @ represents element-wise multiplication or Hadamard product. s/() is calculated as:
( (! n/(
(3)
(6) ( SiLU ( [ConvD (2) (2)
(0) 2
"i ....,a!)))
(4)
(8) c
selective-SSM
(5)
We abstract the operations Equations 4 and 5 as the Conv + SSM operation in Figure At high level, Conv SSM brings information from the past token representations to the current token representation, similar to the attention blocks in transformer LMs. But, unlike attention operation, Conv + SSM scales linearly with the context length and thereby enjoys faster inference speed and longer context limits. See Gu & Dao (2023) for details.
The output of the other path gi (that does not pass through Conv SSM operation) is gating mechanism that regulates the information flow. This gating mechanism resembles parts of LSTM (Hochreiter & Schmidhuber, 1997) and GRU (Cho et al., 2014) networks, where similar gates control selective updates of recurrent state.
(E) = SiLU (0
(6)
In the remainder of the paper, we aim to characterize the role of the components of Mamba in factual recall by adapting tools that have previously been used to analyze transformers.
In Section 3, we apply activation patching to localize factual recall as in Meng et al. (2022a), testing the roles of states 5; 8i, and o at all layers. Section 4, following Meng et (2022a); Hase et al. (2024), we test rank-one edits of facts across components W., Wg and W. at each layer. And in Section we address the challenge of applying attention patching in Mamba, used in Geva et al. (2023) to isolate information flow in GPT LMs. We include a further investigation of Mamba's factual recall mechanisms Appendix E, following Hernandez et al. (2023) in collecting Jacobians within Mamba to test the linearity of its relational encodings.
3 Locating Key States for Factual Recall
We begin with activation patching, seeking to understand there are specific hidden states which play important roles during factual recall. We select fact (s,r, that the LM knows, where r is relation that associates a subject entity with an object entity o. To estimate each state's contribution towards correct factual prediction (s Michael Jordan, = professionally played, basketball), we collect model activations across three different runs:
3

---

Published as a conference paper at COLM 2024
Average restoring 0, over
Average of restoring
(c) Average restoring
0.5
0.5
0.5
prefix
subject [=]]*
subject 2nd last"
earyy sile
0.3
curly sine
0.3
0.3
subject last"
0.2
0.2
02
futter wkkens
Iale site
0.1
Soio k/cc
list token
0.0
0.0
0.0
10 25 30.35 40.45 50 55
1. 20 30 40
5 10 15 20 25 30 35 40 50 55
on patchedl layers
interval patchad laycrs
interval
layors
Figure 3: Average indirect effect of different states of 8i and over 400 facts from the RELA- TIONS dataset (see Appendix A.2). For each layer f, states for window of 10 layers around & are restored from the clean run
clean run G: In the clean run, we run the model on a prompt specifying the fact we are interested in. For example, x = (s, Michael Jordan professionally played. We cache all the hidden states during the clean run to be used later: { (f) (2) | [1,T], , L [1, []
corrupted run G*: that the LM gives a different answer o² o'ooccer) to the modified prompt x = (s*, (i.e., o" O).
In the corrupted run, we swap with different subject s* (Pelé) such
This subject-swapping approach follows the recommendation of Zhang & Nanda (2023) and has the ad advantage of using natural text perturbations to avoid introducing out-of-domain states to the model's computation, as may happen when corrupting s embeddings with Gaussian noise (the method used in Meng al. (2022a)).
patched run G*[+ ] ]: In the patched run, we run the model on the corrupted prompt
x*, but intervene on n, by replacing its value with the corresponding state cached from the clean run G. The remainder of the computation is run normally, meaning that the patched state have downstream effect of potentially changing all the states that depend it. See Figure 2a.
Let p(o), (o), and p*[(- (0) denote the probability assigned to the correct answer o in G, , and G uy respectively. To measure the contribution of I1 in recalling the fact (s),, we define its indirect effect (IE) as:
p-[+ (0)) (p) - -
IE
In Figure 2b we plot the average indirect effect of restoring the residual states (1)) across different layer-token positions over 400 facts from the RELATIONS dataset (Hernandez et al., 2023). The high IE observed at the late site (later layers at the last token) position is natural, as restoring clean ! there will restore most of the model computation from G. However, Mamba also shows high causality at the early site (early-middle layers the last subject token position). This is consistent with what Meng et al. (2022a) observed in GPT LMs.
In Figure 3 we plot the average IE for o. 8i and s! The plot for ( (Figure 3a) looks very similar to Figure 2b, confirming that the output from MambaBlock has strong causal effects at both early and late sites. Interestingly, Figure 3c shows that the selective-SSM
(2)
outputs s (f) have high IE only at the late site, resembling the behavior of attention modules in GPT models (Meng et al., 2022a). However, there is state that appears to do the opposite; in other words, there is no state with strong effects at the early site and not at the
late site (The gate output (*) does have stronger IE at the early site, but these effects are very weak). To compare with autoregressive transformer LMs, activation patching results for Pythia-2.8b is shown in Figure 9 in Appendix D. This comparison reveals a key way how
Mamba differs from transformers: while transformer ML. outputs have effects in the early
4

---

Published as a conference paper at COLM 2024
IE with by restored
s blocked
8, blocked
0 blocked
1,00
subj last
0.75
Average
0.50
0.25
0.00
(b) j prompt last
1,00
Average
0.50
0.25
0.80
15
35
50
Layers
Figure 5: Impact of ablating % 8t and 0; IE 18,5 for (a) subject last and (b) prompt last token positions. Taken together (a) and (b) show a clear separation roles between early-mid and later layers in Mamba- 2.8b. 1/!") up to layer 46 only show strong IE the subject last token position and have negligible impact after that. Whereas TE of kg jumps to 1.0 after layer 46. (a) also shows that, at the subject last token, before layer 27 28, 1 is significantly reduced by blocking either 0, 8i, or paths (sorted
in descending order of damaging ya)) (b) At the prompt last token, ablating o; paths can significantly reduce 18,0 in layers 47 50.
site and not the late site, Mamba there no similar state that specializes only at the early site, at which factual recall would be expected to occur. This presents the question: which parameters in Mamba mediate factual recall?
replicate an experiment from Meng et (2022a) to probe path-specific effects (Pearl, 2022) by sev- ering a path from the causal graph and mon- itoring its effect. Here, we interested in understanding the effect of the contri- butions from gi, S;, and o; (ie. states that are processed by Wg. Conv SSM, and W respectively) while recalling a fact. First, in the corrupted run G', at token position i, we cache all the contributions from the paths as 87 {5,
To investigate this question,
G
ha) from G-
(b) Blocking &
G'-A1
(0) E [], Then
in the patched run G*(+- 1,,",, we restore
g
( that was cached from the clean run G into its corresponding state (as in Figure 2a), but with an additional modification: derstand the contribution from the s; paths, we sever those paths by also patching se (cached from the corrupted run G*) to their corresponding locations (see Figure 4). The same experiment replicated to understand
Figure 4: To probe for path-specific effects, (a) n restored from the clean run G (as in Figure 2a). (b) Then, reveal the role of the Conv SSM contribu- tions, s states from the corrupted run G* are also patched to block the contributions from paths.
un-
the contributions of gi and o states. We note that severing the of will sever s(0) and g) as well (see Figure 1).
In Figure 5 we plot the average results of this experiment for token positions (a) subject last and ib) = prompt last over 400 examples randomly sampled from the RELATIONS dataset. The key findings can understood by examining the gap between the purple bars and the red, green, and blue bars: a large gap indicates a strong mediating role for Conv + SSM, Wg or W. parameters, respectively. At the early site at the subject last token, both the Conv SSM and W have strong role, but W, plays an either. Yet the strongest mediator at the late site is also W.. This experiment highlights the importance of W, in both stages of predicting fact. But it also suggests that Mamba
larger role than
5

---

ROME generates a rank-one update to wil)) down' the down-projection matrix of the MLP module for the last token of the subject at layer k-which plays the role of the associative Published as a conference paper at COLM 2024
does not separate early-site factual recall between these groups of parameters cleanly as transformers. However, Figure 5 reveals a clean separation of roles between early to mid and later layers, analogous to the findings of Hernandez et al. (2023) in transformer LMs. We also note that this division of responsibilities between layers can be more sharply noticed in Mamba when compared to transformers LMs (compare Figure 5 with Figure 10).
Editing Facts With ROME
Having begun to characterize the locations of important states for factual recall, we now investigate whether factual recall behavior can be edited. In particular, we apply the ROME (Rank One Model Editing, Meng et al., 2022a) technique to Mamba. ROME begins with the observation that any linear transformation can be considered as an associative memory (Anderson, 1972; Kohonen, 1972), mapping a set of keys K = [kg k2]k2]... . to their corresponding values V = [ag [2]2 ...], and uses this to edit factual associations in trans- former LMs. Here, we apply the technique to particular set of linear transformations within Mamba, and report our editing success each.
The input to ROME is a prompt x = rs, where s (Emmanuel Macron) is a subject entity and (is the President of) is a relation. ROME also takes counterfactual object o* (England), meant to replace the correct object o (France) in the model's output. To effect that change,
memory. In generating the rank-one update, ROME considers the input to w(() "doum as the key (k.). Then, with gradient descent ROME calculates value (v.) such that, when v. is inserted as the output of w(t)
downn' the model will output 0". Importantly, while optimizing V+, ROME attempts to minimize unrelated changes in model outputs (Joe Biden, for example, should still be mapped to the United States post-edit). Finally, ROME adds a rank-1 matrix A
to ((() such that (wlt)
down + + = va (See Meng et al. (2022a) for details.)
down
4.1 Applying ROME in Mamba
We apply ROME on the three different projection matrices of Mamba: w(O which affects only the Conv SSM path, w( which affects only the gating path, and w! the final output of the MambaBlock, which is added to the residual state. We plot ROME performance on different projection matrices (w(+), w(f), and w!()) across all the layers in Figure 6a.
To evaluate editing performance, we use the COUNTERFACT dataset from Meng et al. (2022a). COUNTERFACT contains 20K counterfactual examples in the form (s,r,0 0²), where is the correct answer to the prompt x = rs,r), and o* is the object which is to be inserted as the new answer to the prompt (See Appendix A.1 for details). We select the first 2000 examples
from this dataset for our module-layer sweep. We use the original evaluation matrices in Meng et al. (2022a) to measure ROME edit performance. The final score (S) in the ROME evaluation suite is the harmonic mean of three different scores:
1. Efficacy (ES): For an edit request (s,r,0 + 0*), we say the edit is effective if, post-edit, the LM assigns p(o") > p(o) in response to the prompt (s,r). Efficacy reflects the portion of the examples where the edit was effective.
2. Generalization (PS): A successful edit should be persistent across different paraphrases of (s, For each of the request instances (,, 0 + o*), p(o*) > p(o) is checked post-edit with set of different rephrasings Xp (s) of the prompt x (s, r), where P, denotes a set of paraphrased templates for the relation r.
Further motivating these experiments, previous work has shown that the locations identified by activation patching techniques are not necessarily those which have the strongest edit performance (Hase et al., 2024).

---

Figure 6a shows that ROME can achieve high scores (S) for range of early to middle layers by modifying any one of the projection matrices (())) w( w!( matching observations made by Hase et al. (2024) regarding transformer LMs. However, we found that performance Published as a conference paper at COLM 2024
W
W
(a) ROME performance on Mamba-2.8b
Efficacy (ES)
Generalization (PS)
Specificity (NS)
Score
100
100
100
100
75
WMM
75
25
25
25
25
0
6" 10" SO as 30
0
Eddt Layer
Edit Layer
Edit Layer
Edit Layer
W down
(b) ROME performance Pythia-2.8b
Efficacy (ES)
Generalization (PS)
Specificity (NS)
Score (S)
100
75
75
75-
75-
50
50
50-
50-
25
0
15215 16 30
12 24
0
912 18 24
Tayer
Edit Layer
Tayer
Edit Layer
Figure 6: ROME performance in editing facts across different layers (a) by modifying w! W
w! in Mamba-2.8b, and (b) modifying w in Pythia-2.8b. Results are reported the first 2000 examples in the COUNTERFACT dataset.
down
Specificity (NS): Finally, the edit should be specific to (s(s) and should not additionally change the mapping of some nearby subject Sn to o². To evaluate the specificity of an edit we measure p(o,) > p(o") with (s,),) for a set of nearby factual associations ((S,, On # o²}.
does depend on the location of the edit. For example, in the case of w( and W!!), the score (S) and generalization (PS) drops after around layer 43. This is consistent with our
findings from the path-blocking experiment in Figure 5a. We also find that edits to w!( have poor generalization (PS) in early layers, whereas high PS be achieved at early layers by modifying either w(() w! consistent with their higher indirect effects as seen in Figure 5a.
Where is the right place to apply ROME on Mamba? Figure could suggest w the causal effect of gi states mostly concentrated the subject last token, similar to the behavior of MLPs in transformers (Meng et al., 2022a). Consistent with this is the
g since
down connects to attention modules only architectural fact that, just as transformers' (1())
through the residual stream, the output of w!()) does not flow through the Conv + SSM module-a module that other work has suggested might play a role similar to that played by attention heads transformers (Grazzi et al., 2024). And, indeed, we find that ROME
can successfully insert facts by modifying (!()) On the other hand Figure 6a reveals sudden
drops in efficacy and generalization at middle layer gates, suggesting that w(()) may an unreliable mediator at some layers. Our experiments further show that the best performance
for ROME is empirically achieved by modifying w(?). This is consistent with the fact that o; states show stronger causal effect at the subject last token than &; states do (see Figures 3a and 5a). Additionally, ROME achieves better generalization (PS), competitive specificity (NS), and an overall better score (S) with w!!) We hypothesize that the strong performance of w may be due to the the separation of roles between early-mid and later layers observed in Figures 2b, 3a, and 5. Also Appendix C where isolate the contribution of w! by
7

---

Published as a conference paper at COLM 2024
0.50
Relative change in p(o)
0.25
-0.25
-0.50
-0.75
-1.00
20
25
30
35
40
45
50
Layer
subject
subj last
non-subject
prompt last
( future tokens via sj paths is Figure 7: Relative change in p(o) when information flow from dg blocked, with taking the value of either subject, non-subject, or the prompt last token positions. For each layer E, 9; paths were blocked for window of 10 layers around &.
from IE (0), which reveal a critical role of w! in early-mid layers at subject last token position while mediating fact.
subtracting 1E (0) + IE
8.0
We plot ROME performance for a similar sized Pythia model on Figure 6b for comparison.
Attention Knock-out in Mamba?
Attention modules mediate the flow of information across different token positions trans- former LMs. In attention "knock-out" experiments the information that flows through specific edge (from 1.th token 9th token) via certain attention head blocked to under- stand if critical information flows through that edge. This is also form of causal mediation analysis and has been effective in understanding the information flow transformer LMs (Geva et al., 2023; Wang et al., 2022; Todd et al., 2023). In Mamba, information from past tokens is retained in the s; states, with the Conv + SSM operations (see Figure and Equa- we perform experiments similar to attention knock-out experiments in Mamba in order to understand how it moves factual information?
tions 3-5). We ask,
We find that performing similar experiments in Mamba be difficult. The use of Conv with a non-linearity in conjunction with selective-SSM make challenging to remove the information retained in the 4th token from the hth token (see Appendix for details). However, it is possible to block the propagation of information from the k token to all the future tokens via Conv SSM with mean-ablation. Specifically, for a layer L, we set
(8) == B [a(9), where IE [a((1] is the mean of a(() states collected with 10,000 tokens from WikiText-103 by Merity et al. (2016). We recognize that this intervention may not be as surgical as cutting specific edge. However, with some caveats, this experiment suggests
that the factual information flow in Mamba is similar to what Geva et al. (2023) observed in LMs.
We randomly sample 700 facts across factual relations from the RELATIONS dataset. For each those examples we block-out information propagation of the subject, non-subject, and the prompt-last token positions for window of 10 layers around specific layer f. The effect of blocking out Conv SSM information flow for certain layer-token (E k) positions
measured the relative change in p(o) with (p(o)a("-- (10 (((o))pp)))) /p(o). Figure 7 shows the averaged result and it leads us to draw the following conclusions about how factual information flows in Mamba:
8

---

Published as a conference paper at COLM 2024
(a) The purple lines show that blocking out non-subject information flow in early middle layers can bring down p(o) by up to 50%. Non-subject tokens are used to specify the relation r. This observation leads us to believe that Mamba propagates relation specific information to future tokens using Conv+SSM operations in early-middle layers.
(b) Interestingly, the green lines (blocking the subject information flow) shows two valleys:
1. The first valley at the early layers is not surprising as Mamba needs to collate in- formation from all the subject tokens early layers to recognize a subject entity s consisting of multiple tokens.
2. However, the valley at layers 43-48 suggest that Mamba uses Conv + SSM paths in those layers to propagate critical information from the subject to later tokens. This aligns with Figures 5b and 3c, where states in those layers show high indirect effects, indicating their crucial role while recalling a fact.
(c) The blue dashed lines indicate the effect of blocking the information of only the subject last token. If the ablation is performed in very early layers, later layers start to compensate for that. However, the valley around layers 20-21 suggests that Mamba expects to recognize the full subject entity by then in order to recall relevant associations (enrichment). Notably, activation patching results for o; and s--states that we hypothesize take crucial part in the enrichment process-als show strong indirect effect around that region (Figures 3a, 3b, and 5a). The blue line follows the green line after layer 30. The weaker effect observed might be because ablating subject last token is not always enough to remove all the subject information. For example, in Eiffel Tower, Eiffel (tokenized as E, iff, el) is more informative than the last token Tower.
These findings align with how factual information flows through attention modules in autoregressive transformer LMs, as observed by Geva et al. (2023) in GPT. However, unlike Geva et al. (2023), we cannot make strong claims about the unique role of the final token position (prompt-last) with this experiment. As we block out information flow to all future tokens, the intermediate states in between the ablated kd token and the last token are affected as well.
6 Related Works
Mamba. Mamba is recent family of language models that are based on state space models (SSMs). Neural SSM-based models have achieved good performance across different modalities, including vision (Nguyen et al., 2022), audio (Goel et al., 2022), and genomic sequences (Nguyen et al., 2023). Only recently, however, with Mamba, have they become competitive with the language modeling performance of transformers (Gu & Dao, 2023). Like transformers, Mamba contains factual knowledge about real world entities (Grazzi et al., 2024). However, knowledge representation in Mamba (and other LMs based on SSMs) has up to now remained unexplored.
There are few works focused on interpreting Mamba. Ali et al. (2024) identify implicit attention-like matrices formed by Mamba's selective state space layers. Grazzi et al. (2024), while not strictly focused on interpreting Mamba's internals, apply linear probes Mamba's (decoded) intermediate states during in-context regression tasks. They discover substantial similarities between Mamba and transformer models: both architectures pursue "iterative" strategies, with the task loss falling more or less monotonically as the layer index increases. Finally, in concurrent work, Paulo et al. (2024) apply different set of interpretability techniques to Mamba, focusing on activation steering and linear probing methods. Like us, they find that approaches developed for transformer models largely transfer to Mamba.
Locating Factual Knowledge in Language Models. To make factually correct statements about the world, LM has to store factual knowledge about real world entities somewhere in its parameters. Understanding how and where neural network stores knowledge is a core problem for interpretability and it has thus been studied from several perspectives (Ji et af., 2021; Wang et al., 2014). One line of work trains classifiers to probe for properties encoded in model representations (Ettinger et al., 2016; Shi et al., 2016; Hupkes et al., 2018; Conneau et al., 2018; Belinkov et al., 2017; Belinkov & Glass, 2019). However, the flexibility
9

---

Published as a conference paper at COLM 2024
of these classifiers lead to overestimating model knowledge and capabilites (Belinkov, 2022). Causal mediation analysis methods (Pearl, 2022) attempt to measure the causal contribution of intermediate states to task performance. Meng et al. (2022a;b) use activation patching to identify key MLP modules for factual recall, highlighting the middle layers at particular token positions as being especially important. Similarly, Geva et al. (2023) apply causal mediation analysis to attention modules, seeking to understand the mechanism of cross-token factual information flow inside transformer LMs.
7 Discussion
In this paper we have set out to understand whether the analytical methods and tools developed for transformer LMs can also be applied the Mamba LMs based recurrent state-space architecture. Although our experiments have been limited to Mamba-2.88 2.bb, the largest available LM of that family, and comparisons the similarly-sized transformer Pythia-2.8b, the methods we have investigated are general, and can be used to analyze to any state-space model.
Our overall comparisons of Mamba and transformers are positive: with activation patching we have found that, similar autoregressive transformer LMs, Mamba shows signs of localization at the last subject token and at specific layer ranges while recalling a fact. Although, unlike transformers, Mamba has no MLP modules, we find that their W. weights can receive rank-one model editing (ROME) edits with good generalization and specificity at range of layers, similar to W Womn in Pythia and GPT family of LMs. We have also been able to partially adapt the tools of attention knock-out in Mamba by blocking outgoing information from a token, revealing information flows similar to transformer LMs during factual recall.
The similarity that we have observed between factual recall mechanisms transformers and Mamba leads to speculate that the autoregressive language modeling task itself induces a pattern of localized factual recall that is independent of modeling architecture. When constraining a model to process text from beginning to end, the ordering creates a specific bottleneck in the information flows: the end of a subject becomes singular moment at which recognition of the subject is both possible and useful, and we find that both transformers and Mamba arrange their computations to localize factual recall at that moment. We hypothesize that other future autoregressive LMs architectures should expect to see similar locality in factual recall as well.
In summary, we find that many the tools used to interpret and edit large transformers can be adapted to work with Mamba, and we are optimistic that those tools will continue to be useful as architectures continute to evolve.
Ethics
By exploring the factual recall mechanism in Mamba, we potentially improve its trans- parency, enabling oversight and control. However, the ability to modify facts directly in the model brings with it the potential for abuse, such as adding malicious misinformation or bias.
Reproducibility
We ran all experiments on workstations with either 80GB NVIDIA A100 GPUs or 48GB A6000 GPUs, using the HuggingFace Transformers library (Wolf et al., 2019) and PyTorch (Paszke et al., 2019). We make use of publicly available datasets COUNTERFACT and RELATIONS in this work.
10

---

Published as a conference paper at COLM 2024
Acknowledgements
This research has been supported by a grant from Open Philanthropy (DB, AS), and an NSF Computer and Information Science and Engineering Graduate Fellowship Fellowship (DA). We are also grateful to the Center for Al Safety (CAIS) for sharing their compute resources, which supported many our experiments. Some of our initial analyses were conducted with beta version of NNsight Frotto-Kauffmn et al., 2024) an implementation of Mamba instrumented for research by Jaden Fiotto-Kan Noott-Kauffann.
References
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXio preprint arXiv:2403.01590, 2024.
James A Anderson. A simple neural network generating an interactive memory. Mathematical biosciences, 14(3-4):19-2220, 1972.
Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48(1):207-219, 2022.
Yonatan Belinkov and James Glass. Analysis methods in neural language processing: A survey. Transactions of the Association for Computational Linguistics, 7:49-72, 2019.
Yonatan Belinkov, Nadir Durrani, Fahim Dalvi, Hassan Sajjad, and James Glass, What do neural machine translation models learn about morphology? arXiv preprint arXir:1704.03471, 2017.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdannu,, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXivv1449.1220 1259, 2014.
Alexis Conneau, German Kruszewski, Guillaume Lample, Loic Barrault, and Marco Baroni. What you can cram into a single vector: Probing sentence embeddings for linguistic properties. arXiv preprint ar Xiv:18805 000.0000000000 2018
James Durbin and Siem Jan Koopman. Time series analysis by state space methods, volume 38.
OUP Oxford, 2012.
Yanai Elazar, Nora Kassner, Shauli Ravfogel, Abhilasha Ravichander, Eduard Hovy, Hin- rich Schütze, and Yoav Goldberg. Measuring and improving consistency in pretrained language models. Transactions of the Association for Computational Linguistics, 9:1012-1031, 2021.
Allyson Ettinger, Ahmed Elgohary, and Philip Resnik. Probing for semantic evidence of composition by means of simple classification tasks. In Proceedings of the 1st workshop cvaluating vector-space representations for nlp, Pp. 134-139, 2016.
Jaden Fiotto-Kaufman, Alexander R Loftus, Eric Todd, Jannik Brinkmann, Caden Juang, Koyena Pal, Can Rager, Aaron Mueller, Samuel Marks, Arnab Sen Sharma, et al. Nnsight and ndif: Democratizing access foundation model internals. arxiv preprint ar arXit:2407.14561 2024.
Atticus Geiger, Zhengxuan Wu, Hanson Lu, Josh Rozner, Elisa Kreiss, Thomas Icard, Noah D. Goodman, and Christopher Potts. Inducing causal structure for interpretable neural networks. CoRR, abss2211.00022 2021. htttpp://awxiv.org/bbb////////// 00826.
11

---

Published as a conference paper at COLM 2024
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in ruto-regressive language models. arXiv preprint arXio:2344.14767, 2023.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Re. It's raw! Audio generation with state-space models. Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 7616-7633. html.
PMLR, 17-23 Jul 2022. URL https: http:::://////dd///oeeee mlr www.mmrrrrrrrr wrress/v12/ggoell/ooee
Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is Mamba Capable of In-Context Learning?, 2024. URL
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arxiv preprint ar arXiv:2312.00752 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. Xiv preprint arXiv:2111.00396, 2021.
Peter Hase, Mohit Bansal, Been Kim, and Asma Ghandeharioun. Does localization inform editing? surprising differences in causality-based localization vs. knowledge editing in language models. Advances in Neural Information Processing Systems, 36, 2024.
Evan Hernandez, Arnab Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, and David Bau. Linearity of relation decoding in transformer language models. ar Xiu preprint ar Xrxt::308.00111 099224, 2023.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, (8):1735-1780, 1997.
Dieuwke Hupkes, Sara Veldhoen, and Willem Zuidema. Visualisation and diagnostic classifiers' reveal how recurrent and recursive neural networks process hierarchical structure. Journal of Artificial Intelligence Research, 61:907-926, 2018.
Shaoxiong Ji, Shirui Pan, Erik Cambria, Pekka Marttinen, and S Yu Philip. A survey knowledge graphs: Representation, acquisition, and applications. JEEE Transactions on Neural Networks and Learning Systemss, 33(2):494-514 2021.
Teuvo Kohonen. Correlation matrix memories. IEEE transactions on computers, 100(4):
353-359, 1972.
Siem Jan Koopman, Neil Shephard, and Jurgen A Doornik. Statistical a algorithms for models in state space using ssfpack 2.2. The Econometrics Journal, 2(1):107-160, 1999.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35:17359-17372, 2022a.
Kevin Meng, Arnb Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. Mass- editing memory in transformer. arXiv preprint xrXi:22100007229, 2022b.
Stephen
Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843, 2016.
Neel Nanda, Senthooran Rajamanoharan, János Kramár, and Rohin Shah. Fact finding: Attempting to reverse-engineer factual recall on the neu- ron level, 2023. URL fact-finding- -rrrregnnntttteeeeeeeeeeeeeeeere
Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos as multidiimensional signals with state spaces. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 2846-2861. Curran http:::/pprcceeennggs..... / neurips..............m cc/ppp..cc//ppa htcpppper__ff faper__iles/paper/2200 p2p2r/222// .df..
Associates, Inc. 2022. URL https:
c81 1
12

---

Published as a conference paper at COLM 2024
Eric Nguyen, Michael Poli, Marjan Faizi, Armin Thomas, Callum Birch-Sykes, Michael Wornow, Aman Patel, Clayton Rabideau, Stefano Massaroli, Yoshua Bengio, Stefano Ermon, Stephen A. Baccus, and Chris Ré. Hyenadna: Long-range genomic sequence modeling at single nucleotide resolution, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An im- perative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Gonçalo Paulo, Thomas Marshall, and Nora Belrose. Does transformer interpretability transfer to rnns? arXiv preprint arXiv:2404.05971, 2024. URL https://arxxv..or 2404.05971.
Judea Pearl. Direct and indirect effects. In Probabilistic and causal inference: the works of Judea Pearl, pp. 373-392. 2022.
Xing Shi, Inkit Padhi, and Kevin Knight. Does string-based neural mt learn source syntax? In Proceedings of the 2016 conference on empirical methods in natural language processing, pp. 1524-1534, 2016.
Eric Todd, Millicent Li, Arnab Sen Sharma, Aaron Mueller, Byron Wallace, and David Bau. Function vectors in large language models. arXiv preprint ar Xiv:2330015214. 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. Investigating gender bias in language models using causal media- tion analysis. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 12388-12401. Curran Associates, Inc., 2020. URL File//255bb2222777555e2222222222
wewpeeddiiinnnnnniiiseeeeeeeeeee ccppep. laaper_fflese files/ paaless/aper//2228 .ddf.
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: circuit for indirect object identification in gpt-2 small. ar arviv preprint Xit:2211.00593, 2022.
Zhen Wang, Jianwen Zhang, Jianlin Feng, and Zheng Chen. Knowledge graph embedding by translating on hyperplanes. In Proceedings of the AAAI conference artificial intelligence, volume 28, 2014.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. ar arviv preprint arXiv:1910.037711, 2019.
Fred Zhang and Neel Nanda. Towards best practices of activation patching in language models: Metrics and methods. ar Xio preprint ar XrX:::2309.166 16042, 2023.
13

---

Published as a conference paper at COLM 2024
A
Datasets
We use two datasets; COUNTERFACT by Meng et al. (2022a) and RELATIONS by Hernandez et al. (2023) in this work.
A.1 COUNTERFACT
Meng et al. (2022a) developed the COUNTERFACT dataset for evaluating the efficacy of counterfactual edits in language models. It was prepared by adapting PARAREL (Elazar et al. (2021)) and scraping Wikidata'. The dataset contains 21, 919 requests {s,r,0,0*, π²" where the correct answer to the prompt x = (s,r), o* the counterfactual edit request, and π²' ~ rs),r) r) is paraphrase of the prompt = rs, to test for generalizability (PS). Each of the records also contain some neighborhood prompts TEN to test for specificity (NS) and some generation prompts TG to test if LM generation post-edit is fluent and consistent with the edit. Please refer to Meng et al. (2022a) for details on the curation of this dataset.
We evaluate ROME performance in Mamba-2.88 (Figure 6a) and Pythia-2.8b (Figure 6b) the first 2000 records from COUNTERFACT.
A.2 RELATIONS
The RELATIONS dataset introduced in Hernandez et al. (2023) consists of 47 relations of types: factual, linguistic, bias, and commonsense. A relation r is an association between two entities. For example, the relation, r professionlly played the sport connects the subject s = Michael Jordan with the object o = basketball. The dataset contains set of (s,0) for each relation r.
In the scope of this paper, we only utilize the 26 factual relations from this dataset. We evaluate LRE in Mamba and Pythia for all the 26 factual relations. We also use this dataset for locating key fact-mediating states in Section 3 and Appendix D. We randomly sample 400 examples (s,r,o) across different factual relations place in city, country capital city, person occupation, plays pro sport, company hg, and product by company. For each of these examples we randomly select another example within the same relation (s", rr,*) such that F s* and + o²_ The average indirect effect (IE) of applying activation patching over these 400 examples is depicted on Figures 2b, 3, 5 for Mamba-2.. 2.8b) and on Figure (for Pythia-2.8b). We the same set of relations in Section 5 where adapt attention knock-out to Mamba.
B Challenges in Performing Attention Knock-out in Mamba
Attention heads in autoregressive transformer LMs and Conv + SSM operations in Mamba play a similar role: bringing bringing/retaining retaining information from the past tokens. Attention "knock- out" is type of causal mediation analysis that tries to understand information flow in transformer LMs by cutting off information propagation from t token to 9thh token position. In transformers, each of the attention heads in an attention module attn(f) calculates an attention matrix I,, where Lygk quantifies how much attention is being paid to the 1tth token by the 4dl token with this specific attention head (see Vaswani et al. (2017) for details on the attention operation). We can block the information flow from t token to 9th token via specific attention head by simply setting Ly oo in the forward pass.
For Mamba, Ali et al. (2024) show that the amount of information retained in the 9d token state s from the convolved state at thh token cf (where k < 4), after the selective-SSM operation (see Equations 4 and 5) can be visualized as an attention matrix per channel.
(f) (() can Since the selective-SSM operation linear, the information retained in s from
www. ww..idata..org/wwkk//kkkk..
14

---

Published as a conference paper at COLM 2024
be calculated accurately as a - input-dependent parameters for the jth token. See Gu & Dao (2023) and Ali et al. (2024) for details on selective-SSM operation. We ask: can we block the information flow from the k token the 4th token in Mamba by subtracting out a a from s If so, attention knockout experiments in Mamba become feasible.
[T A!
By (.) where A, B(a) and c. are
9
We find that blocking information flow via Conv SSM operation through this specific edge from the k token to the qth token can be challenging in Mamba. Note that, since (O) is
convolued state with a receptive field of size in Mamba-2.8b, the states ki k12' and (f)
(e) (E)
(*) from sq
k-3 also retain information from & Which means that even we subtract
qik
(6) stop this leakage, we would want (f)
these states can "leak" information about dx to sq
to subtract from s/ all the information retained from well. However, accurately calculating this is challenging because of the SiLU non-linearity after Conv1D (see Equation (4)).
( via k+11' k-2 and k+3 states as
(E) (f)
In our initial experiments tested subtracting only (2) from sy But found that qi Mamba-2.8b could often refer to the kthh token from the qdl token in copy and factual recall tasks.
C Isolating The Contribution of w
Recall from Figure and Equation (2) that when 0 (9) restored, the s/ and gi are restored as well. To isolate the contribution of only w!O) we subtract out IE (0) + IE from IE 0 and plot the results on Figure 8. Notice that subtracting IE ( cancels out the high indirect effect at the late site shown by later layers the last token position. But, together IE (n) + IE cannot cancel out high IE (0) observed at the early site, that is early-mid layers at the last subject token. This reconfirms the mediating role of w! the early site while recalling fact.
(2)
8,
;; 8)
10 15
11,--(ur e- m)
st be. let"
beton
30
45
Figure& 8: Isolating the contribution of (!! IE 1i + IE in subtracted from IE 0f0) Notice that IE 19 [EE (e) + IE 8)
(c)) still shows higher causal effect at the early site (more pronounced
than IE (e)) while the high causal effect at the late site cancels out. 8)
15

---

Published as a conference paper at COLM 2024
D Locating Key Modules in Pythia-2.8b
(a) Average IE of restoring 8, over 400 prompts
Avemge IE of restoring MLP
(b) Average IE of restoring ATTN
0.8
0.8
pretix
wbject [:23
0.5
0.6
0.6
subject 2nd Jast
eatly site
early sioe
0.4
0.4
subject lase
site
0.0
0.0
15
1E
20
16
15
25
single restored layer within Pythis
Conner gasched layers
patthee layers
Figure 9: Average indirect effect of residual state, MLP, and attention outputs in Pythia-2.8b over 400 facts. For MLP and attention outputs a window of 10 layers around is restored, as restoring just one layer barely shows visible patterns.
with h restored
ATTN, blocked
MLP. blocked
Average iF
0.75
subj last
0.50
0.25
0.00
(b)
prompt last
Average IR
0.75
0.50
0.25
0.00
30
Layers
Figure 10: Impact of ablating ATTN, or MLP; on 10 for (a) subject last and (b) prompt last token positions Pythia-2.8b
LINEARITY OF RELATION EMBEDDING (LRE)
With activation patching we can identify where facts are located inside LM. We are also interested in understanding how LMs extract this information given r), Figures 2b and show a clear separation of roles in early-mid and later layers Mamba. We observe similar phenomenon in autoregressive transformer LMs (Meng et al., 2022a;b; Geva et al., 2023). According to Geva et al. (2023), in transformer LMs, the subject entity representation at the subject last token position, goes through enrichment process, mediated by the MLP in the early-mid layers, where s populated with different dictsaattribute relevant to the subject entity Then, at the last token position, attention modules perform query on the enriched to extract the answer to the prompt x = (s,r). Hernandez et al. (2023) approximate the query operation performed the enriched for specific relation by taking the first order Taylor series approximation (LRE) of the LM computation as
F(s,r) = BJps +
dF
OF b = Esg F(s,r) ds
where E sur
(8)
ds (spr)
(s)r)
ß is scalar, and p is the rank Jf
Hernandez et al. (2023) show that for range of different relations it is possible to achieve LRE that is faithful to the model computation F by averaging the approximations of and calculated on just examples. We utilize LRF to understand the complexity of decoding
16

---

Published as a conference paper at COLM 2024
factual relations in Mamba. We find the hyperparameters p and the layer & (where to extract the enriched s from) using grid search. For mathematical and implementation details, see Hernandez et al. (2023).
(a) LRE faithfulness for factual relations in Mamba-3B
1.0
0.9
0.8
0.7
Faithfulness
0.6
0.5
0.4
0.3
0.2
(b) LRE faithfulness for factual relations in Pythia-3B
1.0
0.9
0.7
Faithfulness
0.6
0.5
0.4
0.2
0.1
0.0
person native language
landmark continent
country language
country capital city
person university
plays pro sport country largest city person lead singer of bund. person occupation person sport position city country
person plays instuument
president birth year
product by company superhero person country currency
food from country
company hq star constellation name
superhero archnemesis
pokemon evolution
person father
landmark country-
person mother
company CEO
Figure 11: Relation-wise LRE faithfulness to the LM decoding function for (a) Mamba-2.8h and (b) Pythia-2.8b. The relations are sorted according to their LRE faithfulness in Mamba- 2.8b. Horizontal red lines per relation indicate random-choice baseline. We only present results for the factual relations in the RELATIONS dataset.
We plot the faithfulness of LRE with и samples on Figure 11 for Mamba-2.8b and similar sized Pythia. The metric faithfulness represents the portion of facts (s,r,0) that can be correctly retrieved if the LM computation F(s,r) is replaced with LRE(S), a simple affine transformation.
We only calculate LRE for the factual relations in the RELATIONS dataset. Figure 11a shows that only for 10 out of 26 factual relations can linear LRE achieve more than 50% faithfulness. For comparison, in the same sized Pythia-2.8b LRE achives > 50% faithfulness for factual relations (see Figure 11b). And, in both Mamba and Pythia, LRE fails to achieve good fuithfulness for the relations where the range (the number of unique answers) is large. These findings align with what Hernandez et al. (2023) observed on GPT and LLaMA models; suggesting that, similar to transformer LMs, factual knowledge might heterogeneously represented for different relations in Mamba.
17

---

Published as a conference paper at COLM 2024
E.1 LRE Performance Across Different Layers
Besides faithfulness Hernandez et al. (2023) introduced another metric causality to measure the performance of LRE. Since LRE is linear function, it is invertible. Assume that for fact (s) LRE can faithfully replace LM computation F(s, Then given the representation o* of another object (o" o) should give As, such that when added to s, := As, the model computation F(s,r) r) should generate o². See Hernandez et al. (2023) for details on
this.
Mamba-2.8b
country cupital city
landmark country
plays pro sport
occupution
1.0
1.0
1.0
0.5
0.5
0.3
0.8
0.6-
0.6-
0.6
0.4
0.4
0.4
0.4
022
02-
0.2
0.2
0.0
0.0
0.0
plays instrument
company CEO
star constellation name
mother
Scone
0.6
0.2
02-
0.2
16
Pythia-2.85
country cupital city
landmark in country
plays pro sport
person occupation
0.2
person plays instrument
company CEO
star constellation name
person mother
1.0
200
200
Faithfulness
Causality
Figure 12: For Mamba, only perform sweep till layer 48, as Figure 5 suggests negligible activity for later layers at the subject last token
18

---

Published as a conference paper at COLM 2024
F Activation Patching results on Mamba-2.8b
rescoting H,"" of single layer
restoring d, of 10 layers around
restoring 8, of 10 layers around
restoring of 10 layers around
0.5
0.5
0.5
[PAD]
0.4
0.4
Jondande
professioaaly
0.3
0.3
playd-
02
0.2
the
0.0
0.5
[PAD]
Harvard/ Ocford
0.4
0.4
0.4
University Univensity
0.3
0.3
0.3
locatel
0.2
0.2
02
0.2
city
0.1
0.1
30 20 20
45
16
0.9
10
20 90
0.0
40
30
15
30
35
50.50
0.0
0.5
[PAD]
0.4
0.4
0.4
0.4
professino
0.3
0.3
0.3
Druce' Harper
02
0.0
10
20 30
50
IE
5 10
15
45
IE
5. 15 20
30
50.55
IE
5
10
15
20
30
55
IE
0.5
0.5
0.5
[PAD]
0.4
0.4
0.4
0.3
0.3
Canada Jgaan
0.2
0.2
0.1
0.1
0.0
0.0
18
10
40
16
510
15
20
35
40
45
55
0.5
0.5
[PAD]
The
0.4
0.4
head
qparter
0.3
0.3
03
Microsofi/
0.2
0.2
0.2
02
Jocoed
0.1
0.1
0.0
10
30
50.55
20
30
50.55
10
15
20
35
40
18
19


Reviews:
[{"review_id": "9jmceOzt8l", "reviewer": "Reviewer_kU61", "summary": "**Paper summary**\n\nThe paper studies mechanistic interpretability of factual knowledge in the context of a pretrained Mamba language model. It borrows heavily from recent techniques that have been explored in the context of Transformers,  applying or adapting interpretability techniques specifically for the Mamba model. Using a few different probes, the paper provides some analysis comparing specificity of factual knowledge location in Mamba with Transformers as well as highlighting some limitations in reusing existing techniques.\n\n**Review Summary**\n\nOverall, I found this to be a useful study that has some flavor of a replication study (reapplying existing techniques and analysis) to a slightly new model architecture (Mamba vs Transformer). I will recommend acceptance of the paper as the work is clearly written and presented and thorough application of interpretability techniques to the Mamba architecture. The main drawback of the paper is that some of the results are of limited use in their current form. Aside from general insights on layer/neuron representation of entities, it is still unclear how one could make use of this analysis in practice.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review. We are pleased to know that you found our work well-motivated and useful. \n\n> “ … What are practical applications of this knowledge”? \n\nAlthough not asked as a question, reviewer `kU61` was skeptical how or if our investigation on LRE (Section 5) and factual information flow (Section 6) have any practical applications. We believe this is an important question and here is our position on this - \n\nAll of these methods, including understanding information flow and LRE, have potential applications in helping debug model behavior. Understanding the role of different components inside the model while performing a task may reveal potential “bugs” in the mechanism; like: Is the mechanism missing something crucial? Or is the model doing something unnecessary/weird and why? … \n\nAs for LRE, [Hernandez et al, 2023](https://arxiv.org/pdf/2308.09124) present a tool called *attribute lens* based on it, which can be directly applied to the latent of different token-layer positions in a free-form text to visualize the model’s current belief. If a model utters an incorrect fact for which the relation is *linear*, *attribute lens* could be used to visualize whether information retrieved early in the model is incorrect (suggesting the LM might not have learned the association from training data), or if the mistake is happening later (suggesting the model is making a mistake due to something else in the context). However, *attribute lens* will fail to elicit information faithfully for non-linear relations.\n\nWe recognize that full end-to-end debugging will eventually require a better understanding of processing of information at every stage, not just factual association. But partial insights are already helpful in building intuition and they may help direct future research.\n\n\n> \"I could not find if there is a specific dataset for Section 3 ...\"\n\nThanks for bringing this to our attention. Due to space limitations in the main paper we had to move the dataset section in the appendix (Appendix A). The \"400 facts\" mentioned in the results of activation patching experiments (Figures 2b, 3, and 5 in Section 3) were sampled from 6 different relations in the Relations dataset (Hernandez et al, 2023) - *place in city*, *country capital city*, *person occupation*, *plays pro sport*, *company hq*, and *product by company*. We will update the figure captions and text in relevant sections to refer to Appendix A in the next revision to clarify this.", "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": "- The paper is well written and clearly presented. The background on Mamba and the studied techniques are useful.\n\n- The investigation of interpretability methods to novel LM architectures is a useful type of study reproduction. As a field, we should invest in novel architectures and make sure our techniques can be robust to architectural changes.\n\n- The adaptation of the activation patching to Mamba, allowing several techniques to be adapted from their Transformer counterparts, is well motivated and executed.\n\n- The analysis of the results, while not unexpected following prior results on Transformers, is clear and uses sound experiments.", "reasons_to_reject": "- While some of the analysis presented is interesting, it is unclear how to make use of some of the observations in practice. One example is the results from the LRE (Linearity of relation embedding) section. The faithfulness scores of the LRE are pretty low making it hard to claim that these representations truly capture internal representation of what goes inside the neural network.\n\n- Similarly for the attention knock-out investigation. Suppose we can more clearly understand information flow inside of Mamba (or Transformer); what are practical applications of this knowledge? Do we gain better insights into training new models? Fixing factuality errors? Better understand model confidence about its own knowledge?", "questions_to_authors": "- I could not find if there is a specific dataset for Section 3. It mentions that there are “400 facts” analyzed, but I could not find where these facts are derived from.", "extra_scores": {}}, {"review_id": "Qw07n7D4yz", "reviewer": "Reviewer_o7rW", "summary": "The work adapts a series of information tracing and editing techniques originally developed for transformers to the Mambda state-space model. The general finding is that facts are stored (and can be edited) in a similar way in Mamba as they are in transformers (the representatitve transformer model here beign Pythia).\n\nThis is a nice exercise in transfering knowledge across model architectures. The authors have considered a wide range of information tracing methods and a well-established editing method - ROME. The work seems solid and could become quite impactful if state-space models take off.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review. We are glad that you found our work well motivated.\n\n> \"... how would the approaches proposed here generalize to other state-space models?\"\n\nIn this paper we investigate two primary questions - \n\n1.  Can the analytical tools and methods developed to understand factual recall mechanisms in transformer language models be applied across different language modeling architectures?\n2. To what extent do our insights on factual recall mechanisms in transformer LMs (at a high level) generalize across different types of LMs? \n\nWe limit our experiments on Mamba-2.8b, the best performing language model among all the linear-time RNN approaches; including SSMs, RWKV, and other architectures ([Gu and Dao, 2023](https://arxiv.org/abs/2312.00752)). Our experiments yield overall positive results on both questions. This gives us reason to believe that the approaches and findings in this paper are not specific to any language modeling architecture, but will (to some extent) generalize to a range of other SSM architectures; the main assumption is that it is an **autoregressive** language model.\n\n\nAs we discussed in Section 6 and Appendix D, replicating path-dependent attention knock-out experiments in Mamba pose specific challenges due to the use of convolutions and a non-linearity  in conjunction with pure SSM state. Other SSM architectures without this complexity may allow similar experiments as well. \n\n\nThank you for your suggestion. We will clarify our stance in the Discussion section in the next revision.\n\n\n> Typos\n\nThanks for the deep read and catching the typos. We will address them in the revised draft.", "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": "Clear motivation and technical contribution.", "reasons_to_reject": "The authors should include a real discussion (rather than what is now called Discussion but is really a Conclusion) and address the question: how would the approaches proposed here generalize to other state-space models? I.e. how much can a state-space model deviate from the exact Mamba architecture and still be analyzed or manipulated with the methods presented here?", "questions_to_authors": "Typos:\n- \"In this we are able\" \n- \"the evolution of hidden state\" \n- Figure 5 is missing an x-axis label", "extra_scores": {}}, {"review_id": "ujibYFaZ0U", "reviewer": "Reviewer_SQMx", "summary": "This work tries to analyze the mechanism of factual recall in Mamba following the pattern that is previously employed to analyze Transformer-based language models. Specifically, this work reveals that the activation patching technique can be used to specify model components that relates to outputs the correct objects in factual recall. However, applying rank one model editing (ROME), linearity of relation embedding, and knocking out attention in Mamba demonstrate different behavior.", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive review. We are glad to hear that you found our work valuable.\n\n\n> \"I am a bit confused about what narrative this work tries to build ... \"\n\nWe are actually interested in a broader question confronting the field of interpretability - *to what extent do the analytical tools and our insights on certain mechanisms developed for one architecture (such as Transformers) generalize to a very different architecture?* This work is a case study, exploring this question by trying to understand the factual recall mechanism in Mamba with methods that have been effective in understanding how Transformer LMs do so. Our investigation reveals that these methods can often be applied directly or with some adaptations to account for architectural differences. Additionally, we find that, despite significant architectural differences, Mamba shares many high-level similarities with transformer LMs when it comes to factual recall.\n\nWe will update the abstract and relevant sections in the introduction to clarify our motivation and the narrative of this paper.\n\n> \"The term 'early sites' and 'late sites' lack a clear definition throughout the work.\"\n\nWe will annotate Figures 2, 3, and 9 in the revised draft to clarify what we refer to as “early site” and “late site”. Thanks for bringing this to our attention.\n\n\n> \"The claim 'transformer MLP outputs have effects in early site and not the late site' needs reference or connection to other parts to be better understood.\"\n\nThank you for catching this. This claim was made comparing Figure 3 (Mamba 2.8b) in Section 3 with Figure 9 (Pythia 2-8b) in the Appendix. Causal mediation analysis (Figure 9) shows that, when Transformer LMs recall a fact, MLP contributions are pronounced at the early site (subject last token - early to mid layers), while the high causal effect at the late site (prompt last token - later layers) is due to the Attention mechanism. We will revise the text to clarify this.\n\n\nLet us know if you have further concerns or feedback to improve our paper. And kindly consider adjusting your score upwards if you are satisfied with our response.", "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": "This work adds value in aligning the analytical tools in Transformer-based architectures to non-transformer architectures like Mamba.", "reasons_to_reject": "I am a bit confused about what narrative this work tries to build: the previous framework in analyzing how factual recall happens in transformer-based language models partially works in Mamba, or how factual recall happens in Mamba based on some existing works that are done in transformer-based language models? So far it reads like the former, which makes different parts not connected strongly under the architecture of Mamba.", "questions_to_authors": "- The term 'early sites' and 'late sites' lack a clear definition throughout the work.\n- The claim 'transformer MLP outputs have effects in early site and not the late site' needs reference or connection to other parts to be better understood.", "extra_scores": {}}, {"review_id": "zsklgWsbgG", "reviewer": "Reviewer_rwMm", "summary": "This paper presents a thorough investigation of factual knowledge in Mamba. The authors adapt and apply several techniques used in Transformer LMs to investigate where and how Mamba encodes facts. Specifically, the authors apply (i) activation patching to locate facts, (ii) ROME to edit facts, and (iii) LRE to understand how Mamba extracts factual knowledge. In doing this, the authors find many similarities between Transformer LMs and Mamba for factual knowledge, as well as some differences and challenges, such as blocking the information from between two tokens due to the receptive field of the convolutional layer.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review and your suggestions on how to improve our paper. We will try our best to incorporate your feedback in the next revision.\n\n> \"I found it hard to follow the experiments on “path-specific effects” (page 5). This subsection could be improved ... \"\n\nWe will update parts of section 3 to make out path-specific experiment setup easier to understand. Thank you for your suggestions.\n\n\n> \"Which dataset did you use in Sections 3 and 6?\"\n\nThanks for bringing this to our attention! We had to move the dataset section in Appendix A due to the page limit on the main paper. The activation patching results (Figures 2b, 3, and 5 in section 3; also Figs 8, 9 and 10 in the appendix) are averaged over 400 facts sampled from 6 different relations in the Relations dataset ([Hernandez et al, 2023](https://arxiv.org/pdf/2308.09124)) - *place in city*, *country capital city*, *person occupation*, *plays pro sport*, *company hq*, and *product by company*. \n\nThe experiment in Section 6 (continued on Appendix D) was conducted on 700 examples randomly sampled from the same set of relations in the Relations dataset. We will update the text in relevant sections and the figure captions to clarify this.", "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": "1.  A thorough investigation of factual knowledge in Mamba models.\n2.  The paper provides insights across different axes, from where facts are stored, to how they can be encoded, as well as how to edit them.\n3. The paper is generally well written and organized.", "reasons_to_reject": "I do not have any strong reasons to reject the paper. That said, some parts could be improved:\n1. I found it hard to follow the experiments on “path-specific effects” (page 5). This subsection could be improved perhaps by bringing Figure 5 closer to the text and making more explicit connections so the reader can better visualize the process.\n2. Present a succinct table or text that compares Mamba with Transformer LMs so that the reader can quickly glance at your key results before delving into the corresponding sections.", "questions_to_authors": "1. Which dataset did you use in Sections 3 and 6?", "extra_scores": {}}]

Author rebuttal:

Meta-review: The submission applies several recent mechanistic interpretability techniques to Mamba with the key question being *To what extent do those technique generalize between different architectures?*. The reviewers agree that while the submission is not introducing a new method, it is a useful and timely study in interpretability of SSMs. It is recommended that the authors rework the discussion section so that it becomes less of a summary but spends more time on what classes of interpretability and editing methods seem most promising for Mamba and related models. 

The authors promise to open-source the code which I believe would be quite useful for other researchers.
Decision: Accept