MambaLRP: Explaining Selective State Space Sequence Models

2024-06-11

arXiv.org

AI Summary

Selective state space models (e.g., Mamba) suffer from poor interpretability due to their hardware-aware scanning and gated state update mechanisms, which violate standard attribution assumptions. Method: This work introduces MambaLRP—the first layer-wise relevance propagation (LRP) method tailored for Mamba—by customizing LRP rules to respect Mamba’s unique architecture: it rectifies non-conservative components in the scanning and gating operations and enforces strict relevance conservation. Contribution/Results: MambaLRP is the first theoretically sound, stable, and faithful explanation method for Mamba. It achieves state-of-the-art interpretability performance across multiple Mamba variants and benchmark datasets. Empirical analysis reveals internal model biases and validates Mamba’s long-range dependency modeling capability, thereby enhancing decision transparency and enabling trustworthy deployment.

Technology Category

• Machine Learning: Transparent, Interpretable, Explainable ML
• Natural Language Processing: Interpretability, Analysis, and Evaluation of NLP Models
• Computer Vision: Interpretability, Explainability, and Transparency

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

Recent sequence modeling approaches using selective state space sequence models, referred to as Mamba models, have seen a surge of interest. These models allow efficient processing of long sequences in linear time and are rapidly being adopted in a wide range of applications such as language modeling, demonstrating promising performance. To foster their reliable use in real-world scenarios, it is crucial to augment their transparency. Our work bridges this critical gap by bringing explainability, particularly Layer-wise Relevance Propagation (LRP), to the Mamba architecture. Guided by the axiom of relevance conservation, we identify specific components in the Mamba architecture, which cause unfaithful explanations. To remedy this issue, we propose MambaLRP, a novel algorithm within the LRP framework, which ensures a more stable and reliable relevance propagation through these components. Our proposed method is theoretically sound and excels in achieving state-of-the-art explanation performance across a diverse range of models and datasets. Moreover, MambaLRP facilitates a deeper inspection of Mamba architectures, uncovering various biases and evaluating their significance. It also enables the analysis of previous speculations regarding the long-range capabilities of Mamba models.

Problem

• Mamba Model
• Interpretability
• Long Sequence Processing

Innovation

• MambaLRP
• Layer-Relevance-Propagation
• Model-Interpretability