MambaLRP: Explaining Selective State Space Sequence Models
Authors: Farnoush Rezaei Jafari
Grégoire Montavon
Klaus Robert Muller
Oliver Eberle
Venue: NeurIPS 2024 poster (2024)
Keywords: Explainable AI
explainability
interpretability
state space models
Mamba
long-range dependency
Primary area: interpretability_and_explainability

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.

MambaLRP: Explaining Selective State Space Sequence Models
Farnoush Rezaei Jafari 1.2-
Grégoire Montavon 3.2.1
Klaus-Robert Müller | (2,45666
Oliver Eberic1.2+
"Machine Learning Group, Technische Universität Berlin, 10587 Berlin, Germany
*BIFOLD Berlin Institute for the Foundations of Learning and Data, 10587 Berlin, Germany *Department of Mathematics and Computer Science, Freie Universität Berlin,
Arnimallee 14, 14195 Berlin, Germany
DDpaarteett of Artificial Intelligence, Korea University, Seoul 136-713, South Korea Max Planck Institute for Informatics, Stuhlsatzenhausweg 66123 Saarbrüccen, Germany Google DeepMind, Berlin, Germany
Abstract
Recent sequence modeling approaches using selective state space sequence models, referred as Mamba models, have seen surge of interest. These models allow efficient processing of long sequences in linear time and are rapidly being adopted in 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 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 more stable and reliable relevance propagation through these components. Our proposed method is theoretically sound and excels achieving state-of-ther explanation performance across 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.
Introduction
Sequence modeling has demonstrated its effectiveness and versatility across wide variety of tasks and data types, including text, time series, genomics, audio, and computer vision [24, 82, 32, 9, 26]. Recently, there has been surge of interest as structured state space sequence models (SSMs) [35, 66, 33). This is due to their ability to process sequences in linear time, opposed to quadratic time required by the more established Transformer architectures [69]. The recent Mamba architecture, prominent and widely adopted instance of state space models, has demonstrated competitive predictive performance modeling tasks across domains and applications [33, 43, 83, 78, 70], while scaling linearly with sequence length.
class of sequence modeling architectures, known
variety of sequence
Cooresponndecce to: rezaeijafirr @ campus.tu coapppu.tuu-brliiie berlin.de, oliver.eberle@mu bu-berlin.de
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

a)
b)
c)
BASIC LRP
MAMBALRP
fiz)
Block 5
Block 6
Block 5
Block 6
928
H
4
0.0226
veliion
conservation
V f O æ
noisy
explonations
explorations
Figure I: Conceptual steps involved in the design of MambaLRP. (a) Take as starting point basic LRP procedure, equivalent to Gradient x Input. (b) Analyze layers in which the conservation property violated. (c) Rework the relevance propagation strategy at those layers achieve conservation. The resulting MambaLRP method enables efficient and faithful explanations.
As Mamba models, and more generally SSMs, are rapidly being adopted into real-world applications, ensuring their transparency is crucial. This enables inspection beyond test set accuracy and uncovering various forms of biases, including Clever-Hans effects [39]. It is particularly important in high- risk domains such as medicine, where the prediction behavior must be robust under real-world conditions and aligned with human understanding. The field of Explainable AI [48, 36, 8, 58] focuses on developing faithful model explanations that attribute predictions to relevant features and has shown success in explaining many highly nonlinear models such as convolutional networks [20], or attention-baa based Transformer models [3.2]. 2].
Explaining the predictions of Mamba models is however challenging due to their highly non-linear and recurrent structure. A recent study [4] suggests viewing these models attention-based models, enabling the use of attention- by attention- based techniques are often unreliable and exposed potential misalignment between input features and attention scores [75, 38). As an alternative, Layer-wise Relevance Propagation (LRP) [10] decomposes the model function with the goal of explicitly identifying the relevance of input features by applying purposely designed propagation rules at each layer. A distinguishing feature of LRP is its adherence to conservation axiom, which prevents the artificial amplification suppression of feature relevance in the backward pass. LRP has been demonstrated to produce faithful explanations across various domains (e.g. [7, 62, 3, 20]). Nevertheless, the peculiarities of the Mamba architecture are not addressed by the existing LRP procedures, which may lead to the violation of the conservation property and result in unreliable explanations.
explanation methods [1, 18]. Yet, the explanations produced
In this work, we present MambaLRP, a novel approach to integrate LRP into the Mamba architecture. By examining the relevance propagation process across Mamba layers through the lens of conserva- tion, we pinpoint layers within the Mamba architecture that need to be addressed specifically. We propose a novel relevance propagation strategy for these layers, grounded in the conservation axiom, that is theoretically sound, straightforward to implement and computationally efficient. Through number of quantitative evaluations, we show that the proposed Mambal.RP approach allows to ro- bustly deliver the desired high explanatory performance, exceeding by far the performance of various baseline explanation methods as well as a naive transposition of LRP to the Mamba architecture. We further demonstrate the usefulness of MambaLRP in several areas: gaining concrete insights into the model's prediction mechanism, uncovering undesired decision strategies in image classification, identifying gender bias in language models, and analyzing the long-range capabilities of Mamba. Our code is publicly available.'
2 Related Work
Structured State Space Sequence Models (SSMs). Transformers [69] have emerged as the most widely used architectures for sequence modeling. However, their computational limitations, partic- ularly with large sequence lengths, have restricted their applicability in modeling long sequences. Addressing these computational limitations, recent works [34, 35] have introduced structured state

---

space sequence models (SSMs) as an alternative approach. SSMs are a class of sequence model- ing methods, leveraging the strengths of recurrent, convolutional, and continuous-ting methods, demonstrating promising performance across various domains, including language [30, 46], im- age [77. 13, 50], and video [71] processing, and beyond [59, 22, 42]. A recent advancement by Gu and Dao [33] introduced selective SSM, an enhanced data-dependent SSM with a selection mechanism that adjusts its parameters based on the input. Built on this dynamic selection, the Mamba architecture fuses the SSM components with multilayer perceptron (MLP) blocks. This fusion simplifies the architecture while improving its ability to handle various sequence modeling tasks, including applications in language processing [6, 53, 72], computer vision [41, 83, 78], medical imaging [44, 76, 31, 56, 40, 74, 73), and graphs [70, 14]. This fast adoption of SSMs and Mamba models underscores the need for reliable explanations of their predictions.
Explainable AI and SSMs. In efforts to explain Mamba models, [51] analyzed if the interpretability tools originally designed for Transformers can also be effectively applied to architectures such as Mamba. In this context, Ali et al. [4] and Zimerman et al. [84] recently proposed viewing the internal computations of Mamba models previous works that use attention signal as explanation, including Attention Rollout [1] and variants thereof [18, 17]. While these approaches can provide some insight, they inherit the limitations of using attention as an explanation [75, 38), including their inability to capture potential misalignment between tokens and attention scores, and the limited performance in empirical faithfulness evaluations. Alternative Explainable AI methods, not yet applied to Mamba models but in principle applicable to any model, include techniques using input perturbations [81, 85, 29] or leveraging gradient information [11, 64, 68, 65, 63). Despite their wide applicability, these models have certain drawbacks, such as requiring multiple function evaluations for single explanation or being susceptible to gradient noise, resulting in subpar performance, as our benchmark experiment will demonstrate. Alternatively, deriving tailored approaches that reflect the underlying model structure, has shown be a promising direction in developing better attribution methods based on gradient analysis of the prediction function [7. 27, 62, 3). In the Layer-wise Relevance Propagation framework, this necessitates suitable propagation rules, which are currently lacking for the Mamba architecture. To tackle these challenges, we introduce Mambal as an efficient solution for the computation of reliable and faithful explanations that are theoretically grounded in the axiom of relevance conservation.
attention mechanism. This approach builds upon
3 Background
Before delving into the details of our proposed method, we begin with a brief overview of the selective SSM architecture, followed by an introduction to the LRP framework.
Selective SSMs (S6) An important component within the Mamba [33] architecture is the selective SSM. It is characterized by parameters, A. B. and C, and transforms given input sequence (rz)T= into an output sequence of the same size (y))T= via the following equations:
h A Bara
(1)
yt- = C,h,
(2)
where the initial state ha O. What distinguishes the selective SSM from the original SSM (S4) [35] is that the evolution parameter, Ar, and projection parameters, B, and Ct, are functions of the input Tt. This enables dynamic adaptation of the SSM's parameters based on input. This dynamicity facilitates focusing on relevant information while ignoring irrelevant details when processing sequence.
Layer-wise Relevance Propagation Layer-wise Relevance Propagation (LRP) [10] is an Explain- able AI method that attributes the model's output to the input features through a single backward pass. This backward pass is specifically designed identify neurons relevant to the prediction. LRP assigns relevance scores to neurons in given layer and then propagates these scores to neurons in the preceding layer. The process continues layer by layer, starting from the network's output and terminating once the input features are reached. The LRP backward pass relies on an axiom called "conservation' requiring that relevance scores are preserved across layers, avoiding to artificially amplify or suppress contributions. For example, let and be the input and output of some layer, respectively, and let R(x) and (y) represent the sum of relevance scores in the respective layers. The conservation axiom requires that R(x) R(y) holds true.
3

---

4
LRP for Mamba
In this work, we bring explainability, particularly LRP, to Mamba models, following the conceptual design steps, shown in Fig. 1. We start by applying a basic LRP procedure, specifically one corresponding Gradient Input (GD,, to the Mamba architecture. This serves an effective initial step for identifying layers where certain desirable explanation properties, like relevance conservation, are violated. We analyze different layers of the Mamba architecture. derive relevance propagation equations and test the fulfillment of the conservation property. Our analysis reveals three components in the Mamba architecture where conservation breaks: the SiLU activation function, the selective SSM, and the multiplicative gating of the SSM's output. Leveraging the analysis above, we propose novel relevance propagation strategies for these three components, which lead robust, faithful and computationally efficient explanation approach, called Mambal.RP.
4.1 Relevance propagation in SiLU layers
We start by examining the relevance propagation through Mamba's SiLU activation functions. This function is represented by the equation y = x ox), where a denotes the logistic sigmoid function.
Proposition 4.1 Applying the standard gradient propagation equations yields the following result, which relates the relevance values before and after the activation layer:
0f
af
0f 8y
y
d(())
(3)
dx
dy
((
R(V)
The derivation for Eq. 3 can be found in Appendix A.I. We observe that the conservation property, iee. Rxx) R(y), violated whenever the residual term non-zero. We propose restore the conservation property in the relevance propagation pass by locally expanding the SiLU activation function as:
= x. [o(x)]
where |est treats the given quantity as constant. This can be implemented e.g. in PyTorch using the .detcch( () function. Repeating the derivation above with this modification yields the desired conservation property. R(x) = R(y). The explicit LRP rule associated to this LRP procedure is provided in Appendix B.
4.2 Relevance propagation in selective SSMs (S6)
The most crucial non-linear component of the Mamba architecture is its selective SSM com- ponent. It is designed to selectively retain or discard information throughout the sequence by adjusting its parameters based on the input, en- abling dynamic adaptation to each token. To facilitate the analysis, we introduce an incon- sequential modification to the original SSM by connecting C, to ht instead of 24. To do so, we can redefine Ac, B. and C, matrices as blockdiag( Ar, 0), (B,, I), and (C, (0) respec- tively, such that Ie becomes part the state he without altering the overall functionality of the SSM
3/-1
U/
C
h
h
D
Ft-1
P
Figure Unfolded view of SSM. highlighting two subsets of nodes, the relevance of which should be conserved throughout relevance propagation.
The unfolded SSM, with the aforementioned modification, is illustrated in Fig. 2. The complex relevance propagation procedure in the SSM component can be further simplified by considering two groups of units, illustrated in red and orange in Fig. 2. In these two groups, there are no connections within units of the same group, all the relevance propagation signals from the first group are directed towards the second group, and the second group receives no further incoming relevance propagation signal. With these properties, these two groups should, according to the principle of conservation, receive the same relevance scores.

---

Proposition 4.2 Defining 02 = (A,, Br, Ct-1). and working out the propagation equations between
these two groups yields the following relation:
0f 1
of
0f
0f 00 0f a0:
Fx, -hp-1
+
Ohe h4 +
dyc-1 D1-1
00, dr,
1t 00,
dhz-1
h4-1
(5)
R(hc) 1)
R 1)
The derivation for Eq. 5 can be found in Appendix A.2. We note that the residual term which is typically non-zero, violates conservation. Specifically, conservation fails due the dependence of on the input. We propose to rewrite the state-space model at each step in way that the parameters 0 appear constant, i.e.:
h₂ [At]est + [Beeaa 5/+ = (Colea.14
(6)
(7)
These equations can also be interpreted as viewing the selective SSM localized non-selective, Le. standard, SSM. With this modification, conservation holds between the two groups, i.e. R(x,) + R(h--1) R(h,) (y)-))) By repeating the argument for each time step, conservation is also maintained between the input and output of the whole SSM component. Explicit LRP rules are provided in Appendix B.
4.3 Relevance propagation in multiplicative gates
In each block within the Mamba architecture, the SSM's output is multiplied by an input-dependent gate. In other words, =A- =B. where 2 SSM(x)) and = = SiLU(Lineard (x)). Assume that the locally linear expansions introduced in Sections 4.1 and 4.2 are applied to the SSM components and SiLU activation functions, the mapping from y becomes quadratic.
Proposition 4.3 Applying the standard gradient propagation equations establishes the following relation between the relevance values before and after the gating operation:
af
0f
af
y
y
Ox
Dy
dy
R(x)
R(y)
The derivation for Eq. 8 and explicit LRP rules can be found in Appendix A.3 and Appendix B, respectively. this equation, we observe a spurious doubling of relevance in the backward pass. This can be addressed by treating half of the output constant:
0.5 (ZA : =B) + 0.5 [RA 2
As for the previous examples, this ensures the conservation property R(x) = R(y). An alternative would have been to make y linear by detaching only one of the terms in the product, as done for the SiLU activation or the SSM component. However, the strategy of Eq. better maintains the directionality given by the gradient. We further compare these alternatives in an ablation study presented in Appendix C.5, demonstrating empirically that our proposed approach performs better.
4.4
Additional modifications and summary
The propagation strategies developed for the Mamba-specific components complement previously proposed approaches for other layers, including propagation through RMSNorm layers [3] and convolution layers via robust LRP-y rules [49, 25] and their generalized variants. A summary of these additional enhancements is provided in Appendix C.2. Furthermore, our proposed propagation rules are generally applicable to other models that utilize similar components, such as multiplicative gates in recent architectures [54, 52, 45, 21].
straightforward implementation of the propagation rules can be achieved by computing MambaLRP via Gradient x Input, where the gradient computations are modified to align with the proposed rules. The procedure consists of two main steps:
1. Perform the detach operations of Eqs. (4), (6). (7), and (9) (as well as similar operations for RMSNorm and convolutions).
2. Retrieve MambaLRP explanations by computing Gradient x Input on the detached model.
5

---

Experiments
To evaluate our proposed approach, we benchmark its effectiveness against various methods previously proposed in the literature for interpreting neural networks. We empirically evaluate our proposed methodology using Mamba-130M, Mamba-1.4B and Mamba-2.8B language models [33], which trained on diverse text datasets. The training details can be found in Appendix C.I. For the vision experiments, we use the Vim-S model [83]. Moreover, we perform several ablation studies to further investigate our proposed method.
Datasets In this study, we perform experiments on four text classification datasets, namely SST-2 [67], Medical BIOS [28]. Emotion [60], and SNLI [16]. The SST-2 dataset encompasses around 70K English movie reviews, categorized into binary classes, representing positive and negative sentiments. The Medical BIOS dataset consists of short biographies (10K) with five specific medical occupations as targets. The SNLI corpus (version 1.0) comprises 570k English sentence pairs, with the labels entailment, contradiction, and neutral, used for the natural language inference (NLI) task. The Emotion dataset (20K) a collection of English tweets, each labeled with one of six basic emotions. For the vision experiments, we use ImageNet dataset [23] with 1.3M images and 1K classes.
specific
Baseline methods We compare our proposed method with several gradient-based model- aggnosicc explanation techniques: Gradient x Input (GI) [11, 64], SmoothGrad [65], and Integrated Gradients [68]. Furthermore, we evaluate the performance of our proposed method against a naive implementa- tion of LRP, i.e. LRP (LN-rule). where the LRP-0 rule is used in all linear and convolution layers, along with the LN-rule [3] normalization layers.
We further compare the performance of our proposed method with two attention- approaches, Attention Rollout (AttnRoll) and MambaAttr [4], which are recently proposed for Mamba models. Both methods are extensions of techniques originally developed for Transformer models: Attention Rollout [1] and Gradient AttentionxAttentioon Rollout [18].
5.1 Conservation property
To verify the fulfillment of the conservation property, on which our method is based, we compare the network's output score with the of relevance scores attributed the input features, for both the GI baseline and the proposed Mambal.RP. The analysis is performed for Mamba- 130M and Vim-$ models trained the SST-2 and ImageNet datasets, respectively. Full conservation is achieved if the output score equals the sum of relevance, as indicated by the blue line in Fig. 3. Our results show that conservation is severely violated by the GI baseline, and is addressed a large extent by MambaLRP, Residual lack of conservation is due to the presence of biases in linear and convolution layers, which are typically non aon-attributable.
a) SST-2
MRmbaLRP
b) ImageNet
Mam
10
Output
Output
3
-30
30
10
I.R(K)
2,R(x)
X,R(x)
X,R(x)
Figure 3: Conservation property. The x-axis represents the sum of relevance scores across the input features and the y-axis shows the network's output score. Each point corresponds to example and its proximity to the blue identity line indicates the extent to which conservation is preserved, with closer alignment suggesting improved conservation.
5.2 Qualitative evaluation
In this section, we qualitatively examine the explanations produced by various explanation methods for Mamba- 130M and Vim-S models. Fig. 4 illustrates the explanations generated to interpret the Mamba- I 33M model's prediction on a sentence from the SST-2 dataset with negative sentiment. We note that all of the explanation methods attribute positive scores to the word 'disgusting', which appears reasonable given the negative sentiment label. However, it is notable that the explanation

---

generated by MambaLRP is more sparse and focuses particularly on the terms "so" and 'disgusting', In contrast, the explanations produced by the gradient- based methods and AttnRoll appear be quite noisy. Furthermore, we show the explanations produced to interpret the Vim-S model's predictions
MambaLRP at least
disgusting
that
may
hard pressed
retain
their
lunch
at least one
Scene
50
disqusting
that viewers may be hard pressed to retain their lunch
SmoothGrad at least one Scene
dlsgusting that viewers may be hard pressed retain their
lunch
one scene
disqusting
viewers may
hard pressed to retain
AttnRoll
at
least
one scene
50
disgusting
thet
viewers
may
be
hard pressed to retain
their
lunch
MambaAttr
disgusting
that
may
pressed
retain
their
lunch
Figure Explanations generated for sentence of the SST-2 dataset. Shades of red represent words that positively influence the model's prediction. Conversely, shades of blue reflect negative contributions. The heatmaps of attention-baut based methods are constrained to non-negative values.
images of the ImageNet dataset in Fig. 5. Purely gradient-based explanations tend to identify unspecific noisy features, while both attention-based approaches, AttRRoll and MambaAttr, are more effective highlighting significant features. Among these methods, MambaLRP stands out for its ability to generate explanations that are particularly focused on key features used by the model to make a prediction. Take, for instance, the first image classified under the 'African elephant' category. We can see that the explanation generated by MambaLRP not only includes all occurrences of the 'African elephant' object but also highlights its distinctive features, such as the tusks. the second image labeled 'wild boar', despite the presence of multiple objects in the image, Mambal explanation remains focused on the 'wild boar' object, disregarding other objects. Moreover, in the third instance, Mambal.RP uncovers a spurious correlation, the presence of watermark in Chinese, influencing the model's prediction, a subtlety overlooked or not fully represented by other methods. Further qualitative results can be found in Appendix C.6.
Image
GI
SmoothGrad
IG
AttnRoll
MambaAttr
MambaLRP
African elephant
wil boar
cerfon
-1.0
-0.5
0.0
0.5
1.0
Ours
Figure 5: Explanations produced by different explanation methods for images of the ImageNet dataset.
AttnRoll and MambaAttr are limited non- non-negative heatmap values.
5.3 Quantitative evaluation
To quantitatively evaluate the faithfulness of explanation methods, we employ an input perturbation approach based on ranking input features by their importance [57], which can be done using either Most Relevant First (MoRF) or Least Relevant First (LeRF) strategy. Ranked features are iteratively perturbed through a process known as flipping We monitor the resulting changes in the output logit, fc, for the predicted class c, and compute the area under the perturbation curve. The areas under the contrast,
curves for LeRF and MoRF strategies are denoted by Af..R and A A respectively. the insertion method starts with fully perturbed input and progressively restores important features. The areas under the curves for this method are indicated by A and A for the MoRF and LeRF strategies, respectively. A reliable explanation method is characterized by low values of
7

---

AMORE or ALeRF, and large values of AfeRF or AMORE. In an effort to minimize the introduction of -ittr--distribttinn manipulations, the recent study by Blücher et al. [15] advocates for harnessing both insights to derive more resilient metric. Therefore, we follow the same strategy as [15, 2] to evaluate explanation methods. The evaluation metrics are defined as A AfeRP - AMORF and AA' AMERE AfoRE. For both metrics, a higher score is preferable, as it signifies a more accurate and reliable explanation method.
The outcomes of this analysis are represented in Table 1. Mambal consistently achieves highest faithfulness scores in comparison to other baseline methods. We observe that GI struggles with noisy attributions, leading to low faithfulness scores. However, methods like Integrated Gradients and MambaAttr have shown improvements in this regard. We note that LRP (LN-rule) outperforms most methods across the majority of the text classification tasks. Nevertheless, its performance is notably inferior compared to MambaLRP. Overall, we observe that MambaLRP significantly outperforms all other methods by a substantial margin. In both vision and NLP experiments, attention-based methods have shown superior performance compared to the purely gradient- based approaches.
Table 1: Evaluating explanation methods. Higher scores A AF indicate more faithful explanations.
SST-2
Med-BIOS
SNLI
Emotion
ImageNet
Methods
Mamba
Mamba
Mamba
Mamba
Mamba
Mamba
Mamba
Mamba
Mamba
Mamba
Mamba
Mamba
Vim-S
130M
1.4B
288B
130M
1.48
2,8B
130M
1.48
2.8B
130M
1:48
2.8B
Random
-0.012
-0.106
0.007
0.044
-0.014
-0.037
0.010
0.002
0.000
-0.001
0.000
0.000
-0.001
GI [64]
0,078
-0.106
-0.043
0.200
-0.634
-1.434
-0.039
-0.039
0.083
-0.787
-0.409
-1.533
-0.018
Smoothhradd [65]
1.377
-0.383
-0.655
1.661
-2.300
0.486
-0.687
00.47
1.808
-1.852
-4,228
0.209
KG [68]
0.857
0.216
0,322
1,296
1.065
1.937
0.453
0.218
0.331
1,808
2,010
4.314
1.217
AtmRoll [4]
0.657
0.431
0.452
2,228
1.076
2.241
0.242
0.371
0.292
0.389
1.483
0.530
2.427
MambaAtir [4]
1.190
0.626
0.341
3.126
3.006
5,326
0.513
0.554
0,343
2.003
4.706
3,849
2.676
LRP (LNNrule, [3])
0,877
0.961
0,820
2.217
3.456
5,305
0.673
0.656
0.731
3.079
5.199
5.094
2.548
MannaaLRP (ours)
1.978
3,906
4.234
0.989
0.897
3.523
5.397
4.715
Runtime comparison We report the runtimes of MambaL along with other methods used in this study in Appendix C.9. As shown in Table 10, our method's 's runtime is comparable to GI and can be implemented via a single forward and backward pass. Since approaches like Integrated Gradients require multiple function evaluations, their runtimes are considerably higher than MambaLRP.
Ablation study In Section 4, we proposed techniques for handling different non-linear components within the Mamba architecture. This ablation study aims to assess the significance of cach technique by testing the effect of their exclusion on faithfulness. Table 2 shows that all three modifications are essential for achieving competitive explanation performance, with our proposed method for handling the SSM component being the most critical. Further experiments, comparing different strategies for handling the Mamba block's multiplicative gate, are detailed in Appendix C.5.
Table 2: Analyzing the impact of ablating the three proposed propagation rules on AA" for the components in MambaLRP.
Table 3: Frequency of gendered words in expla- nations for 'Nurse" and 'Surgeon' classes of the Medical BIOS dataset across language models.
SiLU
SSM
Gate
SST-2
ImageNet
Models
Surgeon
Nurse
x
0.577
0.144
GPT2-base
0.14
0.24
1.721
4.022
TS-base
0.10
0.11
x
4
1.943
4.618
RoBERTa-base
0.01
0.06
1.978
4.715
Mamba- 130M
0.009
0.058
Mamba-1.44
0.001
0.042
6 Use cases
Uncovering gender bias in Mamba. Explanation methods serve as tools to uncover biases in pretrained vision and language models. Using our proposed method, we examine Mamba- Namba-130M and Mamba-1.44 models, trained on the Medical BIOS dataset, to investigate the potential presence of gender biases. Following the methodology in [28], we use Mambal_R2 to identify the top-5 tokens of highest importance and to quantify the prevalence of gendered words within these tokens. We find that the model exhibits pronounced preference for female-gendered words in the 'Nurse' class
8

---

Next generaled loken: [1972]
context: 5775 lakens
Pae
The Good Giet 1977 Anwricn meveris comeey -d drana fim dinenee Bethorr Rose, critonn "Ndd Steoo and stariing Reterd syy russ Muss La Masos Quinn Curre ings end Paul Beaoditt The fikn. prodeced by Ray Sturk cestes odd inio chuncars i struiing sser who he sab ci Maniataa spatiicai froe u frinad, the canan songuat the fricnl is gir foxed. who bas yor beto sbaddoed and her pee ocoss yoong dasgtee.
10
[ss] wene falal allergg
salf- loss sderissss
accoddigg Guilllan''s Kar. secodd
tent pize.
Count
Valeodias Midheel
and dironed
10.
poodace your later sluninn B en Willirs was dircced Charlotz Boon frmm place which sover aine agole stured Valastne and was diociall by lay Send rich.
Brows and Pa audo. The taard
[103; why song writeee compawed by Americn dage soeg wetter Duodd Cures and or glaaty mocndel by wifh ook aroep of which Cato wi loakr prinaay mase poodace J appoared Bread alburs shar Mas
10'-
single
aktng memtbrr
wnng only
pop
reacbed
conseppere
New Zoelood
rrached nonter
10a
(176. Lis mumber
unnacrpamary engloval 1908 (U
(231) tit board prssed the iite made calling 1t ae attoccm rnd wet boll ad shat mom thor ssee reaniess pop
2000
6000
slasics the cntly 10
Pasition difference
1323 Summary
Apeil 19
Figure 6: Analysis of the position of tokens relevant for next token generation. Left: Distribution of absolute position of the ten most relevant tokens for the prediction of the next word. Right: Long-range dependency between tokens of the input and the predicted next token (here: 1972).
Trp of cocument
100 100 100
Next generated (157)
context: 2048 tokeess
100
There important hidklen inside lor irmolevant Find and about the important information there..
will quie you
90
pues key Remember it. 9 the pass key. grass green. blue.
docum depen
yellow Here go. There and back again. mounthin high. snow cold. thin Up ward
Ne postion
20%
Pleyeaace
back again grass green
sun yellow we go. back again.
mountain is high. The snow is
100%
Qu What pass key? The pass key
Botton of document
- corfocx longoh
Figure 7: Explanation-based retrieval accuracy in the needle-in-a-haysttac test verifying model reliance on relevant features for different context lengths.
(e.g. the proportion of gender-specific words is 0.058 for females, compared to 0.0 for males in M3mba-130M.). We also compare the results of our analysis with those achieved for the GPT2-base, T5-base, and RoBERTa-bass models as mentioned in [28]. As shown in Table 3, both Mamba models are less dependent on gendered tokens compared to GPt2-base, T5-base, and RoBERTa-baa models, with the Mamba-1.4B model showing a further decrease in bias compared to the Mamba-130M, suggesting improvements in reducing gender bias with increased model size.
Investigating long-range capabilities of Mamba. The ability of SSMs to model long-range dependencies is considered important improvement over previous sequence models. In this use case, we analyze the extent to which the pretrained Mamba- 130M model can use information from the entire context window. We use the HotporQA [79] subset from the LongBench dataset [12], 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. While observe pronounced pattern of attributing to the last few tokens, as seen in prior language generation studies [80, 61], the extracted explanations also identified relevant tokens across the entire context window, as presented for one example in Fig. 6 (right). This suggests that the model is indeed capable of retrieving long-range dependencies. We clearly see that in order to complete the sentence and assign a year to the album release date, the model analyzes previous occurrences of chronological information and MambaL identifies evidence supporting the decision for the date being '1972' as relevant. Our analysis demonstrates the previously speculated long-range abilities of the Mamba architecture [33], which we further explore in comparison to Transformers in Appendix C.8.
Needle-in-a -naeiii----haaackkke test. To assess the model's ability in retrieving relevant pieces of information from a broader context, we perform the beddl--n--a-hasstade test [47]. Our test involves extracting a single passkey (the "needle") from a collection of repeated noise sentences (the 'haystack'), as described in [37]. We run this test at eleven different document depths with three different context

---

lengths. We use an mnstruetion-ineetueed Mamba-2.8B model in this experiment. To analyze the performance of the model, we introduce the explanation- retrieval accuracy (XRA) metric. In this approach, we first identify the positions of the top-K relevant tokens by Mambal .RP, and then, caleulate the accuracy by comparing those positions to the needle's position. As shown in Fig. 7, Mambal RRP accurately captures the information used by the model to retrieve the needle. In this case, the model could accurately retrieve the needle based on relevant information within the text. However, in more realistic and complex scenarios, the model may depend on irrelevant data yet still generate the correct token. This issue can be analyzed using XRA but cannot be evaluated by conventional retrieval accuracy metrics. Such cases and also further details about this experiment are shown in Appendix C.7.
Discussion and conclusion
Mamba models have emerged as efficient alternative Transformers. However, there limited works addressing their interpretability [4, 84]. To address this issue, we proposed MambaL within the LRP framework, specifically tailored the Mamba architecture and built upon the relevance conservation principle. Our evaluations across various models and datasets confirmed that Mambal adheres to the conservation property and provides faithful explanations that outperform other methods while being more computationally efficient. Moreover, we demonstrated how MambaLRP can help users debug state-of-the-un vision and language models while building trust in their predictions through various use cases. Future research can explore its potential across a broader range of applications and Mamba architectures, providing reliable insights into sequence models.
Limitations As propagation-base explanation method, MambaLRP requires storing activations and gradients, leading to memory usage that depends on the model architecture and input sequence length. To reduce memory consumption, techniques such as gradient checkpointing can be utilized, which are applicable other gradient-basee methods as well. However, a limitation of these methods, including MambaLRP, is the potential inaccessibility of gradient information due to proprietary constraints. such cases, approximating gradient information may offer a viable solution.
Acknowledgments and Disclosure of Funding
This work was funded by the German Ministry for Education and Research (refs. 011S UISS14013A-E 01GQ1115, O1GQ0850, 01151, 031L0007 01IS K.R.M. was partly supported by the Institute of Information & Communications Technology Planning & Evaluation (ITPP) grants funded by the Korea government (MSIT) (No. 2019-0-00079, Artificial Intelligence Graduate School Program, Korea University and No. 2022-0-00984 Development of Artificial Intelligence Technology for Personalized Plug-and and--aay Explanation and Verification of Explanation).
References
[1] S. Abnar and W. Zuidema. Quantifying attention flow in transformers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4190-4197, Online, July 2020. Association for Computational Linguistics.
[2] R. Achtibat. S. M. Hatefi, M. Dreyer, A. Jain, T. Wiegand, S. Lapuschkin, and W. Samek. AttaLRP: Attention-aware layer-wise wise relevance propagation for transformers. arXiv:2402.05600 2024.
[3] A. Ali, T. Schnake, O. Eberle, G. Montavon, K.-R. Müller, and L. Wolf. XAI for transformers: Better explanations through conservative propagation. International Conference on Machine Learning, ICML 2022, volume 162 of Proceedings of Machine Learning Research, pages 435 4511 PMLR, 2022.
[4] A. Ali, I. Zimerman, and I.. Wolf. The hidden attention of mamha models. arXiv:2403.01590 2024.
[5] C. An, F. Huang, J. Zhang, S. Gong, X. Qiu, C. Zhou, and L. Kong. Training-free long-context scaling of large language models. In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver. J. Scarlett, and F. Berkenkamp. editors, Proceedings of the 41st International Conference on Machine Learning, volume
235 of Proceedings of Machine Learning Research, pages 1493-1510. PMLR, 21-27 Jul 2024. URL https://procoediig mlr. pprsss/v333 html.
[6] Q. Anthony, Y. Tokpanov, P. Glorioso, and B. Millidge. BlackMamba: Mixture of experts for stace models. 4rXiv:2402.01771, 2024.
10

---

[7] L. Arras, J. Arjona-Medina. M. Widrich, G. Montavon, M. Gillhofer, K.-R. Müller, S. Hochreiter. and W. Samek. Explaining and interpreting LSTMs. Explainable AI: Interpreting, explaining and visualizing deep learning, pages 21 211-238. 2019.
[8] A. B. Arricta, N. D. Rodríguez, J. D. Ser, A. Bennetot, S. Tabik, A. Barbado, S. García, S. Gil-Lopez, D. Molina, R. Benjamins, R. Chatila, and F. Herrera. Explainable artificial intelligence (XAI): concepts, taxonomies, opportunities and challenges toward responsible AI. Inf. Fusion, 58:82-115, 2020.
[9] Z. Avsec. V. Agarwal, D. Visentin, J. R. Ledsam, A. Grabska-Barwinska. K. R. Taylor, Y. Assael, J. Jumper. P. Kohli, and D. R. Kelley, Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 18(10):1196 12006,1112222222 2021.
[10] S. Bach, A. Binder, G. Montavon, F. Klauschen, K.-.. Müller, and W. Samek. On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation. PloS one, 10(7)e01 30140, 2015.
[11] D. Baehrens, T. Schroeter, S. Harmeling, M. Kawanabe, K. Hansen, and K.-R. Müller. How to explain individual classification decisions. The Journal of Machine Learning Research, 11:1803- 11118001181,,,, 2010.
[12] Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang. Z. Huang. Z. Du, X. Liu, A. Zeng. L. Hou, Y. Dong. J. Tang, and J. Li. LongRench: A bilingual, multitask benchmark for long context understanding. arXiv:2308. 14508, 2023.
[13] E. Baron, I. Zimerman, and L. Wolf. 2-D SSM: A general spatial layer for visual transformers. MiXiv:2306.06635 2023.
[14] A. Behrouz and F. Hashemi. Graph Mamba: Towards learning graphs with state space models.
arXiv:2402.08677 0867788.8888, 2024.
[15] S. Blücher, J. Vielhaben, and N. Strodthoff. Decoupling pixel flipping and occlusion strategy for consistent XAI benchmarks. arXiv:2401. 06654, 2024.
[16] S. R. Bowman, G. Angeli, C. Potts, and C. D. Manning. A large annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference Empirical Methods in Natural Language Processing, pages 632-642, Lisbon, Portugal, Sept. 2015. Association for Computational Linguistics.
[17] H. Chefer, S. Gur, and L. Wolf. Generie attention-model explainability for interpreting bi-modal and encoder-decoder transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 397-406, 2021.
[18] H. Chefer, S. Gur, and L. Wolf. Transformer interpretability beyond attention visualization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 782-791, 2021..
[19] G. Chen, X. Li, 7. Meng, S. Liang, and L., Bing. CLEX: Continuous length extrapolation for large language models. In The Twelfth International Conference on Learning Representations, 2024. URL attps://openrerier net/ weot//foommiiee weweummdddxxxSSiPpe...
[20] P. Chormai, J. Herrmann, K.-R. Müller, and G. Montavon. Disentangled explanations of neural network predictions by finding relevant subspaces. IEEE Trans. Pattern Anal. Mach. Intell., 2022.
[21] T. Dao and A. Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp, editors. Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 10041-1007/ PMLR, 21-27 Jul 2024. URL mlr, prwsss/eevve.. 2355 dt///d.2224aat...a
[22] S. B. David, I. Zimerman, E. Nachmani, and L.. Wolf. Decision S4: Efficient sequence- based rl via state spaces layers. In The Eleventh International Conference on Learning Representations, 2022.
[23] Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In IEEE Computer Vision and Pattern Recognition (CVPR), pages 248-22 225,22552222 2009.
[24] Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics, pages 4171 4181-418,, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[25] A.-K. Dombrowski, C. J. Anders, K.R. Müller, and P. Kessel. Towards robust explanations for deep neural networks. Pattern Recognition, 121:108194, 2022. ISSN 0031-3203.
[26] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenbomn, X. Zhai, T. Unterthiner, M. Dehghani, M. Min- derer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[27] O. Eberle, J. Büttner, F. Kräutli, R. Müller, M. Valleriani, and G. Montavon. Building and interpreting deep similarity models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(3):1149 1161, 2020.

---

[28] O. Eberle, I. Chalkidis, L. Cabello, and S. Brandl. Rather nurse than a physician - contrastive explanations under investigation. In Conference on Empirical Methods in Natural Language Processing (EMNLP). pages 6907-6920. Association for Computational Linguistics, 2023.
[29] R. C. Fong and A. Vedaldi. Interpretable explanations of black boxes by meaningful perturbation. In 2017 IEEE International Conference on Computer Vision (ICCV), pages 3449 3457, 2017.
[30] D. Y. Fu, T. Dao, K. K. Saab, A. W. Thomas, A. Rudra, and C. Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv:2212. 14052, 2022.
[31] H. Gong, L. Kang. Y. Wang, X. Wan, and H. Li. nnMamba: 3D biomedical image segmentation,
classification and landmark detection with state space model. arXiv:2402.03526 2024.
[32] Y. Gong, Y.-A. Chung, and J. Glass. AST: Audio Spectrogram Transformer. In Proc. Interspecch 2021, pages 571-575, 2021.
[33] A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv:2312.00775. 2023.
[34] A. Gu, Johnson, Goel, K. Saab, T. Dao, Rudra, and C. Ré. Combining recurrent, convolutional, and continuous time models with linear state space layers. Advances in Neural Information Processing Systems, 34:572-585, 2021.
[35] A. Gu, K. Goel, and C. Re. Efficiently modeling long sequenees with structured state spaces. In International Conference on Learning Representations, 2022.
[36] D. Gunning. DARPA's explainable artificial intelligence (XAI) program. In Proceedings of the 24th International Conference on Intelligent User Interfaces, IUI '19, page ii Association for Computing Machinery, 2019.
[37] C.-P. Hsieh, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, Y. Zhang, and B. Ginsburg. RULER: What's the real context size of your long-context language models? arXiv:2404.00 06654, 2024.
[38] S. Jain and B. C. Wallace. Attention is not Explanation. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics, pages 3543-3556, Minncapolis, Minnesota, June 2019. Association for Computational Linguisties.
[39] S. Lapuschkin, S. Wäldchen, A. Binder, G. Montavon, W. Samek, and K.-R. Müller. Unmasking clever hans predictors and assessing what machines really learn. Nature communications, 10(1): 2019.
[40] Liu, H. Yang, H. Zhou, Y. Xi, Yu, Y. Yu, Y. Liang, G. S. Zhang, H. Zheng, et al. Swin-UMamba: Mamba-based unet with imagenet- pretraining, arXiv:2402.0m 033002 2024.
[41] Y. Liu, Y. Tian, Y. Zhao, H. Yu, L.. Xie, Y. Wang, Q. Ye, and Y. Liu. VMamba: Visual state space model. arXiv:2401.10166 2024.
[42] C. Lu, Y. Schroecker, A. Gu, E. Parisotto, J. Foerster, S. Singh, and F. Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 36. 2024.
[43] J. Ma, F. Li, and B. Wang. U-mamba: Enhancing long-range dependency for biomedical image segmenta- tion. arXiv:2401.0e4 04472, 2024.
[44] J. Ma, F. Li, and B. Wang. U-mamba: Enhancing long-range dependency for biomedical image segmenta- tion. arXiv preprint arXiv:2401.0472 2024.
[45] X. Ma. C. Zhouo, X. Kong, J. He, L. Gui, G. Neubig, J. May, and L. Zettlemoyer. Mega: Moving average equipped gated attention. In The Eleventh International Conference on Learning Representations, 2023. URL httpps:///ppt httpa/////ppeeeeeeeeeee net.iiewweeeeeeeee.eeeeeeeeeeeee wttp:/rrrrmdddleeeei...03.
[46] H. Mehta, A. Gupta. A. Cutkosky, and B. Neyshabur. Long range language modeling via gated state spaces.
arXiv:2206.1399 13947, 2022.
[47] A. Mohtashami and M. Jaggi. Random- access infinite context length for transformers. In Advances in Newral Information Processing Systems, 2023.
[48] G. Montavon, W. Samek, and K.-R. Müller. Methods for interpreting and understanding deep neural networks. Digital signal processing, 73:1-15, 2018
[49] G. Montavon, A. Binder, S. Lapuschkin, W. Samek, and K.-R. Müller. Layer-wise relevance propagation: An overview. Explainable Al: interpretiin, explaining and visualizing deep learning, pages 193-209, 2019.
[50] E. Nguyen, K. Goel, A. Gu, G. Downs, P. Shah, T. Dao, S. Baccus, and C. Ré. S4nd: Modeling images and videos as multidimensional signals with state spaces. Advances in Neural Information Processing Systems, 35:2846 1522446-2661, 2022.
[51] G. Paulo, T. Marshall, and N. Belrose. Does transformer interpretability transfer to rnns?, 2024. URL https: //tpx//./vvv///////////..///wl obg/abs/ /204 05971
12

---

[52] B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao. X. Cheng, M. Chung.
L. Derezynski, X. Du, M. Grella, K. Gv, X. He, II. Hou, P. Kazienko, J. Kocon, J. Kong. B. Koptyra. Mantri, E. Mom, A. Saito, G. Song, X. Tang, Wind. S. Wožmiak, Z. Zhang. Q. Zhou, J. Zhu, and R.-J. Zhu. RWKV: Reinventing RNNs for the transformer era. In H. Bouamor, J. Pino, and K. Bali, editors, Findings of the Association for Computational Linguisties: EMNLP 2023, pages 14048-14077, Singapore, Dec. 2023. Association for Computational Linguisties. doi: 10.18653/v1/ 2023. 1023.fnddigs-emmll..0.
H. Lau, J. Lin, K.
URL httpstttppp/... / hatp//aaaaattoogg.....oo http://ggg...//2/....ionggggege. Fiidings--mmll.-0
[53] M. Pióro, K. Ciebiera, K. Król, J. Ludziejewski, and S. Jaszczur. MoE-Mamba: Efficient selective state space models with mixture of experts. arXiv:2401.4408/ 2024.
[54] Z. Qin. S. Yang, and Y. Zhong. Hierarchically gated recurrent neural network for sequence modeling. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL httpss//pppeeereeiien het/forum?id-Poe PiC.iddddPTC....BB........
[55] N. Rajani, L. Tunstall. E. Beeching, Lambert, A. M. Rush, and T. Wolf. No robots. https:// huggingface.. wwwpgggfaeeccoeeeeeeeee aaataasts//uggingaaaettttt htcl//ooorroee /tobotsssoooooooeeoeo/e 2023.
[56] Ruan and S. Xiang. VM-UNet: Vision mamba UNet for medical image segmentation. arXiv:2402.0249), 2024.
[57] Samek, A. Binder, G. Montavon, S. Lapuschkin, and K.-R. Müller. Evaluating the visualization of what deep neural network has learned. IEEE Transactions on Neural Networks and Learning Systems, 28(1 I): 2660-2673, 2017.
[58] W. Samek, G. Montavon, S. Lapuschkin, C. J. Anders, and K.-R. Müller. Explaining deep neural networks and beyond: A review of methods and applications. Proc. IEFE, 109(3):247-278 2021.
[59] G. Saon, A. Gupta, and X. Cui. Diagonal state space augmented transformers for speech recognition. In ICASSP 2023-2023 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pages 1-5. IEBE, 2023.
[60] E. Saravia, H. T. Liu, Y. Huang. J. Wu, and Y. Chen. CARER: contextualized affect representations for emotion recognition. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 November 4, 2018, pages 3687-3697 Association for Computational Linguistics, 2018.
[61] G. Sarti, N. Feldhus, L. Sickert, O. van der Wal, M. Nissim, and A. Bisazza. Inseq: An interpretability toolkit for sequence generation models. In Proceedings of the ó1 Annual Meeting of the Association for Computational Linguistics (Volume 3: System Demonstrations), pages 421 435, Toronto, Canada, July 2023. Association for Computational I. Linguisties.
[62] T. Schnake, O. Eberle, J. Lederer, S. Nakajima, K. T. Schüttt, K.-R. Müller, and G. Montavon. Higher-order explanations of graph neural networks via relevant walks. IEFF Transactions on Pattern Analysis and Machine Intelligence, 44(11) (1:1)81115999 2022.
[63] R. R. Selvaraju, M. Cogswell, A. Das, R. Vedantam, Parikh, and Batra. Grad-CAM: Visual explanations from deep networks via gradient-based localization. In 2017 IEEE International Conference on Computer Vision (ICCV), pages
2017.
[64] A. Shrikumar, P. Greenside, and A. Kundaje. Learning important features through propagating activation differences. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML' page
2017.
[65] D. Smilkov, N. Thorat, B. Kim, F. B. Viégas, and M. Wattenberg. SmoothGrad: removing noise by adding noise. arXiv: 03825. 2017.
[66] J. T. Smith, A. Warrington, and S. Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.
[67] R. Sncher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Y. Ng, and C. Potts. Recursive deep models for semantic compositionality over sentiment treebank. In Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 1631-1642. ACL, 2013.
[68] M. Sundararajan, Taly, and Q. Yan. Axiomatic attribution for deep networks. In Proceedings of the 34th International Conference on Machine Learning, 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 3319-3328. PMLR, 2017.
[69] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkorcit, L. Jones, A. N. Gomez, L. u. Kaiser, and Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[70] C. Wang. O. Tsepa, J. Ma. and B. Wang. Sraph-Mamin Towards longerange graph sequence modeling with selective state spaces. arXiv:2402.00789, 2024.
13

---

[71] J. Wang, W. Zhu, P. Wang. X. Yu, L. Liu, M. Omar, and R. Hamid. Selective structured state-spaces for long- fong-form/oe video understanding. In Proceedings the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 6387-63977 2023.
[72] J. Wang. T. Gangavarapu. J. N. Yan, and A. M. Rush. MambaByte: Token-free selective state space model.
1r360 2024.
773 Z. Wang and C. Ma. Semi-Mamba- UNet: Pixel-level contrastive cross-supervised visual mamba-based unet for semi-supervised medical image segmentation. arXiv:2402.07245, 2024.
[74] Z. Wang, J.-Q. Zheng, Y. Zhang. Cui, and L. Li. Mamba-UNer UNet-like pure visual mamba for medical image segmentation. arXiv:2402.0507 2024.
[75] S. Wiegreffe and Y. Pinter. Attention is not explanation. In Proceedings the 2019 Conference Empirical Methods Natural Language Processing and the 9th International Joint Conference on Natural I anguage Processing (EMNIP-IJCNIP)) pages 11-20, Hong Kong, China, 2019. Association for Computational Linguistics.
[76] Z. Xing, T. Ye, Y. Yang. G. Liu, and Zhu. SegMamba: Long-range sequential modeling mamba for 3d medical image segmentation. arXiv:2401.13566 2024.
[77] J. N. Yan, J. Gu, and A. M. Rush. Diffusion models without attention. arXiv:2311.18233 2023.
[78] Y. Yang. Xing. and Zhu. Vivim: video vision mamba for medical video object segmentation.
1rXiv22401.41668, 2024.
[79] Z. Yang. Qi. S. Zhang, Y. Bengio, W. W. Cohen, R. Salakhutdinov, and D. Manning, HotpotQA: dataset for diverse, explainable multi-hop question answering. Conference on Empirical Methods in Natural Language Processing (EMNLP). 2018.
[80] K. Yin and G. Neubig. Interpreting language models with contrastive explanations. In Proceedings of the 2022 Conference on Empirical Methods Natural Language Processing, pages 184-198, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics.
[81] M. D. Zeiler and R. Fergus. Visualizing and understanding convolutional networks. In Computer Vision- ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part I 13, pages 818-833. Springer, 2014.
[82] H. Zhou, S. Zhang, J. Peng, S. Zhang, J. Li, H. Xiong, and W. Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. In The Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Virtual Conference, volume 35, pages 11106-J 1166-11115..,.1111111111111 AAAI Press, 2021.
[83] L. Zhu, B. Liao, Q. Zhang, X. Wang. W. Liu, and X. Wang. Vision Mamba: Efficient visual representation learning with bidirectional state space model. arXiv:2401.094 09917, 2024.
[84] I. Zimerman, A. Ali, and L. Wolf. A unified implicit attention formulation for gated-linear recurrent sequence models, 2024. URL pttps://arxils org/oogg//e......oo...... //s////////// /a/////b///200/ 2205/ 16504.
[85] L. M. Zintgraf. Cohen. T. Adel, M. Welling. Visualizing deep neural network decisions: Prediction difference analysis. In International Conference on Learning Representations. 2017.
14

---

A Proofs
In the following, we provide derivations for the conservation analysis performed in Section 4.
A.1 Derivations for SiLU
We first consider the SiLU activation function. As mentioned in Section 4.1, this function is represented by the equation y = x (2), with a being the logistic sigmoid function. By applying the standard gradient propagation equations, get the conservation equation:
R(a)
0f I dy dx
0f dy
dz
Os (o(x) xo (z) x
(10)
0s o (x) x
05
xo (x) [x)(
dy
dy
0s
01
dy +
dx)x) x²
8y
R())
A.2 Derivations for selective SSM
In Section 4.2, we introduced an inconsequential modification to the original selective SSM architec- ture by connecting the matrix C to the state h, instead of the input Fg. The unfolded view of the SSM component with this modification is represented in Fig. 2. We can observe two subsets of nodes in this figure. The relevance scores of these two subsets should be equal if the conservation property holds. Computing the relevance propagation equation between these two groups, we obtain:
R(ze)+((bee)))
05
0f h--1=
Of 0+ h+
0f 00.
D: + thi-1
2x + 00, dx: 21
dz,
Th, dx,
0f 0 atyy-1 h1-1
0f Th, h-11
Of 00:
h(-1
dyt-1 Shy-1
Th, Th,-1
00, th:-1
Os -+ hr
05 The lu-
05 8 yt- 11-11)
(11)
Thy dz,
θhy Thi-1
0tji-1 δhy-1
0f 00,
0f 00,
11--1)
+
00, Oe
00, Jh4-1
0f h4 +
0f
9 + 0f 00,
0f 00: h
Th Oy--1
00, Oz
00,
Oh,-1
R((he) R(ve-1)
A.3 Derivations for multiplicative gate
Mamba is composed of several blocks. In each block, the selective SSM's output is multiplied by an input-dependent gate. other words, y = zA=B with = = SSM(2) and žB = ()ine()))a))())). By applying the standard gradient propagation equations, get the conservation equation:
15

---

R(a)
05 x
05 dy
dy dx x
Or
05 ŌzA
8zB)
2B 2A Oz
dy
Or
(12)
0f (ZARB + =22B))
dy
0f
0f
Oy Oy 1/
1/
R(y)
B Explicit propagation rules for MambaLRP
Whereas Mambal.PP is more easily implemented via the modified gradient-based approach described the main paper, we provide below explicit relevance propagation equations for better comparability with other works. We refer to Sections and of the main paper for the definition of the notation.
B.1 SiLU
Explicit LRP rule for SiLU layers:
R(x;)= = R(()))
(13)
B.2
SSM
Using the shortcut notations a, = [AA(=))); byj =
[Be(x)) and cy (Cr-- we can
write the propagation of relevance the previous state space activations explicitly as:
R(h Σ
h R(y, 1--1)
((-1)
Σ
(1-1)
aij
10 R(h1))
(14)
Caj
E.
11- [t-
Oij + Ex
Fxp
and the propagation of relevance the SSM input as:
24 big
R(x!")) = Σ E, x beg
R(h()))
(15)
+ Σx n1
(1-1) 02j
B.3 Multiplicative Gate
Explicit LRP rule for the multiplicative gate:
R([=a]k = (.5 R([=n]) = 0.5 R()))
(16)
(17)
C Experimental details
In this section, we provide experimental details on our experiments that allow reproducibility of our results.
C.1 Models and datasets
For the NLP experiments, we fine-tuned all parameters the pretrained Mamba- 30M, Mamba-1.4B, and Mamba-2.8P models² on four text classification datasets: SST-2, SNLI, Medical BIOS, and Emotion. The data statisties can be seen in Table 5. For the vision experiments, we used the pretrained Vim-S model?, trained on the ImageNet dataset.
http:////....ttceccccc/attttttes hhttps:/ggtthub http://hhhhb.cccuuuus////aaaa
16

---

Training details During training, we used batch size of 32. To train the Hamba-1,44p and Mamba- 2.8B models on the SNLI dataset, batch size of is used. We employed the (EleutherAl/gpt-neee 20b}* tokenizer. The models' parameters were optimized using AdamW optimizer with a learning rate set at 7e 5. Additionally, we used linear learning rate scheduler with an initial factor of 0.5. All models were trained for maximum of 10 epochs. We employed early stopping and ended training as soon as the validation loss ceased to improve. The top-1 accuracies of the models on each dataset are detailed in Table 4.
Table 4: The accuracies Mamba-I 13M, Mamba- 1.4B, and Mamba-2.88 models on the validation sets four text classification datasets.
Table 5: Data statistics.
Dataset
Train
Test
Validation
Dataset
130M
1,4B
2.8B
SST-2
68K
2K
1K
SST-2
91.97
94.15
94.26
Mcd-BIOS
8K
1K
1K
Med-BIOS
89.10
90.30
90.50
Emotion
16K
2K
2K
Emotion
93.45
93.65
94.10
SNLI
550K
10K
10K
SNLI
89.57
91.05
91.14
ImageNet
1.3M
50K
100K
C.2 MambaLRP details
In this section, we begin by showing how Mambal.RP can be implemented through the following algorithms. Then, explain the generalized LRP-y rule, provide details regarding hyperparameters used our implementation, and outline the hyperparameter selection procedure.
Algorithm 1: Mambal in SiLU activation layer
Data: Input: x LB,
& (a)) return z @ (Si).(()) @ 2]. detachh)
Algorithm 2: MambaLRP in Mamba block
Data: Input: x (B, L, D)
Data: Output: y (B, L.
x': (B, L, B) SiLU(ConvId(z))
g: (B, L, E) <- SiL.U(Linear(z))
9 is an input-dependent gate
A: (E, N) < Parameter
B: (B, L, N) < Linear(r" (x)e
C: (B, L, N)
Linear(at
C input
A: (B,
L, E) + Coftplus(Parameter + Linear(z'))
and B are input-dependent YSaN: (B, L. B) SSM(AAdetache detach(), B detach(), C. ())
A, B: (B, I., F, N) <- discretize(A, A, B)
y/: (B, L, B) - 0.5(ysaM 3 g) + 0.5[yasM o detach
y: (B, L, D) < Linear(y')
11 return y
The following list represents the hyperparameters of the above-mentioned algorithms:
batch size
L
sequence length
hidden dimension
E
expanded hidden dimension
N
SSM dimension
hattp:///gttabbb.oe//leettttt///
17

---

Explanations generated by propagation-baser methods rely on gradient computations, which can result in noisy explanations models with many layers. This due the phenomena of gradient shattering and the presence of noisy gradients, which are more common in deep complex models [25, 2]. To mitigate this, we apply the generalized LRP->-rule to the convolution layers of the Vision Mamba model to improve the signal to noise ratio, thereby enhancing explanations. The generalized LRP-y rule is defined in Eq. 18:
E, # (Wiey + (o ) +2, (wey+ (
(wiy mwm (wss, + R(ys)
if 23
>
R(x;)
(18)
E, x: (uusstwwe, 1 +2 +
(yj))
else
Σ, ( + +w +2 (0ns) + +
where (.)* max(0.)) and (.) min(O, (0, .). and 2j Σ Wy In our experiments, the parameter y is set 0.25. Our observations reveal that applying this rule the language models does not lead to any discemibbe improvements. Therefore, we use the LRP-0 rule in these models.
C.2.1 LRP composites for Vision Mamba
As mentioned Section C.2, we apply the generalized LRP-y rule the convolution layers of the Vim-S model to produce more faithful explanations. In this experiment, we justify this choice. Vision Mamba is composed of number of blocks and in each block, there are several linear and convolution layers, where the generalized LRP-y rule can be used. As can be seen in Table 6, the LRP-0 rule is sufficient to produce meaningful explanations. However, we can perform hyperparameter search by applying the LRP-y rule across different layers of the model to find the most accurate LRP composite.
Table 6: Finding the best LRP composite for Vision Mamba. The layers in which the generalized LRP-y rule applied are represented with LRP-y and the ones in which the basic LRP rule, i.e. LRP-0, is used represented with LRP-0.
ImageNet (AAf 7)
in-proj
out-proj
convid
LRP-Y
LRP-Y
LRP-y
4.196
LRP-y
IRP-Y
LRP-0
4.218
LRP-0
LRP-7
LRP-y
4.283
LRP-0
IRP--
IRP-0
4.336
LRP-Y
LRP-0
LRP-7
4.597
LRP-Y
LRP-0
LRP-0
4.599
LRP-0
LRP-0
LRP-0
4.684
LRP-0
LRP-0
LRP-y
4.715
We apply the LRP-y rule across different combinations of the input projection (in-proj), output projection (out-proj), and convolution layers of each block. Subsequently, we perform the perturbation experiment to analyze the faithfulness of each combination. We can observe that the best result can be achieved when the IRP-y rule is only used in convolution layers. In all of these combinations, the value of is set to 0.25.
C.3 Further details of other explanation methods
Some of the explanation methods that used in this study have set of hyperparameters. Table 7 provides further details on the specific values assigned to these hyperparameters, chosen based on the values suggested in the original papers [68, 65].
In the vision experiments, we used the original implementations of the AunRoll and MambaAttr methods, provided to explain the Vim-S model. Given the unavailability of code for adapting these approaches the language models, namely Mamba- 130M, Mamba-1.44 1.4B, and Mamba-2.8B. we have developed our own implementation. In the vision case, the authors obtain the final relevance map by
htttps://github http:///ww..aaaaa//aaaaaaaaaaaaa
18

---

Table Hyperparameters other explanation methods. The parameters y and represent the mean and standard deviation of noise, respectively, while the parameter m denotes the sample size.
Method
Hyperparameters 0, a 0.15, m 30
SmoothGrad
Integrated Gradients
m
Table 8: Comparing the proposed strategies for managing the Mamba block's multiplicative gate.
Strategies
SST-2
ImageNet
Detaching the multipiccative gale Half-relevunce propagation
1.577
3.592
1.978
4.715
extracting the row associated with the CLS token in the attention matrix. However, since our language models lack CLS token, we get the final relevance map from the associated with the last token in the attention matrix. This is because predictions are based on the last state in these models. For the gradient-based methods, use the implementations available in the Captum library
C.4 Further details evaluation metrics
As mentioned in Section 5.3, the flippig and insertion metrics can be used to evaluate the quality of the generated explanations. It is important to note that starting with unperturbed images and gradually applying perturbations until fully perturbed images are obtained yields identical results for both the flipping and insertion metrics. Therefore, we have only reported the results for the flipping experiment.
In our flipping evaluations, we calculate the area under the curve (AUC) by starting from full images and progressively masking the pixels with zeros until we reach completely masked images. The perturbation steps defined using np wwl.iissacee(. (0, 1, 11) In vision experiments, Images are normalized using the ImageNet mean and standard deviation, and the explanations are generated and evaluated for the predicted class. Unlike [4], which tracks changes in the model's top-1 accuracy, we monitor changes in the output logit of the predicted class.
C,5 Further ablation experiments
Comparing strategies for managing the Mamba block's multiplicative gate: In Section 4.3, we proposed several strategies to mitigate conservation violation the Mamba block's multiplicative gate. In this experiment, evaluate the proposed approaches. As can be in Table 8, detaching the multiplicative gate 2g leads to lower faithfulness scores compared to the half-relevance propagation approach. To retain conservation, limits capturing long-range dependencies, a task for which this branch is designed for. Detaching may result in loss of valuable information used by the model to make predictions. Therefore, this approach is not considered in Table 8.
alternative approach is detach the SSM's output 2,, which
C.6 Additional qualitative results
In Section 5.2, we qualitatively evaluated the explanations produced by Mambal and other bascline methods. In the following, we demonstrate further qualitative results.
C.6.1 Natural language processing
In the following figures, we represent explanations produced by Mambal and other baseline methods to interpret the Mamba- 130M models trained on various datasets. the visualizations, shades of red represent words that positively influence the model's prediction. Conversely, shades of blue reflect negative contributions. The heatmaps of the AttnRoll and MambaAttr methods are constrained to non-negative values.
'https://cappe ai/
19

---

MambaLRP
Two women are embracing while holding to packages Two woman are holding packages
GI
Two women are embracing while holdinn go packages
Two woman are holding packages
SmoothGrad
Two women are embracing while holding to go packages
Two woman are holding packages
IG
Two women embracing while holding to packages Two woman are holding packages
AttnRoll
Two women
embracing while holding to packages Two woman are holding packages
MambaAttr
Two women embracing while holding to packages
Two woman are holding packages
Figure Explanations generated by different explanation methods for sentence of the SNLI validation set. This sentence belongs to the 'entailment' class.
A woman Is doing cart
wheel
while wearing
the sand next to the beach
MambaLRP
woman is fixing her home
woman doing
wheel wearing
sand next the beach
GI
woman fixing her home
SmoothGrad woman doing
wheel while wearing ini sand next the beach
woman is fixing her home
woman Is doing cart wheel (while) wearing ik ini the sand next to the beach IG woman fixing her home
woman is doing cart: wheel while wearing ik Ini in the sand next to the beach woman is fixing her home
AttnRoll
A woman is doing cart wheel while wearing ik ini in the sand next to the beach woman is tixing her home
MambaAttr
Figure 9: Explanations generated by different explanation methods for sentence of the SNLI validation set. This sentence belongs to the 'contradiction' class.
MambaLRP
am feeling very
blessed
today
they share close
bond
am feeling very blessed today they GI
SmoothGrad am feeling very blessed today that they share close bond
IG
feeling blessed today that share such close bond
AttnRoll
am feellng blessed today that they share such close bond
MambaAttr
am feeling blessed today that they share such close bond
Figure 10: Explanations generated by different explanation methods for sentence of the Emotion
validation set. This sentence belongs to the 'joy" class.
C.6.2 Computer vision
In this section, show explanations generated by MambaLRP alongside other baseline methods to interpret the predictions of the Vim-S model on several images of the ImageNet dataset. As can be seen, explanations generated by purely mradient-based explanation methods are very noisy. In contrast, attention-basee attribution methods have offered more focused and less noisy heatmaps. However, the last two images labeled 'paint brush' and 'flag pole', they could not faithfully explain the model's predictions. Among these approaches, Mambal RP stands out with its ability to generate sparse explanations, offering more faithful explanations how different image patches contribute to the final predictions.
C.7 Additional case results
For the
experiment in Section 6, we use synthetic dataset 7, In this dataset, single passkey (the 'needle") is inserted at different locations within collection repeated noise
20

---

She graduated with honors 2013 Having more than years of diverse experiences especially in N UR SE PR ACT (TON ER Eric a J H oyer affiliates with Swedish American
MambaLRP Hospital and cooper ates with other doctors and specialists in medical group Swedish americ
an Hospital Call Eric a H oyer on phone number 779 6 96 50 more Information and ises or book an appointment
She graduated
Having more than years diverse experiences PR (ATT ITON ER Eric J oyer affiliates with Swedish American
especially
Gl
Hospital and cooper ates with other doctors and specialists in medical group Swedish americ
phone number 779 96 86 50 for more information and adv ises or to book an appointment
Hospital Call Eric H
She graduated with hanors in 2013 Having more than years of diverse experiences especially in N UR SE PR ACT ITION ER Eric a J H oyer affiliates with Swedish American SmoothGrad Hospital and cooper ates with other doctors and specialists in medical group Swedish americ
an Hospital Call Eric a H oyer on phone number 779 6 96 50 for more information and adv ises or book an appointment
She gradusted honors 2013 Having more than years diverse experiences
especially
ITON Eric oyer affiliates Swedish American
IG
Hospital and cooper ates other doctors and specialists medical group Swedish americ
Hospital Call Eric H oyer phane number 779 96 B6 50 for more information and adv ises to book appointment
She graduated with honors in 2013 Having more than years of diverse experiences
especially In N UR SE PR ACT ITINN ER Eric 8 j
H oyer affillates with Swedish American
AttnRoll Hospital and cooper aates with other doctors and specialists in medical group Swedish americ
an Hospital Call Eric a H oyer on phone nnmbber 779 96 86 50 for more information and adv ises or to book an appointment
She graduated with honors 2013 Having more than years diverse experiences especially
moN Eric a oyer MambaAttr Hospital and cooper ates other doctors and specialists medical group Swedish americ Hospital Call Eric H oyer phone number 779 96 50 Information and adv ises book appointment
Swedish
Figure 11: Explanations generated by different explanation methods for sentence of the Medical BIOS validation set. This sentence belongs to the 'nurse' class.
sentences (the 'haystack'), as described in [37]. The dataset is composed of sequences with different context lengths. In our experiment, we use sequences with context lengths of 512, 1024, and 2048. We restrict the maximum context length to 2048 tokens align with the model's training configuration, as this experiment is not designed to evaluate extrapolation beyond this limit. The goal is to focus on certain limitations of the retrieval accuracy metric and and the solutions provided by Mambal We use Mamba-2.8D 2.B model 8 which is finetuned on the No Robots dataset [55] using context length of 2048. Then, we prompt the model extract the passkey hidden among irrelevant text by completing the phrase "The passkey is
Retrieval accuracy metric, which commonly used in the needle-in-at gredle-in-a-haysteee a-aasssyttackk experiment to analyze the model's performance. The synthetic dataset used for this experiment can be designed to include misleading information, which may cause the model generate the correct passkey based on incorrect evidence. In such cases, simply evaluating the retrieval accuracy may be insufficient. This issue can also arise when dealing with more realistic haystacks. Therefore, we introduced explanation-basee retrieval accuracy (XRA) in Section 6. Mambal and the XRA metric designed upon it can help to better examine the evidence the model relies on to retrieve the needle. this approach, we first identify the positions of the top-K relevant tokens by MambaLRP, and then, calculate accuracy by comparing those positions to the needle's position. We set the value of K to 2. This is because MambaLRP identifies the token immediately preceding the generated token the
"atrps://hugginggace harp://ca...c/llllaaaaaa/mmaaaaa 2. chat-no_robote
21

---

Image
GI
SmoothGrad
IG
AttnRoll
MambaAttr
MambaL RP
quail
nail
tench
hartebeest
paint
flag pole
-1.0
-0.5
0.0
0.5
1.0
Ours
Figure 12: Explanations produced by different explanation methods for images of the ImageNet dataset. Explanations produced by AttnRoll and MambaAttr are limited to non-negand non-negative values, whereas those generated by gradient-based techniques and MambaLRP includes both positive and negative contributions.
important one in most the examples and the evidence used for the passkey retrieval usually the second most important token.
The sample in Fig. 13 represents scenario where XRA approach proves valuable. In this the next token generated by the model the second part of the correct passkey (300). However, the model has incorrectly focused the number 300 in the phrase "Pass the key room 6300" to generate this token. Simply looking at the retrieved token might suggest that the model successfully retrieved the correct information. However, examining the Mambal _RP's explanation heatmaps provides deeper insights into the model's behavior. This helps us to debug the model more effectively and design better tests to analyze its capabilities.
C.8 Long-range dependeny comparison Transformers
To explore the capabilities of different model architectures in handling long-range dependencies, performed a direct comparison between Mamba and state-of-the Transformers (Llama-2" and Llama-310) focusing on their performance with inputs exceeding typical context lengths.
ct/.:/a.ccoa//eeeaaa 3
22

---

Next generated token: [300]
context: 2048 tokens
[1] There is za important info hidden inside a lot of irrelevant text. , Find and memor ize them. I will quiz you about the important information there.
[3] Text Pass the key to room 600 The grass is green. The sky is blue. The sun is yollow Here we go There and back again. The mountain is high. The snow cold. The air is thin Up ward we climb From base to peak The focest dense . The trees are tall The path current strong The banks
stremm Back and furth The ocean is vast The waves are cal ming The sand is warm. On ward we sail From coast to coast Here we go There and back again The grass green The sky blue. The suu yellow Hare
winding Ielo woods
and back. The ziver
undkly Down
is deuse There and back spain, The mountain is high. The snow cold The is thin The path winding. 300 is the pass key. Into the woods. Out and back. The river is wide The curzent is strong The banks are maddy Down the stream. Back The
Upp ward we slimb From
peak. The
tree
pass key is
Remember
forth. The occan
cal ming The sand
On ward From coest. Hore go. There and back
Haystack
The sir is thin. Up ward we climb From base peak. The forest is dense. The trees are tall. The path winding. Into the woods Out and Hack and forth, The
back The river ming. The sand is warm On ward we sail. . F'rm coust to coust. Here we go There and back aguin. The grass is green The sky is blue. The
The
strong The banks madity. Down the
The
is yellow. Here we go. There and back again. The mountain is high. The snow cold.
Qn ir: f61 The pass key is
the pass key
Figure 13: Detecting Clever-Hans effect in the needle-in-a heedll-in-a-haysttckk test. Given the 2K context length in this example, visualizing the entire text could be confusing. Therefore, we have removed most of the haystack from the visualization. In this example, the model has generated the correct passkey but the generation is not based on truly relevant information in the text.
For Llama-2 and Llama-3, we extract attributions using LRP for Transformers [3]. As in the Mamba experiment, we generate 10 additional tokens from the HotPoiQA dataset input and explain the prediction for each generated token. The results are shown in Table 9. For Llama-2, which was trained with a context length of 4096, the generated text becomes increasingly less sensible and repetitive for contexts longer than 4k, limitation noted also in [19, 5J. When analyzing the histogram distribution over tokens considered relevant to predict the next token, it appears 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, as presented in Table 9, its output becomes nonsensical for context lengths exceeding 4K tokens, characterized by the use of rare vocabulary and repeated tokens. Thus, the identified relevant tokens are mostly non-semantic such as new line token '<Ox0A>* in I.lama-2 and beginning of sentence token *
" 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 relevant. Given Llama-3's s much larger size (8B) compared to Mamba (130M) and their different training settings, this analysis supports that Mamba indeed uses long-range information. We also find that this ability is not exclusive to SSMs and can in principle also be achieved by Transformer models.
To what extent these findings depend on the amount of training data and model complexity remains an open research question. Our investigation of long-range dependencies in recent sequence genera- tion models highlights the value of faithful attribution methods like Mambal RPP in examining the capabilities and mechanisms utilized by models during generation.
C.9 Runtime comparison
In this section, we report the time required for each explanation method to generate its respective explanation. These times, measured in seconds, are averaged over samples from the Medical BIOS dataset. All baseline methods are evaluated on single A100-40GB GPU with a batch size of 1. All methods are applied to the 10mba-130M model. The results without fast CUDA kernels are shown in Table 10, while the results with fast CUDA kernels are presented in Table 11. We can observe that the runtime of Mambal RP is comparable to GradientxInput Since algorithms like Integrated Gradients and SmoothGrad require multiple function evaluations, their runtimes are significantly higher than MambaLRP and Gradientx
23

---

id length
Llama-2
Llama-3
Mamba
Summary <0x0A>>
Summary Ó The genus Dict per ma is mon
Summary Ô The species member of the genus Ap oll
The genus D ict
y os per ma _6
41
4k
Summary <0x0A Ohio Ohio.
Summary C The fol- lowing is summary of Finn 's head coaching
Summary C The fol- lowing is summary of the history of the
Ohio Obio _Ohio
109
Summary 5x00A> The Jb Jb Jb Jb Jb 5b Jb Jb Jb
Summary The mall is anchored by Hud- son 's Bay Walmart
Summary C The mall is cated in the heart the city
101
10
Count
loz
Count
102
Count
102
10*
10'
10
10a
10a
10a
2000
Position difference
Position difference
Position difference
Table 9: Long-Range dependency experiment, comparing Transformers and Mamba for different
context lengths.
Table 10: Runtime comparison. The time needed for each baseline method to generate its explanations. The times, measured in sec- onds, are averaged over the samples from the Medical BIOS dataset. The model used this experiment is Mamba-130M without using fast CUDA kernels.
Table 11:
Runtime comparison. The time needed for each baseline method to generate its explanations. The times, measured in sec- onds, averaged over the samples from the Medical BIOS dataset. The model used in this experiment is Mamba-130M using fast CUDA kernels.
Methods
Runtime
Methods
Runtime
Gradient x Input
0.7556
Gradient x Input
0.0335
SmoothGrad
22.9772
SmoothGrad
0.9785
Integrated Gradients
22.8071
Integrated Gradients
0.9742
AttRoll
2.1558
AttRRoll
MambaAttr
2.6661
MambaAttr
MambaLRP
0.4345
Mam
0.0306
24

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All claims are addressed in Sections 4, 5, and of the main paper and further details regarding each section can be found the supplemental material.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the computational efficiency of method including runtime comparisons in Section We further critically assess the performance of our approach via number of ablation studies presented in Section 5 in the main paper and in the supplemental materials Section C.2.1 and Section C.5. Limitations regarding scope and possible future directions mentioned in Section 7.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate ""Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex- system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
25

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We state our assumptions and theoretical results Section of the main paper alongside proofs given in Sections A.1, A.2, and A.3 in the supplemental material.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide necessary information for reproducing our results in Section 4 the main paper and Section C.2 of the supplemental material. Details of our proposed method and how it can be implemented can be found in Section C.2 of the supplemental material. Please refer to Section 5 of the paper and Section C of the supplemental material for more details regarding the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
26

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The code for reproducing our results and implementation of our proposed method is publicly accessible and link to code repository is added to the main paper. The datasets used in this work are publicly available, and comprehensive instructions for implementing MambaLRI method can be found in Section C.2 of the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
olicc) for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: // //nips.coa wwtpcc//ppppppeeeeeeeceeeeeeeooe wwtpc///ggiseseeeee httpg//ddde//eeeuummiieeeeeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please check Section 5 in the paper for details about experimental setting and evaluation metrics. Further details can be found Section C. of the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
27

---

Justification:
Please check Section 5.3 of the paper. Further experimental results can be found Section C.5 in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please check Appendix C.9.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that NeurlPS Code of Ethics https http///neuippppeeeee wttpppp.....ppaeaeeeeeeee wwtppblic//eececccccccccccccctlt htt.:a/ieeehhciiuielleeeee..
didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification: We do anticipate any harmful risks from the methods and analyses presented in this work.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
28

---

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The positive societal impacts of our work regarding model explainability are mentioned in multiple parts of the paper, in particular Section 6. do not anticipate any negative societal impact from the methods and analyses presented in this work.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All datasets and models used the paper are properly cited. We have also provided links to libraries and code repositories used, in the supplemental material.
29

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
30

---

Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
31


Reviews:
[{"review_id": "J5dzqEOjBP", "reviewer": "Reviewer_wgXY", "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.", "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?", "limitations": "No. The authors should discuss the limitation of the usefulness of their methods.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The paper is well written and easy to follow.\n2. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and useful comments. We will address them in our response below.\n\n> Regarding the main novelty of applying LRP to the Mamba\n\nAs 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.\n\nWe 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.\n\n> Regarding MambaLRP explanations not being surprising, without new insights into the behavior of architectures.\n\nIn 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.\n\nAs 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**.\n\n> Regarding leveraging MambaLRP explanations for model improvement\n\nPlease 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "coQeSdb3ga", "reviewer": "Reviewer_4ExW", "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.", "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?", "limitations": "The paper briefly explained the limitation in Section 7.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is easy to follow. \n- The introduced LRP framework is faster and provides more faithful explanations than other explanation methods. \n- 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.\n\n- 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. \n\nOverall, the paper lacks a more interesting and useful analysis of the methods using LRP as explained above. This limits the contribution of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> Regarding the paper being limited to the explainability of Mamba\n\nThe 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]**.\n\n**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).\n\n> Regarding comparing to the behavior of Transformers\n\nWhile 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.\n    \nTo 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.\n\n> Comparing long-range dependencies of Mamba and Transformers\n\nThank 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. \n\nFor 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 ``
`` 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.  \n\n 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.\n\n> Regarding the reduced speed-up of MambaLRP compared to Gradient × Input, SmoothGrad, and Integrated Gradients.\n\nWe 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.\n\nPlease 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VcbTTaFhjI", "reviewer": "Reviewer_mtfa", "summary": "### **Post-rebuttal update**\nGiven the authors' additional experiments and changes-to-be-made to the manuscript, raising my original score from a 5 to a 6.\n\n### **Original review**\nThe 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.", "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)?", "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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\n**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).\n\n**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).\n\n**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\nAs 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.\n\n### Clarity\nFor 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.:\n>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.\nHow are the relevant tokens determined?  Upon first pass, I wondered whether they were somehow derived from the\nsupporting fact labels in the HotpotQA dataset, but it seems to be calculated as the difference\tbetween 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?\n\nMost 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.\n\nOther portions which are unclear:\n- \"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?\n- \"All models were trained for a maximum of 10 epochs, with an early stopping mechanism in place.\" <- What is the early stopping criterion?\n- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your detailed review and useful suggestions. We will address your points below.\n\n> Evaluation of Mamba LLMs also on larger 2.8B models\n   \nWith 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. \n\n  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.\n\nTo 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.\n\nMamba-2.8B results:\n| Methods  | $\\Delta A^F$  |\n|-------------------|----------------|\n| Random            | 0.007          |\n| GI                | -0.043         |\n| SmoothGrad        | -0.675         |\n| IG                | 0.322          |\n| AttnRoll          | 0.452          |\n| G$\\times$AttnRoll (Mamba-Attribution) | 0.341          |\n| LRP (LN-rule)     | 0.820          |\n| MambaLRP (Ours)   | **1.157**      |\n\n\n> How are the relevant tokens in Figure 6 determined and how is the histogram calculated? Why are there 127 of 200 HotpotQA samples considered?\n\nAs 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.\n\n> Clarification regarding comparison to [MambaAttr2024]\n\nWe 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. \n\n> What are fast CUDA kernels in this context?\n\nWhen 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.\n\n> What is the early stopping criterion?\n\nIn 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.\n\n> Adding references to Appendix C.1 and adding details about instruction-finetuned model\n\nThanks 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.\n\n> Regarding including foreseeable limitations of our work and potential memory limitations, clarification on comparison to [MambaAttr2024]\n    \nRegarding 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. \n    \n 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.\n    \n 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.\n\n[MambaAttr2024] A. Ali, et al. The hidden attention of mamba models. arXiv:2403.01590, 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PQqGFxIXiq", "reviewer": "Reviewer_dNt6", "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.", "questions": "1.\tMamba 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.\n2.\tMinor: 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.\n3.\tIn 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).\n4.\tMinor: Perhaps a relevant work that is missing is “Does Transformer Interpretability Transfer to RNNs?” by Paulo et al.\n5.\tThere 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.\n6.\tI 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?\n7.\tHalf 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$$\nwhere 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?", "limitations": "From my perspective, the authors address most of the limitations, except those pointed out in the weaknesses section and question 6 (failure cases).", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n\n2) **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). \n\n3) **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.\n\n4) 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:**\n\n- 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?\n\n- 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]).\n\n2) **Insufficient empirical analysis:**\n\n- 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).\n\n- 2.2)  \n>  “Residual lack of conservation is due to the presence of biases in linear and convolution layers, which are typically non-attributable” \n\n 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. \n\n3) **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.\n\n[Uni] A Unified Implicit Attention Formulation for Gated-Linear Recurrent Sequence. Models. Zimerman et al.\n\n[Sca] An Empirical Study of Mamba-based Language Models. Waleffe et al.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the valuable comments. We address them below and add more discussions as an official comment.\n\n1.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].\n\nAs 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].\n\n1.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.\n\n2.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.\n\n2.2: Please find the results in Fig.A of PDF; conservation is fully preserved for biases set to 0.\n\n3: 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.\n\n**Questions**\n\n1: 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.\n\n3: 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.\n\n5: $D$ is used in MambaLRP. As it does not violate conservation, no modifications are necessary. \n\n6: 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).\n\n7: 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.\n\nWe will include your suggestions in the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their comments and valuable feedback. We responded to the comments and made the following changes to our submission. In particular:

- We extended our faithfulness evaluation in Table 1 (main paper) to the larger Mamba-2.8B model trained on SST-2, confirming MambaLRP consistently outperforms all other approaches. To our knowledge, this is currently the most extensive evaluation of XAI methods for Mamba in the literature.

- Reviewer 4ExW and Reviewer dNt6 expressed interest in a direct comparison to Transformers to investigate long-range dependencies. We performed this experiment for Llama-2 and Llama-3 using LRP for Transformers. Our analysis reveals that Mamba effectively uses long-range information, matched by the more sophisticated and larger Transformer model (Llama-3), and notably not Llama-2. Due to space constraints, details for this experiment are given in response to Reviewer 4ExW. 

> Clarification regarding the novelty and contributions of MambaLRP

As new classes of deep neural networks are developed, such as selective state space sequence models, XAI needs to keep up with these model innovations by providing reliable explanations. We contribute to this goal via:

A novel, thorough analysis of the relevance propagation through Mamba components, guided by the conservation axiom (see Section 4 and Appendix A).
Proposal of new LRP rules to mitigate violation of conservation, leading to state-of-the-art explanations for Mamba (see Appendix B). 
Efficient implementation of these proposed rules (see Section 4).
Thorough evaluation of attribution methods for Mamba (see Section 5).
Insightful and practical use cases, addressing aspects of AI safety and transparency (see Section 6).

While we focused on Mamba, the proposed propagation rules are generally applicable to other models using these components, e.g. multiplicative gates in recent models HGRN [18], RWKV [19] and MEGA [20]. 
We summarized our contributions and methodological novelty more clearly in Section 1, 2 and 4 now.

We also appreciate the positive remarks, such as:

- The improvement of the faithfulness is significant compared to the baselines. (Reviewer wgXY)
- The introduced LRP framework is faster and provides more faithful explanations than other explanation methods. (Reviewer 4ExW)
- The experiments show interesting interpretations of Mamba specifically for gender bias, long-context capabilities, and needle-in-a-haystack test in Mamba. (Reviewer 4ExW)
- Explainability of Mamba models is a very important topic. The proposed adaptation to Mamba models is new and an excellent application of this methodology. (Reviewer mtfa)
- This work has significant opportunity for significant impact as the need to explain the decision making processes behind Mamba models grows. (Reviewer mtfa)
- LRP for Mamba is an important direction that facilitates the community in improving and understanding these models. (Reviewer dNt6)
- The method is applicable and allows the authors to provide insightful analyses about Mamba models. (Reviewer dNt6)
- The paper is well written and easy to follow. (Reviewer wgXY)

In summary, additional experiments prompted during rebuttal (1) confirmed the robustness of our evaluation, (2) provided detailed descriptions addressing soundness and novel contributions of MambaLRP, and (3) extended our use cases with additional comparisons, demonstrating the reliability and practical benefits of MambaLRP to the community.

We responded to each reviewer's initial comments below and hope our answers appropriately address all their questions and concerns. We are happy to answer any remaining questions during the discussion phase.

To avoid redundancies, we added references from the individual rebuttals below.

**Rebuttal References**

[1] Ali et al. The hidden attention of mamba models. arXiv, 2024.

[2] Zimerman et al. A Unified Implicit Attention Formulation for Gated-Linear Recurrent Sequence Models. arXiv, 2024.

[3] Ali et al. XAI for transformers: Better explanations through conservative propagation. ICML, 2022. 

[4] Blücher, et al. Decoupling Pixel Flipping and Occlusion Strategy for Consistent XAI Benchmarks. TMLR, 2024.

[5]  Achtibat et al. AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers. ICML, 2024.

[6] Samek et al. Evaluating the visualization of what a Deep Neural Network has learned. IEEE TNNLS, 2017.

[7] Chen et al. Clex: Continuous length extrapolation for large language models. ICLR, 2024.

[8] Huang et al. Training-free long-context scaling of large language models. ICML, 2024.

[9] Bach et al. On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation. PloS one, 2015.

[10] Schnake et al. Higher-Order Explanations of Graph Neural Networks via Relevant Walks. PAMI, 2022.

[11] Anders et al. Finding and removing clever hans. Inf Fusion, 2022.

[12] Przemyslaw et al. Marrying Fairness and Explainability in Supervised Learning. ACM Conf on Fairness, Accountability, and Transparency, 2022.

[13] Samek et al. Explaining deep neural networks and beyond: A review of methods and applications. Proc of IEEE, 2021.

[14] Arrieta et al.  Explainable Artificial Intelligence: Concepts, taxonomies, opportunities and challenges toward responsible AI. Inf fusion, 2020.

[15] Ali et al. Explainable Artificial Intelligence: What we know and what is left to attain Trustworthy Artificial Intelligence. Inf Fusion, 2023.

[16] Roscher et al. Explainable machine learning for scientific insights and discoveries. IEEE Access, 2020.

[17] Letzgus et al. Toward explainable artificial intelligence for regression models: A methodological perspective. IEEE Signal Proc, 2022.

[18] Qin et al. Hierarchically gated recurrent neural network for sequence modeling. Neurips, 2023.

[19] Peng et al. RWKV: Reinventing RNNs for the Transformer Era. EMNLP, 2023

[20] Ma et al. Mega: Moving average equipped gated attention. ICLR, 2023.

Meta-review: After a fruitful discussion, the reviews agree that the paper is well written and clear (especially when improved with the details provided in the rebuttal) and that the proposed approach (MambaLRP) is a significant contribution to understanding an increasingly popular class of models. The proposed approach provides significant faithfulness improvements, which are very welcome, as showcased by the extensive experiments.
I strongly encourage the authors to release the source code, as promised, with the final version of the paper, as it is essential to enable reproducibility.
Decision: Accept (poster)