The Hidden Attention of Mamba Models
This work offers a novel perspective for understanding Mamba models, potentially aiding researchers in improving sequence modeling techniques, though it is incremental as it builds on existing Mamba frameworks.
The paper reveals that Mamba models, known for efficient selective state space modeling, can be reinterpreted as attention-driven models, enabling empirical and theoretical comparisons with transformer self-attention and providing explainability insights.
The Mamba layer offers an efficient selective state space model (SSM) that is highly effective in modeling multiple domains, including NLP, long-range sequence processing, and computer vision. Selective SSMs are viewed as dual models, in which one trains in parallel on the entire sequence via an IO-aware parallel scan, and deploys in an autoregressive manner. We add a third view and show that such models can be viewed as attention-driven models. This new perspective enables us to empirically and theoretically compare the underlying mechanisms to that of the self-attention layers in transformers and allows us to peer inside the inner workings of the Mamba model with explainability methods. Our code is publicly available.