MambaLRP: Explaining Selective State Space Sequence Models

Summary

The paper presents a method to correctly apply LRP to Mamba models. Through careful analysis, the authors demonstrate that applying LRP directly results in poor performance and propose modifications to recover the propagation rules. The obtained method outperforms the alternatives, grounded by theory, and the authors show several applications of their method, including identifying gender biases and measuring the long-range abilities of S6.

Strengths

1) **Impact:** Mamba is an emerging architecture (~600+ citations, 10K+ stars on GitHub). Thus, providing an attribution method for these models is crucial. In this essence, correctly applying LRP is an important direction that facilitates the community in improving and understanding these models. 2) **Simplicity:** The method and modifications are very simple, with few hyper-parameters. While this might be seen as a lack of novelty, I view it as an advantage, making the method easy to use and adaptable to various applications, such as other variants of Mamba (Mamba2) and other domains (DNA, speech and more). 3) **Informative Ablation studies** and justification of decision choices are insightful. For example, Figure 3 and Tables 1, 2, 6, 9 provide comparisons with naive LRP and also present ablation studies that allow the reader to measure the contribution of each modification to the method. 4) Section 6 (use cases) shows that the method is **applicable** and allows the authors to explore the gender bias and long-range abilities of Mamba models and provide insightful analyses about Mamba models.

Weaknesses

1) **The comparison with previous work should be improved:** - 1.1) Metrics and Benchmarking: Can the authors highlight which results are reproduced by them and which results are taken from previous work? Moreover, as far as I understand, the method in [4] is the only method that developed an interpretability method for Mamba before (perhaps in parallel) to this work. If I understand correctly, although the two methods share and use the same pre-trained models, there is no overlap in the metrics. Am I correct? If so, is there a reason for this discrepancy? Can the authors compare results directly with previous work using previously proposed metrics to ensure the gap doesn’t arise from employing [4] incorrectly? - 1.2) Informative Comparison: Additionally, I think the comparison with [4] can be more accurate. The method in [4] applies its method only to S6 layers (without gating and convs), while Mamba LRP is an end-to-end method (which is a strength of Mamba-LRP). However, it is still important to make some apple-to-apple comparisons. Can the authors check both methods on a model without conv and gating layers to determine which approach is better for providing an explanation to S6 layers? Alternatively, it seems that the issue with [4] is fixed in [Uni]. Would the authors be able to compare their method to [Uni]? Providing this comparison would be highly valuable to the community. (I understand that this is a very new paper, so I will not decrease my score if the results are less favorable than those of [Uni]). 2) **Insufficient empirical analysis:** - 2.1) For the long-range experiments (Figure 6), it would be very insightful to compare the behaviour of Mamba to Transformers (Pythia pre-trained models can serve as the transformer baselines since they are trained on the same data). Additionally, analyzing the trend with larger models could yield valuable information on how increased model size enhances long-context capabilities in practice, perhaps the 7B models from [Sca] can be used (which exist in Hugging Face). - 2.2) > “Residual lack of conservation is due to the presence of biases in linear and convolution layers, which are typically non-attributable” Can the authors empirically justify this claim? This can be easily validated by taking a pre-trained model, omitting the biases, fine-tuning it for several epochs, and then checking the conservation again. 3) **Novelty:** The novelty of the method is somewhat limited. One could argue that it merely involves a few applications of detach and the simple half-propagation rule (in addition to previously proposed contributions), which can be easily summarized in a few lines of code (as detailed in Algorithm 2). However, I believe this is not a significant drawback, particularly given the demand for such simple tools in the community. Additionally, the thorough evaluation, including insightful ablation studies, novel test case evaluations, and theoretical justification, is sufficiently robust. [Uni] A Unified Implicit Attention Formulation for Gated-Linear Recurrent Sequence. Models. Zimerman et al. [Sca] An Empirical Study of Mamba-based Language Models. Waleffe et al.

Questions

1. Mamba LRP uses "Detach" on the SSM recurrent matrices. Similarly, [4] uses an attention matrix determined by A, B, and C parameters. It seems that both methods ignore the influence of the input on the system matrices (which is the core of the selection mechanism of Mamba). Am I right? Could addressing this issue provide a way to improve both methods? I would be glad to hear what the authors think about it. 2. Minor: Traditionally in the SSM domain, Delta denotes the step size. I suggest the authors replace \Delta with \delta when discussing the differences between the two scores. 3. In the Needle-in-a-haystack experiment (Figure 7), is there a reason not to increase the context length? It seems that the most interesting part of the figure is missing (which could show if there are edge cases, for example, regimes where the model succeeds in finding the needle, but Mamba-LRP fails). 4. Minor: Perhaps a relevant work that is missing is “Does Transformer Interpretability Transfer to RNNs?” by Paulo et al. 5. There is a standard trend in the SSM literature to omit D (and treat it as a parameter-based skip connection). Is it used in Mamba LRP, or is it ignored (like other biases, which are typically non-attributable)? It would be better if it were written explicitly in the paper. 6. I wonder if the authors can explore the potential limitations or failure cases of the proposed Mamba-LRP. Such information can help the community improve the method in the future. Are there cases where [4] or naive LRP might be better than Mamba-LRP? 7. Half propagation: While I'm not an expert in LRP, I suspect there are more effective methods to manage the gating mechanism. For instance, instead of normalizing the scores by averaging (0.5(x + y)), it may be better to use a weighted approach such as $$(1 - a) \cdot x + a\cdot y$$ where the value of a is determined by the actual norms of x and y, with different a values for each channel. Is there something I'm missing? Can this method improve the conservation properties?

Rating

6

Confidence

5

Soundness

3

Presentation

4

Contribution

3

Limitations

From my perspective, the authors address most of the limitations, except those pointed out in the weaknesses section and question 6 (failure cases).

Summary

### **Post-rebuttal update** Given the authors' additional experiments and changes-to-be-made to the manuscript, raising my original score from a 5 to a 6. ### **Original review** The authors tackle the problem of explainability in Mamba SSMs, which have been recently proposed and widely adapted. Towards this end, they leverage layer-wise relevance propagation (LBP) and derive the necessary procedures/equations to enable LBP in Mamba models, similar to those previously proposed in Conservative Propagation paper for Transformers (Ali et al, 2022). They benchmark their procedure on both Mamba language and vision models across a range of benchmarks and other widely used explainability methods to demonstrate the effectiveness of their approach.

Strengths

**Originality** - Given the wide-spread adaptation of Mamba SSMs since their recent debut, explainability of these models is a very important topic. While the specific detach approach to enable LRP in LLMs is not new, the adaptation to Mamba models is new and an excellent application of this methodology. **Quality** - The proposed approach is a good addition to the growing list of Mamba works, and will make a good tool for practitioners/researchers to explain the decisions of Mamba models going forward. The extensive benchmarks and evaluation of both NLP and vision models and tasks are also solid contributions. However, the evaluation of Mamba LLMs was somewhat limited (see weaknesses for further discussion). **Clarity** - For the most part, the paper is well written. However, important details regarding evaluation are either missing or could be better highlighted (discussed in weaknesses). **Significance** - As previously mentioned, Mamba models continue to gain traction, particularly with the recent release of Mamba 2. Thus, this work has significant opportunity for significant impact as the need to explain the decision making processes behind Mamba models grows.

Weaknesses

### Quality As previously noted, there is significant room for improvement regarding the evaluation of Mamba LLMs. In particular, only the smallest and second largest checkpoints are used for the majority of experiments (with the exception of the 2.8B checkpoint used in the needle in the haystack tests). However, for a paper dedicated to enabling the capability of this line of LLMs, this seems insufficient; why are the 130M and 1.4B checkpoints used for the majority of tests? In particular, the 130M model performs the worst among the suite of Mamba 1 models, but is used to demonstrate qualitative claims about the method (e.g., Figure 3, Figure 4, and reported runtimes). Given the 2.8B is the most accurate and, thus, most desirable for practical use, this should ideally be evaluated along with the 130M (and 1.4B) models to demonstrate the efficacy of this approach at 2.8B Mamba-parameter size. ### Clarity For both the presented experiments and compared methods, several important details were not clear in the paper. In particular, for Figure 6, how is position difference calculated? E.g.: >We use the HotpotQA [70] subset from the LongBench dataset [11], designed to test long context understanding. After selecting all 127 instances, containing sequences up to 8192 tokens, we prompt the model to summarize the full paragraph by generating ten additional tokens Fig. 6, shows the distribution of the positional difference between a relevant token and the currently generated token. How are the relevant tokens determined? Upon first pass, I wondered whether they were somehow derived from the supporting fact labels in the HotpotQA dataset, but it seems to be calculated as the difference between the tokens generated position vs position in the input. Can you clarify this, as well as how (exactly) the histogram in Figure 6 is calculated? Also, the HotpotQA subset of LongBench contains 200 instances, yet the paper evaluates 127; what is the source of this discrepancy? Most importantly, it is not clear whether the author's are comparing to Attention Rollout, which the Ali et al 2024 Mamba explainability paper was based on, or Mamba-Attribution (i.e., the new mechanism derived for Mamba models in the Ali et al 2024 paper). As this is the most relevant Mamba-specific method given the proposed method, this point requires clarification within the main text. Other portions which are unclear: - "The results without fast CUDA kernels" <- what are fast CUDA kernels in this context? The hardware-optimized selective SSM scan kernel included in Mamba (in contrast to the full PyTorch SSM kernel)? Or are these kernels the authors are contributing to speed up their described method? - "All models were trained for a maximum of 10 epochs, with an early stopping mechanism in place." <- What is the early stopping criterion? - Lines 192-193 should forward reference fine-tuning details are further described in "C.1 Models and dataset." Otherwise, it is easy to miss that small excerpt on line 193 which alludes to 130M and 1.4B being fine-tuned for the various tasks (which could lead to thinking pretrained Mamba models are being evaluated).

Questions

"We use an instruction-finetuned Mamba-2.8B model" <- In order to make the paper as self-contained as possible, could the authors summarize the fine-tuning recipe for this model within the supplementary (rather than forward referencing to the huggingface page)?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Could the authors include foreseeable limitations of their work? E.g., potentially large memory utilization to enable explainability in Mamba models? Also, in the case that the compared method AttnRoll was not Mamba-Attribute (or not evaluated using the released code of the Ali et al 2024 paper), these would be important evaluation limitations to list.

Summary

The paper introduces an LPR framework for Mamba. The method breaks the Mamba architecture into three parts (SiLU activation, selective SSM, and multiplicative gating) and analyzes the layers using relevance scores. The evaluations on languages and images show that the proposed method is more precise and faithful than other explanation methods. They also show several use cases including gender bias, long-range capabilities, and the needle-in-a-haystack test.

Strengths

- The paper is easy to follow. - The introduced LRP framework is faster and provides more faithful explanations than other explanation methods. - The experiments show interesting interpretations of Mamba specifically for gender bias, long-context capabilities, and needle-in-a-haystack test in Mamba.

Weaknesses

- The current paper is limited to the explainability of Mamba only. I believe LRP can be applied to Transformers. It would be more interesting to compare the behavior of different methods instead of showing the behavior of one method. - This point is related to the first point. The use case experiments are interesting, but it would be more useful if this could be compared with different methods. For instance, in the long context capability of Mamba, the paper mainly shows that Mamba can use earlier context. Instead, the paper can show that one method has better long context capability than the other in terms of accuracy, and it is verified by LRP. Also, the behavior of different methods can be compared such as 1) if the Transformer-based method is not as good as Mamba in terms of long-range modeling, what does the Transformer model focus on 2) how does the behavior change with tasks (generation vs QA) 3) how does it change with the training datasets, etc. Overall, the paper lacks a more interesting and useful analysis of the methods using LRP as explained above. This limits the contribution of the paper.

Questions

1. Tables 10 and 11 show the runtime comparisons with and without using fast CUDA kernels. When using fast CUDA kernels, other methods including gradient × Input, SmoothGrad, and Integrated Gradients improve the speed a lot (22-25x faster), but the proposed MambaLRP has only 1.3x speed-up. What is the reason?

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The paper briefly explained the limitation in Section 7.

Summary

The paper applies Layer-wise Relevance Propagation to Mamba layers. To maintain the relevance conservation law, the authors propose three fixes to SiLU activation, S6 and the multiplicative gating operators respectively with the technique of gradient blocking. The proposed method improves the faithfulness of LRP explanantion on Mamba significantly.

Strengths

1. The paper is well written and easy to follow. 2. The improvement of the faithfulness is significant compared to the baselines without the fixes proposed in the paper.

Weaknesses

1. The contribution of the work seems incremental and the techniques used are well established. The main novelty here is applying the existing LRP techniques to the Mamba architecture. 2. The explanation generated by MambaLRP is not surprising and does not bring any new insights into the behavior of architectures. It is unclear how these explanations can be used to improve the model's performance on alleviating gender bias and retrieving long-range information.

Questions

Can we leverage the explanations generated by MambaLRP to improve the model 's behavior on the downstream gender debiasing and the long-context retrieval tasks?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

2

Limitations

No. The authors should discuss the limitation of the usefulness of their methods.

Response to Reviewer dNt6: Further Clarification

Due to space constraints, we had to keep our earlier responses brief. We would now like to take this opportunity to expand on some of our answers to ensure greater clarity and convenience for you. > Further elaborations on the experimental results of long-range dependencies (LRD) experiment Inspired by your comments, we conducted comparisons using two state-of-the-art Transformers: Llama-2 and Llama-3. **Setup:** We employed LRP propagation rules of [3] to extract explanations for the Llama models. As in the Mamba experiment, we generated 10 additional tokens from HotPotQA's input and analyzed the prediction of the generated token at each step. **Results:** Our findings are summarized in Table A of the attached PDF. Notably, Llama-2, which was trained with a context length of 4096 tokens, begins to produce less sensical text when given inputs exceeding this length. This behavior aligns with observations in recent studies [7,8]. In contrast, Llama-3 and Mamba are capable of generating sensible text even with longer context lengths, as shown in Table A of the PDF. Although Llama-2’s histogram suggests it uses the full context window and appears to identify more relevant long-range input tokens than Llama-3 and Mamba, closer inspection reveals that this long-range capability is based on unspecific token information. The most relevant tokens for Llama-2 are often non-semantic, such as the newline token `<0x0A>` and the beginning-of-sentence token `<s>`, which often appear at the beginning of the context paragraph. This reliance on non-semantic tokens explains the long-range information retrieval seen in the histogram, particularly in the context size exceeding 4K tokens (the shaded areas in the histogram of Llama-2). In contrast, Llama-3 and Mamba primarily attribute relevance to semantically meaningful tokens, particularly those near the end of the input context window. Llama-3’s histogram indicates it uses more intermediate mid-range dependencies compared to Mamba. Given Llama-3’s significantly larger size (8B) compared to Mamba (130M), our initial analysis shows that Mamba effectively uses long-range information. We also find that this ability is not exclusive to SSMs but can also be achieved by larger Transformers (e.g. Llama-3). In this initial investigation, MambaLRP has proven valuable in comparing the long-range context capabilities across models. We hope our work will facilitate further comparative studies in this area. We will include a summary of this comparison in our final paper. ---- **Questions** 3. > Is there a reason not to increase the context length in needle-in-a-haystack of Figure 7, and are there any failure cases? We did not include context lengths beyond 2048 in the paper, as the model was trained using a context length of 2048 and the study was not designed to assess extrapolation beyond this limit. Our study highlights some limitations of the retrieval accuracy metric and introduces an explanation-aware measure. This new measure ensures that the retrieved information is not only correct but also correct for the right reasons. As requested, we extended the experiment to context lengths beyond 2048, up to 4096, and included the results in Figure C of the PDF. As can be seen, the model's retrieval performance begins to drop as the context length goes beyond 2560. In our analysis, we did not encounter any edge cases where the model successfully retrieved the needle but MambaLRP failed to provide the correct explanations. However, we have shown cases where the model found the needle based on incorrect features, which MambaLRP was successfully able to detect. Please refer to Appendix C.7 (Figure 13). 4. > Including "Does Transformer Interpretability Transfer to RNNs?" Thanks for your suggestion. We will add it to our related works. 5. > Is $D$ used in MambaLRP? Skip connections are not ignored in MambaLRP. As they do not violate the conservation property, no modifications are necessary. This is why we did not mention them in the paper. However, based on your suggestion, we will include it in the paper. If you have any further comments or questions, we are happy to address them during the discussion period.

Reply to rebuttal

I thank the authors for their detailed response. My concerns have been addressed, and the authors have agreed to incorporate these changes into the next draft of the paper. As previously stated, the submitted work is especially important as the Mamba family of models continues to be widely adapted. > It is important to note that our experiments aim to evaluate the effectiveness of our explanation method against other methods applicable to Mamba-based models, rather than evaluating the capabilities of the Mamba architecture itself. To our knowledge, this is the most comprehensive evaluation of different attribution methods for Mamba models in the literature. Previous work [MambaAttr2024] was tested on the Mamba-130M and Vim-S model. The breadth of evaluation is appreciated. Note that MambaAttr2024 is an unpublished manuscript, so their evaluation has not faced scrutiny for NeurIPS publication. For NeurIPS, rigor is necessary. Evaluating the proposed method on the largest checkpoint is representative of how well MambaLBP will work on the Mamba model most likely to be used by LLM practitioners; in personal experience, the behavior/performance of the 5 Mamba LLM checkpoints varies significantly. Thus, generalization across the checkpoints is not a given, and the authors' additional experiments to confirm this greatly help to address this concern. I look forward to the other NLP experiments evaluated on the 2.8B model, and I plan to raise my score from a five to a six. As noted by another reviewer, the reproducible implementation of this method is a critical contribution, as well as the benchmarking suite. Mamba and Mamba2 have proven temperamental for features considered standard for Transformer models (e.g., unmodified use in Huggingface's SFTTrainer), so there is a large margin for error involved in researchers reproducing these results from scratch. What are the authors' plans for releasing code to reproduce results from the paper?

response to the rebuttal

Thank you for the answers. The additional analysis provided during the rebuttal is helpful. Although the rebuttal response partially addresses my concerns (limited empirical analysis), it certainly improves the quality of the paper. I also believe that the proposed method would be useful to the community. I increase my rating from 4 to 6. Please add the full analysis between Transformers and Mamba (like Figures 6 and 7 in the paper) to the revised version.