Understanding Input Selectivity in Mamba: Impact on Approximation Power, Memorization, and Associative Recall Capacity

2025-06-13

AI Summary

This work investigates how Mamba’s input-selectivity mechanism affects its function approximation, long-term memory retention, and associative recall capabilities—and elucidates the underlying mechanisms. Methodologically, we employ state-space model (SSM) analysis, Haar wavelet theory, and dynamical systems modeling. We first prove that the S6 layer exactly implements Haar wavelet projection, revealing a dynamic mechanism that suppresses memory decay; further, we derive an analytical solution and tight theoretical guarantees for the Memory-Query Associative Recall (MQAR) task. Theoretically, Mamba is strictly superior to S4D in approximating discontinuous functions. Empirically, it significantly improves long-range dependency modeling and associative recall performance. Crucially, our theoretical constructions align closely with empirical results, establishing the first unified analytical framework for understanding Mamba’s selectivity mechanism.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Multiagent Systems: Multiagent Learning

Application Category

• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

State-Space Models (SSMs), and particularly Mamba, have recently emerged as a promising alternative to Transformers. Mamba introduces input selectivity to its SSM layer (S6) and incorporates convolution and gating into its block definition. While these modifications do improve Mamba's performance over its SSM predecessors, it remains largely unclear how Mamba leverages the additional functionalities provided by input selectivity, and how these interact with the other operations in the Mamba architecture. In this work, we demystify the role of input selectivity in Mamba, investigating its impact on function approximation power, long-term memorization, and associative recall capabilities. In particular: (i) we prove that the S6 layer of Mamba can represent projections onto Haar wavelets, providing an edge over its Diagonal SSM (S4D) predecessor in approximating discontinuous functions commonly arising in practice; (ii) we show how the S6 layer can dynamically counteract memory decay; (iii) we provide analytical solutions to the MQAR associative recall task using the Mamba architecture with different mixers -- Mamba, Mamba-2, and S4D. We demonstrate the tightness of our theoretical constructions with empirical results on concrete tasks. Our findings offer a mechanistic understanding of Mamba and reveal opportunities for improvement.

Problem

• Understand how Mamba uses input selectivity for performance
• Analyze impact of input selectivity on function approximation
• Investigate Mamba's associative recall and memory capabilities

Innovation

• Input selectivity enhances Mamba's approximation power
• S6 layer dynamically counteracts memory decay
• Mamba architecture solves MQAR task analytically