Locating and Editing Factual Associations in Mamba

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. Third, we examine the linearity of Mamba's representations of factual relations. 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 architectural approach, when it comes to factual recall, the two architectures share many similarities.

Paper

Similar papers

Reviewer rwMm7/10 · confidence 4/52024-05-11

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.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

1. A thorough investigation of factual knowledge in Mamba models. 2. 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. 3. 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: 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 to authors

1. Which dataset did you use in Sections 3 and 6?

Authorsrebuttal2024-06-02

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!

Reviewer SQMx5/10 · confidence 3/52024-05-11

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.

Rating

5

Confidence

3

Ethics flag

1

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

Authorsrebuttal2024-06-02

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!

Reviewer o7rW7/10 · confidence 3/52024-05-13

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.

Rating

7

Confidence

3

Ethics flag

1

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: - "In this we are able" - "the evolution of hidden state" - Figure 5 is missing an x-axis label

Authorsrebuttal2024-06-02

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!

Reviewer kU617/10 · confidence 3/52024-05-16

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.

Rating

7

Confidence

3

Ethics flag

1

Reasons to accept

- The paper is well written and clearly presented. The background on Mamba and the studied techniques are useful. - 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. - The adaptation of the activation patching to Mamba, allowing several techniques to be adapted from their Transformer counterparts, is well motivated and executed. - 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. - 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.

Authorsrebuttal2024-06-02

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!

Authorsrebuttal2024-06-05

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!

Reviewer SQMx2024-06-06

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.

Reviewer rwMm2024-06-06

Acknowledgement of rebuttal

Dear authors, thank you for clarifying my questions. I hope you will add the missing details in the extra page provided for accepted papers.

Program Chairsdecision2024-07-10

Decision

Accept

© 2026 NYSGPT2525 LLC