Elamur Research⌕
☰
home / papers / 2024 / Locating and Editing Factual Associations in Mamba

Locating and Editing Factual Associations in Mamba

Paper · 2024 · Arnab Sen Sharma, David Atkinson and David Bau

Summary

Studies of transformer LMs have found that factual recall is concentrated in particular modules at specific token positions, and that one rank-one weight change can overwrite a stored fact. The authors ask whether this also holds for Mamba, a recurrent state-space LM built from convolutions, gates and selective-SSM modules instead of attention and MLP blocks. They study Mamba-2.8b alongside the similarly sized transformer Pythia-2.8b. They adapt four transformer techniques: activation patching (causal tracing), ROME-style rank-one editing of MambaBlock projections, Jacobian-based tests of whether relations are linear (LRE), and a replacement for attention knockout that blocks the information a token sends forward.

Mamba behaves much like the transformer. Middle layers have strong causal effects at the subject's final token, and later layers matter at the prompt's final token. Edits to the output projection W_o generalize to paraphrases and leave unrelated facts intact across a range of layers. Many relations are well approximated linearly. The authors suggest that most interpretability tools carry over to new architectures, and that localized recall may come from autoregressive modeling itself rather than from attention.

Key points

  • Models: Mamba-2.8b, the largest Mamba LM available at the time, compared with Pythia-2.8b. Data: the RELATIONS dataset (causal tracing averaged over 400 facts) and COUNTERFACT.
  • Corruption is done by swapping the subject (e.g., Michael Jordan → Pelé), following Zhang & Nanda (2023). This avoids the Gaussian-noise corruption used by Meng et al. (2022a).
  • Patching targets the residual states and the internal states s, g and o. Rank-one edits were tested on W_a, W_g and W_o. W_o behaves like W_down in Pythia and GPT models.
  • Some methods transfer only partly. Path-dependent attention patching has no direct Mamba equivalent, so information flow was studied by blocking a token's outgoing information.
  • Code is published at romba.baulab.info. Experiments ran on A100 (80GB) and A6000 (48GB) GPUs using HuggingFace Transformers and PyTorch.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Northeastern UniversityBoston · United States57 works · 6,533 citations🎓Northeastern UniversityShenyang · China42 works · 1,940 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

Cite

Cite as Arnab Sen Sharma, David Atkinson and David Bau. 2024. Locating and Editing Factual Associations in Mamba. arXiv:2404.03646.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2404.03646 · arXiv 2404.03646

@misc{sharma2024locating,
  title = {{Locating and Editing Factual Associations in Mamba}},
  author = {Arnab Sen Sharma, David Atkinson and David Bau},
  year = {2024},
  doi = {10.48550/arxiv.2404.03646},
  eprint = {2404.03646},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2404.03646}
}

References

41 references, 12 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint arXiv:2403.01590, 2024.
  2. James A Anderson. A simple neural network generating an interactive memory. Mathematical biosciences, 14(3-4):197–220, 1972. find
  3. Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48(1):207–219, 2022.
  4. Yonatan Belinkov and James Glass. Analysis methods in neural language processing: A survey. Transactions of the Association for Computational Linguistics, 7:49–72, 2019.
  5. Yonatan Belinkov, Nadir Durrani, Fahim Dalvi, Hassan Sajjad, and James Glass. What do neural machine translation models learn about morphology? arXiv preprint arXiv:1704.03471, 2017.
  6. Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on M… find
  7. Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv:1409.1259, 2014.
  8. Alexis Conneau, German Kruszewski, Guillaume Lample, Loïc Barrault, and Marco Baroni. What you can cram into a single vector: Probing sentence embeddings for linguistic properties. arXiv preprint arXiv:1805.01070, 2018.
  9. James Durbin and Siem Jan Koopman. Time series analysis by state space methods, volume 38. OUP Oxford, 2012. find
  10. Yanai Elazar, Nora Kassner, Shauli Ravfogel, Abhilasha Ravichander, Eduard Hovy, Hinrich Schütze, and Yoav Goldberg. Measuring and improving consistency in pretrained language models. Transactions of the Association for Computational Linguistics, 9:1012–1031, 2021. find
  11. 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 on evaluating vector-space representations for nlp, pp. 134–139, 2016.
  12. 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 to foundation model internals. arXiv preprint arXiv:2407.14561, 2024. find
  13. 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, abs/2112.00826, 2021. URL https://arxiv.org/abs/2112.00826. ↗
  14. Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models. arXiv preprint arXiv:2304.14767, 2023. find
  15. Karan Goel, Albert Gu, Chris Donahue, and Christopher Re. It’s raw! Audio generation with state-space models. In 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 Pro… find
  16. Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is Mamba Capable of In-Context Learning?, 2024. URL http://arxiv.org/abs/2402.03170. — [arXiv:2402.03170](https://arxiv.org/abs/2402.03170) ↗
  17. Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. find
  18. Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021. find
  19. 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. find
  20. 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. arXiv preprint arXiv:2308.09124, 2023. find
  21. Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735–1780, 1997. find
  22. 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. find
  23. Shaoxiong Ji, Shirui Pan, Erik Cambria, Pekka Marttinen, and S Yu Philip. A survey on knowledge graphs: Representation, acquisition, and applications. IEEE Transactions on Neural Networks and Learning Systems, 33(2):494–514, 2021. find
  24. Teuvo Kohonen. Correlation matrix memories. IEEE transactions on computers, 100(4):353–359, 1972. find
  25. Siem Jan Koopman, Neil Shephard, and Jurgen A Doornik. Statistical algorithms for models in state space using ssfpack 2.2. The Econometrics Journal, 2(1):107–160, 1999. find
  26. 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. find
  27. Kevin Meng, Arnab Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. Mass-editing memory in a transformer. arXiv preprint arXiv:2210.07229, 2022b. find
  28. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843, 2016.
  29. Neel Nanda, Senthooran Rajamanoharan, János Kramár, and Rohin Shah. Fact finding: Attempting to reverse-engineer factual recall on the neuron level, 2023. URL https://www.lesswrong.com/posts/iGuwZTHWb6DFY3sKB/ fact-finding-attempting-to-reverse-engineer-factual-recall. ↗
  30. Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos as multidimensional signals with state spaces. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Proce… find
  31. 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. find
  32. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
  33. Judea Pearl. Direct and indirect effects. In Probabilistic and causal inference: the works of Judea Pearl, pp. 373–392. 2022. find
  34. 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. 1526–1534, 2016.
  35. Eric Todd, Millicent L Li, Arnab Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau. Function vectors in large language models. arXiv preprint arXiv:2310.15213, 2023. find
  36. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  37. Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. Investigating gender bias in language models using causal mediation analysis. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processi… find
  38. Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593, 2022. find
  39. Zhen Wang, Jianwen Zhang, Jianlin Feng, and Zheng Chen. Knowledge graph embedding by translating on hyperplanes. In Proceedings of the AAAI conference on artificial intelligence, volume 28, 2014. find
  40. 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. arXiv preprint arXiv:1910.03771, 2019.
  41. Fred Zhang and Neel Nanda. Towards best practices of activation patching in language models: Metrics and methods. arXiv preprint arXiv:2309.16042, 2023. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/sharma-2024-locating-editing-factual-associations-mamba/