Locating and Editing Factual Associations in Mamba

TL;DR

This paper explores how factual knowledge is stored and retrieved in the Mamba model, revealing localization of recall mechanisms, successful fact editing, and similarities to transformer LMs in factual recall processes.

Contribution

It demonstrates that Mamba's factual recall mechanisms are localized, editable, and comparable to transformer models, providing new insights into its internal processes.

Findings

Localization of factual recall to specific Mamba components

Successful insertion of facts via model editing methods

Similarity in factual recall mechanisms between Mamba and transformer LMs

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