Research map / Architectures and efficiency / Attention and transformers
MambaLRP: Explaining Selective State Space Sequence Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MambaLRP, a Layer-wise Relevance Propagation variant for Mamba (selective SSM) models that restores relevance conservation via detach-based rules at SiLU, selective SSM, and multiplicative gate layers, yielding more faithful explanations.
Problem
Explanations for Mamba (selective state space) models are unreliable because standard LRP/gradient methods violate relevance conservation at Mamba-specific components (SiLU, selective SSM, multiplicative gating), and existing attention-based explanations are unfaithful.
Approach
Identify where conservation breaks in Mamba (SiLU activation, selective SSM, multiplicative gate) and fix each with detach-based modifications (treating parts of the expression as constant) so that Gradient x Input on the modified model yields conservative LRP relevance.
Key results
- MambaLRP achieves the highest Delta A^F on all text datasets and ImageNet, e.g., 1.978 on SST-2 for Mamba-130M vs 0.877 for LRP (LN-rule) and 1.190 for MambaAttr, and 4.715 on ImageNet vs 2.548 for LRP (LN-rule).
- Ablation shows removing the SSM propagation rule drops SST-2 score from 1.978 to 0.577 and ImageNet from 4.715 to 0.144, making it the most critical component.
- MambaLRP's runtime is comparable to Gradient x Input and needs a single forward and backward pass, while Integrated Gradients is considerably slower.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MambaLRP instead of the strongest baseline explanation method | Mamba models and Vim-S · SST-2, Med-BIOS, SNLI, Emotion, and ImageNet | ΔA^F: SST-2: 0.877 -> 1.978 (Mamba-130M); Med-BIOS: 3.126 -> 3.906 (Mamba-130M); SNLI: 0.673 -> 0.989 (Mamba-130M); Emotion: 3 | Helped |
| Add the proposed SSM propagation rule (ablation) | MambaLRP · SST-2 and ImageNet | ΔA^F: 0.577 -> 1.978; 0.144 -> 4.715 | Helped |
| Add the proposed gate propagation rule (ablation) | MambaLRP · SST-2 and ImageNet | ΔA^F: 1.721 -> 1.978; 4.022 -> 4.715 | Helped |
| Add the proposed SiLU propagation rule (ablation) | MambaLRP · SST-2 and ImageNet | ΔA^F: 1.943 -> 1.978; 4.618 -> 4.715 | Helped |
| Use half-relevance propagation instead of detaching gate z_B (ablation) | MambaLRP multiplicative gate · SST-2 and ImageNet | ΔA^F: 1.577 -> 1.978; 3.592 -> 4.715 | Helped |
| Apply LRP-γ only to conv1d rather than LRP-0 in all listed layers (ablation) | Vim-S · ImageNet | ΔA^F: 4.684 -> 4.715 | Helped |
| Apply LRP-γ to language-model layers | Mamba language models · Language classification datasets | ΔA^F | No clear effect |
Surprising findings
- Applying the LRP-gamma rule to language models gave no discernible improvement, so LRP-0 is used there.
- Mamba-130M and Mamba-1.4B rely less on gendered tokens than GPT2-base, T5-base and RoBERTa-base when predicting the Nurse class in Medical BIOS.
Limitations
- Memory usage of propagation-based methods scales with model architecture and sequence length, since activations and gradients must be stored.
- Requires gradient access, which may be unavailable for proprietary models.
- Attention-based baselines (AttnRoll, MambaAttr) had to be reimplemented for language models since no code was available.
Experimental setup
Datasets: SST-2 Medical BIOS Emotion SNLI ImageNet HotpotQA (LongBench) No Robots needle-in-a-haystack synthetic dataset
Benchmarks: SST-2 Medical BIOS SNLI Emotion ImageNet
Metrics: Delta A^F (flipping AUC difference between LeRF and MoRF) conservation (output score vs sum of relevance) runtime explanation-based retrieval accuracy (XRA) frequency of gendered words in top-5 tokens
Baselines: Gradient x Input SmoothGrad Integrated Gradients Attention Rollout MambaAttr LRP (LN-rule)
Builds on: Layer-wise Relevance Propagation Gradient x Input LRP-gamma Vision Mamba Mamba
Closest papers by meaning
- LongMamba: Enhancing Mamba's Long-Context Capabilities via Training-Free Receptive Field Enlargement (ICLR 2025 · similarity 0.89)
- Block-Biased Mamba for Long-Range Sequence Processing (NeurIPS 2025 · similarity 0.89)
- Routing Mamba: Scaling State Space Models with Mixture-of-Experts Projection (NeurIPS 2025 · similarity 0.88)
- Achilles' Heel of Mamba: Essential difficulties of the Mamba architecture demonstrated by synthetic data (NeurIPS 2025 · similarity 0.88)
- Mamba Modulation: On the Length Generalization of Mamba Models (NeurIPS 2025 · similarity 0.87)
- Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.