Locating and Editing Factual Associations in Mamba

Authors: arnab sen sharma
david  atkinson
david  bau
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 19:02:17

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.

Main paper: 
 introduction: 
 Studies of autoregressive transformer language models' (LMs) processing of factual statements 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 a 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-such as recurrent networks (RNNs)-has not yet been investigated.In this paper we investigate the internal mechanisms of Mamba (Gu & Dao, 2023) , a recently-proposed state-space language model, a 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 a 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 a 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-knockout) 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 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 a 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 to Mamba. While studies of transformers (Meng et al., 2022a; b; Hase et al., 2024) have found that there are a 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 to 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 on 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 et 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 a 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 MambaBlock1 , which is stacked homogeneously, replacing both attention and MLP blocks used in transformer layers. Here, we focus on the different operations performed inside a MambaBlock. Formally, Mamba is an autoregressive language model: M : X → Y over a vocabulary V that maps a sequence of tokens after the ℓ th (1-indexed) MambaBlock is computed as follows:Conv + SSM σ W a ℓ W g ℓ W o ℓ h i ℓ h i ℓ-1 a i-1 ℓ a i ℓ s i ℓ g i ℓ o i ℓ Figure 1:x = [x 1 , x 2 , . . . , x T ] ∈ X , x i ∈ V to y ∈ Y ⊂ R |V |h (ℓ) i = h (ℓ-1) i + o (ℓ) i (1)where o o(ℓ) i = MambaBlock (ℓ) h (ℓ-1) 1 , h (ℓ-1) 2 , . . . , h (ℓ-1) i = W (ℓ) o s (ℓ) i ⊗ g (ℓ) i (2)Here, ⊗ represents element-wise multiplication or Hadamard product. s(ℓ) i is calculated as: a (ℓ) i = W (ℓ) a h (ℓ) i (3) c (ℓ) 1 , c (ℓ) 2 , . . . , c (ℓ) i = SiLU Conv1D a (ℓ) 1 , a (ℓ) 2 , . . . , a (ℓ) i (4) s (ℓ) i = selective-SSM c (ℓ) 1 , c (ℓ) 2 , . . . , c (ℓ) i (5)We abstract the operations in Equations 4 and 5 as the Conv + SSM operation in Figure 1 . At a 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 g (ℓ)i (that does not pass through Conv + SSM operation) is a 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.g (ℓ) i = SiLU W (ℓ) g h (ℓ-1) i (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 s i , g i , and o i at all layers. In Section 4, following Meng et al. (2022a) ; Hase et al. (2024) , we test rank-one edits of facts across components W a , W g , and W o at each layer. And in Section 5 we address the challenge of applying attention patching in Mamba, as used in Geva et al. (2023) to isolate information flow in GPT LMs. We include a further investigation of Mamba's factual recall mechanisms in Appendix E, following Hernandez et al. (2023) in collecting Jacobians within Mamba to test the linearity of its relational encodings. 
locating key states for factual recall: 
 We begin with activation patching, seeking to understand if there are specific hidden states which play important roles during factual recall. We select a fact (s, r, o) that the LM knows, where r is a relation that associates a subject entity s with an object entity o. To estimate each state's contribution towards a correct factual prediction (s = Michael Jordan, r = professionally played, o = basketball), we collect model activations across three different runs: i , and s (ℓ)i over 400 facts from the RELA-TIONS dataset (see Appendix A.2). For each layer ℓ, states for a window of 10 layers around ℓ are restored from the clean run G.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, r) = Michael Jordan professionally played. We cache all the hidden states during the clean run to be used later:h (ℓ) i , a (ℓ) i , s (ℓ) i , g (ℓ) i | i ∈ [1, T], ℓ ∈ [1, L] .corrupted run G * : In the corrupted run, we swap s with a different subject s * (Pelé) such that the LM gives a different answer o * (soccer) to the modified prompt x * = (s * , r) (i.e., o * ̸ = o).This subject-swapping approach follows the recommendation of Zhang & Nanda (2023) and has the 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 et al. (2022a)) .patched run G * [← h (ℓ)i ]: In the patched run, we run the model on the corrupted prompt x * , but intervene on h (ℓ) i 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 can have a downstream effect of potentially changing all the states that depend on it. See Figure 2a . i in recalling the fact (s, r, o), we define its indirect effect (IE) as:EQUATIONIn Figure 2b we plot the average indirect effect of restoring the residual states h(ℓ)i 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 a clean h (ℓ) i there will restore most of the model computation from G. However, Mamba also shows high causality at the early site (early-middle layers at 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 i (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 outputs s (ℓ) i have high IE only at the late site, resembling the behavior of attention modules in GPT models (Meng et al., 2022a) . However, there is no 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 g (ℓ)i 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 MLP outputs have effects in the early site and not the late site, in Mamba there is 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?G * G * Blocking [←h ] ℓ i h ℓ i g ℓ i s ℓ i h ℓ i from G s i o ℓ i (a) (b)Figure 4 : To probe for path-specific effects, (a) h (ℓ) i is restored from the clean run G (as in Figure 2a ). (b) Then, to reveal the role of the Conv + SSM contributions, s * i states from the corrupted run G * are also patched to block the contributions from s i paths.To investigate this question, we replicate an experiment from Meng et al. (2022a) to probe path-specific effects (Pearl, 2022) by severing a path from the causal graph and monitoring its effect. Here, we are interested in understanding the effect of the contributions from g i , s i , and o i (i.e. states that are processed by W g , Conv + SSM, and W o respectively) while recalling a fact. First, in the corrupted run G * , at token position i, we cache all the contributions from the s i paths ass * i = {s * (ℓ) i | ℓ ∈ [1, L]}. Then in the patched run G * [← h (ℓ) i ], we restore h (ℓ)i that was cached from the clean run G into its corresponding state (as in Figure 2a ), but with an additional modification: to understand the contribution from the s i paths, we sever those paths by also patching s * i (cached from the corrupted run G * ) to their corresponding locations (see Figure 4 ). The same experiment is replicated to understand the contributions of g i and o i states. We note that severing the o (ℓ) i will sever s (ℓ) i and g (ℓ)i as well (see Figure 1 ).In Figure 5 we plot the average results of this experiment for token positions (a) i = subject last and (b) i = prompt last over 400 examples randomly sampled from the RELATIONS dataset. The key findings can be 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, W g , or W o parameters, respectively. At the early site at the subject last token, both the Conv + SSM and W g have a strong role, but W o plays an even larger role than either. Yet the strongest mediator at the late site is also W o . This experiment highlights the importance of W o in both stages of predicting a fact. But it also suggests that Mamba does not separate early-site factual recall between these groups of parameters as 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 keysK = [k 1 |k 2 | . . . ] to their corresponding values V = [v 1 |v 2 | . . . ],and uses this to edit factual associations in transformer LMs. Here, we apply the technique to a particular set of linear transformations within Mamba, and report our editing success on each. 2The input to ROME is a prompt x = (s, r), where s (Emmanuel Macron) is a subject entity and r (is the President of ) is a relation. ROME also takes a counterfactual object o * (England), meant to replace the correct object o (France) in the model's output. To effect that change, ROME generates a rank-one update to W (ℓ) down , the down-projection matrix of the MLP module for the last token of the subject at layer ℓ-which plays the role of the associative memory. In generating the rank-one update, ROME considers the input to W We apply ROME on the three different projection matrices of Mamba: W o , 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 (ℓ) a have poor generalization (PS) in early layers, whereas high PS can be achieved at early layers by modifying either W (ℓ)g or W (ℓ)o , consistent with their higher indirect effects as seen in Figure 5a .Where is the right place to apply ROME on Mamba? Figure 3 could suggest W (ℓ) g , since the causal effect of g i states is mostly concentrated at the subject last token, similar to the behavior of MLPs in transformers (Meng et al., 2022a) . Consistent with this is the architectural fact that, just as transformers' W (ℓ) down connects to attention modules only through the residual stream, the output of W (ℓ) g 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 in transformers (Grazzi et al., 2024) . And, indeed, we find that ROME can successfully insert facts by modifying W (ℓ) g . On the other hand Figure 6a reveals sudden drops in efficacy and generalization at middle layer gates, suggesting that W (ℓ) g may be an unreliable mediator at some layers. Our experiments further show that the best performance for ROME is empirically achieved by modifying W (ℓ) o . This is consistent with the fact that o i states show a stronger causal effect at the subject last token than g i states do (see Figures 3a and 5a ). Additionally, ROME achieves better generalization (PS), competitive specificity (NS), and an overall better score (S) with W (ℓ) o . We hypothesize that the strong performance of W (ℓ) o may be due to the the separation of roles between early-mid and later layers observed in Figures 2b, 3a, and 5 . Also see Appendix C where we isolate the contribution of W k to future tokens via s i paths is blocked, with k taking the value of either subject, non-subject, or the prompt_last token positions. For each layer ℓ, s i paths were blocked for a window of 10 layers around ℓ.s (ℓ) i + IE g (ℓ) i from IE o (ℓ) i , which reveal a critical role of W (ℓ)o in early-mid layers at subject last token position while mediating a fact.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 in transformer LMs. In attention "knock-out" experiments the information that flows through a specific edge (from k th token to q th token) via a certain attention head is blocked to understand if critical information flows through that edge. This is also a form of causal mediation analysis and it has been effective in understanding the information flow in 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 i states, with the Conv + SSM operations (see Figure 1 and Equations 3-5). We ask, can we perform experiments similar to attention knock-out experiments in Mamba in order to understand how it moves factual information?We find that performing similar experiments in Mamba can be difficult. The use of Conv with a non-linearity in conjunction with selective-SSM make it challenging to remove the information retained in the q th token from the k th token (see Appendix B for details). However, it is possible to block the propagation of information from the k th token to all the future tokens via Conv + SSM with mean-ablation. Specifically, for a layer ℓ, we set a ℓ) , where E a (ℓ) 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 a 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 GPT LMs.(ℓ) k := E a (We randomly sample 700 facts across 6 factual relations from the RELATIONS dataset. For each of those examples we block-out information propagation of the subject, non-subject, and the prompt-last token positions for a window of 10 layers around a specific layer ℓ. The effect of blocking out Conv + SSM information flow for certain layer-token (ℓk) positions is measured as the relative change in p(o) with p o | a (ℓ) k :=E a (ℓ) -p(o) /p(o). Figure 7 shows the averaged result and it leads us to draw the following conclusions about how factual information flows in Mamba:(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 information from all the subject tokens in 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 s i 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 can 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 i and s i -states that we hypothesize take crucial part in the enrichment process-also 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 k th token and the last token are affected as well. 
related works: 
 Mamba. Mamba is a 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. Locating Factual Knowledge in Language Models. To make factually correct statements about the world, a LM has to store factual knowledge about real world entities somewhere in its parameters. Understanding how and where a neural network stores knowledge is a core problem for interpretability and it has thus been studied from several perspectives (Ji et al., 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 of these classifiers can 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. 
discussion: 
 In this paper we have set out to understand whether the analytical methods and tools developed for transformer LMs can also be applied on the Mamba LMs based on recurrent state-space architecture. Although our experiments have been limited to Mamba-2.8b, the largest available LM of that family, and comparisons to 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 to 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 o weights can receive rank-one model editing (ROME) edits with good generalization and specificity at a range of layers, similar to W down 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 in transformers and Mamba leads us 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 a 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 of 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.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 Wikidata3 . The dataset contains 21, 919 requests {s, r, o, o * , π * } where o is the correct answer to the prompt x = (s, r), o * is the counterfactual edit request, and π * ∼ P (s, r) is a paraphrase of the prompt x = (s, r) to test for generalizability (PS).Each of the records also contain some neighborhood prompts π N to test for specificity (NS) and some generation prompts π G 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.8b (Figure 6a ) and Pythia-2.8b (Figure 6b ) on the first 2000 records from COUNTERFACT.The RELATIONS dataset introduced in Hernandez et al. ( 2023) consists of 47 relations of 4 types: factual, linguistic, bias, and commonsense. A relation r is an association between two entities. For example, the relation, r = professionally played the sport connects the subject s = Michael Jordan with the object o = basketball. The dataset contains a set of (s, o) 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 6 different factual relations -place in city, country capital city, person occupation, plays pro sport, company hq, and product by company. For each of these examples we randomly select another example within the same relation (s * , r, o * ) such that s ̸ = s * and o ̸ = o * . The average indirect effect (IE) of applying activation patching over these 400 examples is depicted on Figures 2b, 3 , 5 for Mamba-2.8b) and on Figure 9 (for Pythia-2.8b). We use the same set of 6 relations in Section 5 where we adapt attention knock-out to Mamba.Attention heads in autoregressive transformer LMs and Conv + SSM operations in Mamba play a similar role: bringing/retaining information from the past tokens. Attention "knockout" is a type of causal mediation analysis that tries to understand information flow in transformer LMs by cutting off information propagation from k th token to q th token position.In transformers, each of the attention heads in an attention module attn (ℓ) calculates an attention matrix L, where L q,k quantifies how much attention is being paid to the k th token by the q th token with this specific attention head (see Vaswani et al. (2017) for details on the attention operation). We can block the information flow from k th token to q th token via a specific attention head by simply setting L q,k := -∞ in the forward pass.For Mamba, Ali et al. (2024) show that the amount of information retained in the q th token state s(ℓ)q , from the convolved state at k th token c (ℓ) k (where k < q), after the selective-SSM operation (see Equations 4 and 5) can be visualized as an attention matrix per channel.Since the selective-SSM operation is linear, the information retained in s (ℓ)EQUATIONi , and C(ℓ)i are input-dependent parameters for the i th 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 th token to the q th token in Mamba by subtracting out α(ℓ) q,k from s (ℓ)q ? If so, attention knockout experiments in Mamba become feasible.We find that blocking information flow via Conv + SSM operation through this specific edge from the k th token to the q th token can be challenging in Mamba. Note that, since c k . Which means that even if we subtract α(ℓ) q,k from s(ℓ)q , these states can "leak" information about a (ℓ)k to s (ℓ)q . To stop this leakage, we would want to subtract from s (ℓ) q all the information retained from a (ℓ)k via c (ℓ) k+1 , c (ℓ)k+2 , and c (ℓ) k+3 states as well. However, accurately calculating this is challenging because of the SiLU non-linearity after Conv1D (see Equation ( 4)).In our initial experiments we tested subtracting only α(ℓ) q,k from s (ℓ)q . But we found that Mamba-2.8b could often refer to the k th token from the q th token in copy and factual recall tasks.Recall from Figure 1 and Equation ( 2) that when o (ℓ) i + IE g (ℓ) i subtracted from IE o (ℓ) i . Notice that IE o (ℓ) i -IE s (ℓ) i + IE g (ℓ) istill shows higher causal effect at the early site more pronounced than IE With activation patching we can identify where facts are located inside a LM. We are also interested in understanding how LMs extract this information given x = (s, r). Figures 2b and 5 show a clear separation of roles in early-mid and later layers in Mamba. We observe a 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 s, at the subject last token position, goes through an enrichment process, mediated by the MLP in the early-mid layers, where s is populated with different facts/attributes relevant to the subject entity s. Then, at the last token position, attention modules perform a query on the enriched s to extract the answer to the prompt x = (s, r). Hernandez et al. (2023) approximate the query operation performed on the enriched s for a specific relation r by taking the first order Taylor series approximation (LRE) of the LM computation F asEQUATIONβ is a scalar , and ρ is the rank of J Hernandez et al. (2023) show that for a range of different relations it is possible to achieve a LRE that is faithful to the model computation F by averaging the approximations of J and b calculated on just n = 5 examples. We utilize LRE to understand the complexity of decoding factual relations in Mamba. We find the hyperparameters β, ρ and the layer ℓ (where to extract the enriched s from) using grid search. For mathematical and implementation details, see Hernandez et al. (2023) . In their paper,Gu & Dao (2023) call this component Mamba-the same name as the LM family.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).www.wikidata.org/wiki/Wikidata:Main_Page 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_summary: **Paper summary**

The 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.

**Review Summary**

Overall, 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.
strength_weakness: strength_weakness: - 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.

- 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: questions: - 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.
paper_summary: paper_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).

This 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.
strength_weakness: strength_weakness: 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: questions: Typos:
- "In this we are able" 
- "the evolution of hidden state" 
- Figure 5 is missing an x-axis label
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: - The term 'early sites' and 'late sites' lack a clear definition throughout the work.
- 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: I do not have any strong reasons to reject the paper. That said, some parts could be improved:
1. 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.
2. 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: questions: 1. Which dataset did you use in Sections 3 and 6?

Comments: title: Acknowledgement of rebuttal
comment: Dear authors, thank you for clarifying my questions. I hope you will add the missing details in the extra page provided for accepted papers.
Thanks for the reply from the authors and more information about the questions. Since the theme of the work is to investigate how much the current analysis tools can be extended to non-Transformer architecture and the results shows a partial matching of the transfer, the work is better to include more high-level guideline about the failure part and how many properties are architecture-agnostic (as well as agnostic to scaling model parameters) a which can be considered for future novel architectures. So far, the work does not seem to bring a sufficient answer to this question. I would rather remain my score unchanged.
We are grateful to our reviewers for their thoughtful and constructive feedback. We plan to revise our paper to incorporate their suggestions, which we believe will strengthen it. Here is a summary of changes we intend to make:

* [`kU61`, `rwMm`] Revise the text and captions of figures in Sections 3 and 6 to clarify the dataset used to perform those analyses.
* [`o7rW`] In the Discussion section, include our opinion on to what extent the approaches explored in this paper are generalizable across different SSM architectures.
* [`kU61`, `o7rW`] Revise the text in the Abstract, Introduction, and relevant sections to clarify the narrative and motivation of our work.
* [`SQMx`] Update the text and annotate Figures 2, 3, and 9 in our paper to clarify “early site” and “late site”.
* [`rwMm`] Revise the description of path-dependent experiment setup in Section 3 to make it easier to understand.


We have responded to each reviewers questions/concerns in separate responses. We invite the reviewers to share additional comments, questions, or suggestions to improve our paper. And kindly consider adjusting scores upwards if your concerns have been addressed.


Thanks!
Dear Reviewer rwMm,

Thank you for your suggestions on how to improve our paper. As we have a short time left in the discussion period, a gentle nudge to please let us know if you are satisfied with our responses. Kindly let us know if you have any further feedback.

Thank you!
Dear Reviewer SQMx,

We have a short time left in the discussion period, so a gentle nudge to please let us know if you are satisfied with our responses. As a reminder, we plan to revise the paper in following ways to incorporate your feedback:

1. Clarify the narrative and motivation in the Abstract And Introduction.
2. "early site" refers to the prominent indirect effects noticed at the subject last token in early to mid layers. And "late site" refers to the high causal effect at the prompt last token in later layers. We plan to annotate Figures 2, 3, and 9; also revise the text in relevant sections to clarify this.

We invite you to share any additional feedback and help us improve our paper. And kindly consider increasing your score if your concerns have been addressed.

Thank you!
Dear Reviewer o7rW,

Thank you again for your suggestions on how to improve the Discussion section in our paper and for catching the typos. We have a short time left in the discussion period, so a gentle nudge to please let us know if we are satisfied with our responses. Kindly let us know if you have any additional feedback.

Thank you!
Dear Reviewer kU61,

We have a short time left in the discussion period, so a gentle nudge to please let us know if we are satisfied with our responses to your concerns. Kindly let us know if you have any further feedback or suggestions to improve our paper.

Thank you!