How does Mamba Perform Associative Recall? A Mechanistic Study

Grégoire LE CORRE, Ningyuan Huang, Alberto Bietti

Keywords: Probing, Circuit analysis, Understanding high-level properties of models

Other Keywords: Mamba, Associative-Recall

TL;DR: We probe Mamba to identify its associative-recall mechanisms

Abstract: Mamba has recently emerged as a promising alternative to Transformers, demonstrating competitive performance in many language modeling tasks with linear-time computational complexity. Theoretical characterization of Mamba has largely focused on its approximation power for solving certain tasks through specific constructions. However, it remains unclear whether Mamba trained with gradient descent can learn such constructions. As a first step to address this gap, we perform a mechanistic study of simplified Mamba models on associative recall tasks. By analyzing the learned model weights and the hidden state evolution, we uncover the mechanisms used by simplified Mamba models to perform associative recall. We complement our study with theoretical analysis on the optimization dynamics of simplified Mamba models that give rise to such mechanisms.

Submission Number: 198