MambaLRP: Explaining Selective State Space Sequence Models

Thanks for the valuable comments. We address them below and add more discussions as an official comment.

1.1: We use the official code of [1] to produce the results of AttnRoll and G$\times$ AttnRoll (Mamba-Attr). We evaluated MambaLRP's performance against these approaches using flipping and insertion metrics that are well-established for analyzing faithfulness [4,6] and are directly related to those used by [1]. In our metrics, $A^F_{MoRF}$/$A^F_{LeRF}$ corresponds to what [1] calls positive/negative perturbation. [4] suggests combining both metrics to derive a more resilient metric, resulting in $\Delta A^F = A^F_{LeRF} - A^F_{MoRF}$, used in [4,5]. The only difference is that [1] tracks accuracy in positive/negative perturbations, while we track changes in output logit, which is a standard practice in the XAI community [4,5,6]. Both metrics highly overlap as shown in Tab.B of the attached PDF. When using faithfulness metrics, various factors can impact the result, e.g. how the input is masked. To ensure a fair and consistent analysis across all models, tasks and XAI methods, we use a unified evaluation metric, as done in [3,5,6], and have not directly used the numbers reported in [1].

As requested, we follow the metric of [1] and mask pixels from least-to-most relevant and vice versa (10%-90%), track the top-1 accuracy and calculate AUC for the prediction on ImageNet val set for Vim-S (Tab.C in the attached PDF). As shown, MambaLRP outperforms [1].

1.2: [1] argues that Mamba uses implicit attention mechanisms and provides XAI tools, which aggregate attention across layers and may not capture the impact of certain layers. In contrast, MambaLRP, as you noted, considers the full model structure, highlighting its strength. Moreover, [1] introduced Mamba-Attr, which as said in the paper "exploits the gradients of both the S6 mixer and the gating mechanism''. In our evaluations, we compared MambaLRP with G$\times$AttnRoll (Mamba-Attr). As the methods of [1] are designed to explain the predictions of Mamba models, it is both valid and fair to compare MambaLRP with them in terms of the faithfulness of the explanations, as done in [3,5]. Despite [2] being released after the NeurIPS deadline, we compared MambaLRP's performance with the numbers reported in [2] for Vim-S in Tab.C of the PDF. This shows that MambaLRP outperforms the recent approach of [2]. We have added [2] to our related works.

2.1: As requested, we did a direct comparison to Transformers for our LRD use case (Tab.A of attached PDF). Many of the widely-used Transformers, e.g. GPTNeoX and Pythia unfortunately, do not allow inputs longer than 2048 or do not generate sensible text from long inputs. Instead, we used Llama-2 and Llama-3. We found that Llama-3 uses information from more intermediate mid-range dependencies than Mamba, though both favor tokens close to the end of context. Given Llama-3's larger size (8B) compared to Mamba (130M) and their different training settings, the analysis supports that Mamba can effectively use long-range information. Please refer to our answer to reviewr 4ExW for details.

2.2: Please find the results in Fig.A of PDF; conservation is fully preserved for biases set to 0.

3: As many XAI methods, including LRP, are not model-agnostic, tailored approaches must be developed for emerging classes of DNNs like Selective State Space Sequence Models (Mamba models). Mamba models include previously unstudied components, making existing LRP rules unreliable for these models, as shown in Tab.1 of paper. Since the explainability of Mamba models using LRP has not been addressed in the literature, we propose novel LRP rules specifically for SSM components, SilU, and multiplicative gates. These rules are not chosen heuristically but are the result of our thorough analysis of how relevance propagates through each component and when conservation breaks. Please refer to Appendix A and B for our theoretical analysis and derived LRP rules. As ease of implementation is desirable for XAI methods, we provided straightforward implementations that bypass the need for the implementation of complex LRP rules. These highlight the strengths of our approach: 1) theoretical soundness, 2) more faithful explanations, and 3) ease of implementation. Our contributions are the theoretical analysis of relevance propagation in different Mamba components, proposing new LRP rules for those that violate conservation, providing straightforward implementations, thorough performance evaluation against other XAI approaches, and insightful use cases, showing the value of MambaLRP for other lines of XAI research.

Questions

1: In forward path, $A$, $B$, and $C$ are calculated based on the input and are treated as constants in backward path. By detaching, we do not consider the gradient flow through them; however, their influence is still addressed by weighting the input in backward pass.

3: We initially limited the context length to 2048 to match the model's training and focus on retrieval accuracy limitations and MambaLRP solutions. At your request, we extended the experiment to 4k context length, with results in the PDF. Performance drops beyond 2560 and we found no cases where the model retrieved correctly but MambaLRP failed to explain reliably.

5: $D$ is used in MambaLRP. As it does not violate conservation, no modifications are necessary.

6: We investigated edge cases using Mamba-1.4B trained on Med-BIOS and found that MambaLRP outperforms Mamba-Attr [1] in about 95% of instances. In the most extreme edge case, both methods mostly select the same relevant tokens, but their score distributions differ slightly (see Fig.B of PDF).

7: In our work, $\alpha$ is not treated as a hyperparameter but set to 0.5 based on theoretical considerations. Specifically, in Equation (12) of the Appendix A.3, we show that $\mathcal{R}(x) = 2\mathcal{R}(y)$, leading to the choice of 0.5.

We will include your suggestions in the paper.

Thank you for your valuable feedback and suggestions. We are happy to see that you have increased your score. We have included the new experiments in our paper.

Regarding making the code accessible

We understand the importance of providing accessible and clean code to the community. Therefore, we are preparing a user-friendly GitHub repository, and we will include the link in the camera-ready version. We further plan to extend MambaLRP and our code to include Mamba variants including the recent Mamba2 model.

Thank you again for the support of our work and we look forward to see how the open access to reliable XAI for Mamba can lead to domain insights and model improvements in our community.

Thanks for your detailed review and useful suggestions. We will address your points below.

Evaluation of Mamba LLMs also on larger 2.8B models

With our empirical evaluations in Table 1, we aimed to cover a representative range of tasks (four different text classification tasks, image classification), models (Mamba, Vision Mamba) and model sizes (130M and 1.4B), which confirmed that our proposed approach consistently have highest faithfulness scores---independent of the performance, architecture and domain of the specific Mamba model, compared to existing explanation methods.

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.

To further extend our evaluations and follow your suggestion, we fine-tuned Mamba-2.8B model on the SST-2 dataset using the same protocol described in Appendix C.1. The following table presents the performance evaluation of different explanation methods on this model, confirming the superior performance of MambaLRP. We are currently running this experiment on the other NLP datasets as well and will add the results to the final version of the paper.

Mamba-2.8B results:

How are the relevant tokens in Figure 6 determined and how is the histogram calculated? Why are there 127 of 200 HotpotQA samples considered?

As mentioned in Section 6, we used the top-k most relevant tokens that received positive relevance scores when applying MambaLRP and plotted the resulting histogram of relevant tokens over the positional differences between the generated and identified relevant tokens. In the paper, we present results for $k=10$. We have further validated that the shape of the histogram distribution of relevant tokens remains consistent for $k \in \{1,3,5\}$. By applying a context length cutoff of 8096 to filter HotpotQA, we obtained the 127 samples used in our experiment.

Clarification regarding comparison to [MambaAttr2024]

We compared our proposed method to both Attn-Rollout and Mamba-Attr introduced for the Mamba models in [MambaAttr2024]. We referred to them in our paper as AttnRoll and G$\times$AttnRoll, in consistency with the respective methods proposed for Transformers. We have modified the name G$\times$AttnRoll to Mamba-Attr in our paper.

What are fast CUDA kernels in this context?

When we mention fast CUDA kernels, we are referring to the hardware-optimized selective SSM scan kernel included in Mamba. Our proposed explanation method does not involve any additional optimizations.

What is the early stopping criterion?

In our experiments, we employ early stopping and end training as soon as the validation loss ceases to improve. We have now added this information to the paragraph on training details in Supplement C.1.

Adding references to Appendix C.1 and adding details about instruction-finetuned model

Thanks for your suggestions. We have now added a reference to enhance accessibility and clarify our experimental setup in Section 5. We will also add more information about the instruction-finetuned model to the supplementary material.

Regarding including foreseeable limitations of our work and potential memory limitations, clarification on comparison to [MambaAttr2024]

Regarding memory consumption, LRP is a backpropagation-based explanation method that requires storing activations and gradients. The memory usage depends on the model and input size. To reduce memory consumption, techniques such as gradient checkpointing can be employed. This is also true for other gradient-based methods.

Regarding the comparison with [MambaAttr2024], we did evaluate the performance of MamabLRP against Mamba-Attr (see also our answer to your previous point above). We have now clarified this potential misunderstanding in the main text and apologize for any confusion about the specific baseline method we used. For our evaluation of Mamba-Attr and Attn-Rollout, we have used the official code provided by the authors, as described in Appendix C.3, providing a fair and reproducible benchmark evaluation.

A potential limitation of gradient-based explanation methods, like MambaLRP, is that the gradient information might not be accessible, e.g., due to proprietary constraints. In these situations, one possible solution is to approximate the gradient information. We will use the additional space to add a dedicated Limitations section to outline and discuss these aspects of our method.

[MambaAttr2024] A. Ali, et al. The hidden attention of mamba models. arXiv:2403.01590, 2024.

Thanks for your detailed feedback. We are happy to see that you appreciate the potential of our approach in providing faithful explanations. We will address your points below.

Regarding the paper being limited to the explainability of Mamba

The class of Selective state space sequence models (Mamba models) present a significant change in model architecture compared to Transformers, necessitating tailored XAI methods to address their unique components. As LRP is not model-agnostic, new LRP rules must be developed for emerging classes of deep neural networks, as they may contain new components for which no existing LRP rules apply. For instance, the LRP framework has already been successfully extended to include Transformers in [3,5].

Mamba models, however, include new components absent in Transformers, and no prior propagation rules exist for them. Consequently, following the work of [3,5,9,10,17] who have extended the LRP framework to new architectures (e.g. Transformers, CNNs, GNNs, regression networks), we derived new LRP rules for Mamba models. Proposing these rules is not trivial; it requires performing a detailed analysis of the relevance propagation process, guided by the conservation axiom. This resulted in more faithful explanations compared to naive applications of existing LRP rules not designed for Mamba and also other approaches, presented in our main paper (Table 1).

Regarding comparing to the behavior of Transformers

While we agree that comparing the capabilities of Transformer and Mamba is an interesting line of research, our work focuses on deriving and thoroughly evaluating our proposed explanation approach for Mamba models. In this, we follow the standard approach to evaluate new explanation methods in terms of faithfulness [4, 6], which is inherently model-specific. Therefore, a comparison to transformer models is not included in the evaluation section of our paper. Instead, to show the versatility and robustness of our proposed method, we tested MambaLRP across various selective state space model architectures and sizes.

To compare the capabilities of Mamba and Transformer using their explanations, we have extended our long-range dependency experiments to include Llama-2 and Llama-3 Transformers (see the attached PDF for results). We describe our findings just below. Moreover, in our bias detection use case in Section 6 of our paper, we have already compared the performance of Mamba-130M and Mamba-1.4B to several Transformer models.

Comparing long-range dependencies of Mamba and Transformers

Thank you for this suggestion. We also found comparing Mamba and Transformer models in terms of their long-range capabilities very interesting. Thus, we investigated this further and did a direct comparison to Transformers for our long-range dependency use case. Many of the widely-used Transformers, e.g. GPT-2, GPTNeoX and also Pythia unfortunately, do not allow inputs longer than 2048 or do not generate sensible text from these long inputs. Instead, we use two state-of-the-art Transformers: Llama-2 and Llama-3, and extract attributions using the LRP rules of [3]. As in our Mamba experiment, we generate 10 additional tokens from the HotPotQA's input and at each step explain the prediction of the generated token. Please find the results in Table A of the attached PDF.

For Llama-2, trained with a context length of 4096, the generated text becomes increasingly less sensible and repetitive for contexts longer than 4k, a limitation noted in [7,8]. When analyzing histogram distributions across models, it seems that Llama-2 uses information more uniformly across the entire context and identifies more relevant long-range dependent tokens compared to Llama-3 and Mamba. However, its output becomes non-sensible for lengths above 4K tokens. Thus, the identified relevant tokens are typically non-semantic such as new line token <0x0A> in Llama-2 and beginning of sentence token <s> found at the start of the context paragraphs. For Llama-3 and Mamba, the attributions can identify meaningful relevant tokens. When directly compared, Llama-3 uses information from more intermediate mid-range dependencies than Mamba, though both favor tokens close to the end of the input as relevant. Given Llama-3's much larger size (8B) compared to Mamba (130M) and their different training settings, this analysis supports that Mamba can effectively use long-range information. We also find that this ability is not exclusive to SSMs but can also be achieved by Transformers.

In this first step, MambaLRP allowed us to compare long-range context capabilities across models and hope our work facilitates the generation and investigation of future comparative studies. We will add a summary of this comparison to our final paper.

Regarding the reduced speed-up of MambaLRP compared to Gradient × Input, SmoothGrad, and Integrated Gradients.

We apologize for the confusion caused by a typo in the number. The correct runtime for MambaLRP listed in Table 11 is 0.03063. This means when using or not using fast CUDA kernels, the runtime of MambaLRP is comparable to GI.

Please let us know if you have further questions or comments. At this stage, we are still able to make revisions to the draft and actively participate in the discussion.

Thank you for your review and useful comments. We will address them in our response below.

Regarding the main novelty of applying LRP to the Mamba

As new model architectures, such as SSMs and Mamba models, are developed, the field of XAI is challenged to develop faithful attribution methods to explain their predictions, especially given that the naive applications of existing XAI methods often fail in this regard. As LRP is not a model-agnostic framework, new LRP rules must be derived as new classes of deep neural networks emerge, since these models may include new layers or components for which no LRP rules exist. LRP was initially introduced in [9] to explain predictions for kernel-based classifiers and for multilayered neural networks. Over the years, it has been extended to accommodate new architectures. For instance, as the existing LRP rules were insufficient for explaining models such as Transformers, [3,5] introduced novel LRP rules specifically for Transformer models. Similarly, with the advent of Graph Neural Networks (GNNs), [10] proposed GNN-LRP to explain their predictions.

We would like to emphasize that for the class of selective state space sequence models (i.e., Mamba models), there exist no propagation-based techniques such as LRP that effectively address the Mamba architecture. A naive application of existing LRP techniques, which are not designed for Mamba, leads to far inferior and unreliable feature attributions as shown in Table 1. To overcome this shortcoming, we have (1) identified the unreliable model components and (2) derived novel propagation rules for them, which form MambaLRP. Our novel LRP rules are grounded in our theoretical analysis of the relevance propagation process through different layers/components within Mamba, which identified those that violate the conservation property. Consequently, we proposed novel LRP rules for the SSM components, SiLU activation functions, and multiplicative gates. The explicit LRP rules that we have proposed for Mamba models can be found in Appendix B. Given that ease of implementation is a key characteristic of XAI methods, we have further contributed by proposing straightforward implementation strategies that bypass the need for the implementation of complex LRP rules in Section 4.

Regarding MambaLRP explanations not being surprising, without new insights into the behavior of architectures.

In the field of XAI, one line of research focuses on developing reliable explanation methods for DNNs [3,5,9,10,13,14] and another complementary line of research on using these methods for in-depth model analysis [12,15] and insight [16]. XAI applications, such as model debugging, identifying biases, examining fairness, and assessing capabilities like long-range dependencies, rely on high-quality explanations. Thus, our study specifically targets the challenge of generating high-quality explanations for the novel class of selective state space sequence models, as their highly non-linear and complex architectures make explaining them a significant challenge.

As we have proposed a novel explanation algorithm for Mamba, our main experiments are focused on analyzing the faithfulness of the explanations generated by MambaLRP against those produced by other methods, shown in Table 1. We demonstrate that our method provides more faithful explanations compared to existing alternatives while being computationally more efficient (see Appendix C.8). We observed that this approach does allow us to identify unexpected model behaviors, as illustrated in our image classification experiment in Figure 5, where the prediction of the class 'carton' is influenced by the presence of a watermark on the image. To bring new insights into the model's behavior, we used MambaLRP to uncover gender bias, investigate long-range capabilities of Mamba, and propose an explanation-aware measure for the Needle-in-a-haystack test, shown in Section 6.

Regarding leveraging MambaLRP explanations for model improvement

Please note that our use cases (Section 6) focus on model diagnosis rather than model improvement. Identifying a model's weaknesses is the essential first step before any improvements can be made. For example, as illustrated in Figure 5, MambaLRP can be employed to analyze the model's sensitivity to Clever-Hans features, such as watermarks in Chinese. Additionally, it can be used to detect gender biases, as discussed in Section 6. In the case of gender debiasing, as an example, if ground-truth explanations are available, the model can be fine-tuned to align the generated explanations with the ground-truth as done in [11,15] for other XAI approaches and bias types. Regarding long-context retrieval tasks, the needle-in-a-haystack test is designed to evaluate such capabilities in LLMs. In Section 6, we introduced a new metric based on the explanations generated by MambaLRP to measure the model's retrieval performance in an explanation-aware manner. This means that instead of solely checking retrieval accuracy, one can also assess if the model retrieves the right information for the right reasons. For more detailed and interesting results, please refer to Figure 13 in Appendix C.7. In that figure, we have shown a case where the model retrieved the correct information based on incorrect evidence, which MambaLRP was successfully able to detect. This failure case could not be identified by the retrieval accuracy metric typically used in the needle-in-a-haystack experiment, highlighting the value of MambaLRP in providing practical insights. Leveraging XAI methods for improving safety aspects of models is a very active line of research and crucially requires reliable explanation methods. With MambaLRP, we have proposed a faithful and robust explanation method for this purpose.