Vision Mamba Mender
Authors: Jiacong Hu
Anda Cao
Zunlei Feng
Shengxuming Zhang
Yi Wang
Lingxiang Jia
Mingli Song
Venue: NeurIPS 2024 poster (2024)
Keywords: Mamba
State Space Models
Computer Vision
Primary area: machine_vision

Abstract: Mamba, a state-space model with selective mechanisms and hardware-aware architecture, has demonstrated outstanding performance in long sequence modeling tasks, particularly garnering widespread exploration and application in the field of computer vision. While existing works have mixed opinions of its application in visual tasks, the exploration of its internal workings and the optimization of its performance remain urgent and worthy research questions given its status as a novel model. Existing optimizations of the Mamba model, especially when applied in the visual domain, have primarily relied on predefined methods such as improving scanning mechanisms or integrating other architectures, often requiring strong priors and extensive trial and error. In contrast to these approaches, this paper proposes the Vision Mamba Mender, a systematic approach for understanding the workings of Mamba, identifying flaws within, and subsequently optimizing model performance. Specifically, we present methods for predictive correlation analysis of Mamba's hidden states from both internal and external perspectives, along with corresponding definitions of correlation scores, aimed at understanding the workings of Mamba in visual recognition tasks and identifying flaws therein. Additionally, tailored repair methods are proposed for identified external and internal state flaws to eliminate them and optimize model performance. Extensive experiments validate the efficacy of the proposed methods on prevalent Mamba architectures, significantly enhancing Mamba's performance.

Vision Mamba Mender
Jiacong Hul Anda Cao', Zunlei Feng2 Shengxuming Zhang', Yi Wang', Lingxiang Jia Mingli Song!.4 1,3,4
'College of Computer Science and Technology, Zhejiang University, School of Software Technology, Zhejiang University.
State Key Laboratory of Blockchain and Data Security, Zhejiang University, Hangzhhuu High-Teck Tech Zone (Binjiang) Institute of Blockchain and Data Security
{j {jiaconghu, www.onggguccacaaaaaaoaamaa zunlelllleee eengggifngg jzju.ce edu. en, (2sxm1998, Be,i:88y__,,iinn lingixiannn brooksong) @zju. edu. cn
Abstract
Mamba, state-space model with selective mechanisms and hardware-aware ar- chitecture, has demonstrated outstanding performance long sequence modeling tasks, particularly garnering widespread exploration and application in the field computer vision. While existing works have mixed opinions of its application in visual tasks, the exploration of its internal workings and the optimization of its performance remain urgent and worthy research questions given its status as novel model. Existing optimizations of the Mamba model, especially when applied in the visual domain, have primarily relied on predefined methods such as improving scanning mechanisms integrating other architectures, often requiring strong priors and extensive trial and error. In contrast to these approaches, this paper proposes the Vision Mamba Mender, systematic approach for understanding the workings of Mamba, identifying flaws within, and subsequently optimizing model performance. Specifically, we present methods for predictive correlation analysis of Mamba's hidden states from both internal and external perspectives, along with corresponding definitions of correlation scores, aimed at understanding the workings of Mamba in visual recognition tasks and identifying flaws therein. Additionally, tailored repair methods are proposed for identified external and in- ternal state flaws to eliminate them and optimize model performance. Extensive experiments validate the efficacy of the proposed methods on prevalent Mamba chitectures, significantly enhancing Mamba's 's performance. For more information, please visit https::///// //vpsi:////iii/ mamba- mender. gi thub. io/.
1 Introduction
Deep learning has demonstrated outstanding performance in various fields of artificial intelligence, with deep neural networks such as convolutional neural networks (CNNs) [1, 2] and Vision Trans- former [3, 4] dominating the field of computer vision. However, the limited receptive field [5, 6] of CNNs and the quadratic computational complexity [7, 8] of Transformers constrain their further development. To overcome these limitations, an increasing number of studies have attempted to propose more advanced models [9-12]. Recently, Mamba [13], based on the state space model [11], has become the focus of these research efforts. Mamba can maintain nearly linear computational com- plexity while achieving global receptive field, leading outstanding performance. Consequently, has been widely adopted in the field of computer vision [14-16].
Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

To enhance the performance of Mamba in visual tasks, existing works mainly focus on improving the architecture of Mamba [17-21]. For instance, Zhu al.[17] proposed Vision Mamba, which adds branches to the original Mamba to simultaneously process image sequences in both forward and backward directions. Nearly simultaneously, Liu et al.[18] introduced VMamba, which adopts a four-way scanning strategy on top of the original Mamba to achieve a more comprehensive global receptive field. Other improvements include PlainMamba [19], Mamba-ND [20], SiMBA [21], and MambaMixer [22], all of which are variants of the original Mamba framework.
However, the aforementioned methods that optimize the Mamba model by improving its architecture are predefined and require strong prior knowledge and extensive trial and error. Additionally, Yu et al. [23] recently pointed out in their latest research that the current improvements made to Mamba- based visual models are unnecessary for visual tasks, especially for visual recognition tasks, This opposing view underscores the necessity and urgency of further optimizing Mamba models in the field of computer vision. Hence, unlike the aforementioned pre-optimization methods, this paper attempts to analyze the working mechanism of Mamba from a post- pessppeespec,,ve identify the reasons of flaws leading to incorrect prediction results, and automatically rectify them to further enhance the performance of Mamba models. Moreover, this approach is applicable to all Mamba-like models.
Based on this idea, in this paper, we propose Vision Mamba Mender, a systematic approach to understanding the working mechanism of Mamba from a post-perspective. identifying flaws within it, and rectifying them to ultimately improve model performance. In understanding the operational mechanism of the Mamba model, we categorize the computational process of Mamba into external state interaction and internal state interaction.
Along these two perspectives of external and internal states, introduce state correlation analysis method tailored for Mamba to establish the correlation between hidden states and predicted results. Additionally, we define external state correlation scores and internal state correlation scores to quantitatively analyze differences in state correlations, revealing flaws existing respectively in the external and internal states. Specifically, external state flaws refer to instances where correct model predictions in certain states are predominantly associated with foreground regions, while incorrect predictions are primarily linked to background regions. Internal state flaws, on the other hand, pertain to cases where correct predictions within a class are correlated with the same regions within the state and exhibit low overall complexity, whereas incorrect predictions within the class focus on different internal regions of the state and demonstrate higher overall complexity.
Furthermore, we propose repair methods tailored for addressing both external state flaws and internal state flaws. Specifically, in the repair of external state flaws, we impose constraints on the external state correlations within certain modules of Mamba to increase their correlation with foreground information during prediction. On the other hand, in the repair of internal state flaws, we impose constraints on the internal state correlations within certain modules of Mamba to enhance their correlation with genuine class-specific internal information during prediction. Through extensive experimentation, we demonstrate that the proposed Vision Mamba Mender is applicable to state-of- the-art Vision Mamba architectures.
The contributions of this paper are summarized follows:
We propose a novel post-hoc optimization method named Vision Mamba Mender. This method is applicable to existing state-of-the-art Vision Mamba architectures, identifying and repairing flaws in the Mamba model's mechanisms for visual recognition tasks from a post-hoc perspective, ultimately enhancing the model's performance.
We introduce a state correlation analysis method and correlation score definitions for Mamba from both external and internal hidden state perspectives. These methods are used to identify flaws. Additionally, we introduce a state correlation constraint method to rectify these flaws.
Extensive experiments demonstrate that the proposed Vision Mamba Mender can effectively identify and repair flaws in the Mamba model without introducing additional parameters, significantly improving the performance of the Vision Mamba.
2 Preliminaries
Mamba [13] is a novel and efficient sequence modeling model composed of multiple uniformly stacked Mamba blocks. Each Mamba block is constructed based on selective state space model

---

(SSM). Unlike the previous time-ine invariant SSM [11], it allows the parameters of the SSM to depend on the input, enhancing the model's 's expressive capacity. Furthermore, inspired by H3 [24] and Gated MI. [25], each Mamba block also incorporates modules such as causal convolution and gating mechanisms.
As illustrated in Figure I, given the hidden state hn ( of the i-th token in the input of the é-th Mamba block, the computation of the hidden state lin of the i-th token in the output of the f-th Mamba block is follows:
W
(6)
(0)1 - ..
n
X
n = cacusal-Cony1DDe
(
(0)
(2)
(C) selective-SSM(c (0) (8)
(0)
(3)
(2) SiLUUM, (!)).
(4)
SSM
/ (0! o ( w(()
(5)
where S.)U(.), causal-Cony ID).,, and ScLcctive-SSM() denote the activation function, the casual ID convolution, and the selective state model, respectively. W((), W and W(( are the projection matrices for the linear operations in the Mamba block, while xr z, and y represent the intermediate states within the Mamba block. Finally, by applying residual connection, the hidden state h! is obtained as follows:
G
Conv
x
2
W
W
M - (1 +yn (e)
(6)
In Eqn. (5), @ denotes the Hadamard product. The term 2n Eqn. (4) computed through separate pathway, serving gating mechanism to regulate the information flow in the main pathway.
The computational process of Mamba block.
Figure 1:
However, the original Mamba block is designed for one-dimensional sequences and is not suitable for handling multidimensional visual data, particularly for vision tasks requiring spatial awareness. Existing vision Mamba architectures enhance the basic Mamba block to accommodate these require- ments, such as ViM [17]. VMmmba [18]. PMamba [19], Mamba-ND [20], and SiMBA [21]. Unlike the predefined optimizations, approach analyzes the working mechanism of the Mamba model post-hoc, identifying flaws and making repairs to further enhance the model's performance.
Where Do Flaws Occur?
Identifying flaws in Mamba first requires an understanding of how Mamba operates. Some studies have provided empirical evidence to elucidate the mechanisms Mamba models in the NL.P domain, such their contextual learning ability [26], factual recall capability [27], and interpretability [28]. However, elucidating the operational mechanisms of Mamba models in the visual domain remains significant challenge. Ali et al. [29] established connection between the selective SSM within the Mamba block and the self-attention mechanism Transformers, allowing the selective SSM to represent the interaction process between any two states by constructing self-attention matrix, which is utilized for image feature attribution. However, focusing solely on the selective SSM within the Mamba block is far from sufficient, causal convolution, shown in Eqn. (2), also participates in the interaction between states. Moreover, other computational modules within the Mamba block used for state interaction must also be considered.
In this section, we first investigate the working mechanisms of Mamba. We summarize the computa- tional processes within Mamba as state interactions! These interactions are categorized into external state interactions (where a state interacts with other states to form a new state, as shown in Eqn. (2) and (3) and internal state interactions (where state interacts only with its internal information to form new state, as shown Eqn. (1), (4), and (5). We explore the operational mechanisms of Mamba from both the external and internal state interaction perspectives to identify flaws.
In Mamba models used for visual recognition, each state corresponds to an image patch or a learnable classification token.

---

8
9
10
12
13
14
15
16
17
18
19
20
21
22
23
24
Figure 2: Visualization of the external state correlation e(f.*) of the output ( from the selective- SSM module in different ViM [17] blocks. The depth of the ViM blocks increases from left to right.
3.1 External State Correlation Analysis
To understand the impact of external states predictions within Mamba's working mechanisms, we propose the Grad-ESC method, inspired by Grad-CAM [30], which uses gradient and activation information to assess the importance of each neuron in Secision-making Grad-ESC calculates the correlation between external states and the model's prediction outcomes, termed external state correlation. Unlike the attention mechanism derived by Ali et al. [29], which applies only to the selective SSM module within the Mamba block, Grad-ESC allows for correlation analysis of the outputs from any module within the Mamba block.
Specifically, given the predictive distribution p output by the Mamba model and the true class &s of the input sample, the calculation process for external state correlation e(f,*) RR (where H and
W represent the height and width of the input image) of the states {sh 1 In=1 output by the seleetive
SSM module (where (2) RRR, N and D the number and dimensionality of the states) is as follows:
1
N
Op
e (e.+)) = R(3 52
5 = (D(g o (s
= N
Osn (e)
(7)
n=1
where ((E..) E RD represents the weights used to weight the state
( along its dimensions based on
gradient information, En(-) denotes the average across all dimensions, 8n E IR signifies the degree of correlation between the nith state and the prediction after integrating gradient information, and RR denotes the operation of reshaping and scaling to the original image size after retaining only the states corresponding to the image patches.
As illustrated in Figure 2, taking the state 8/ as an example, we visualize the external state correla- tions of s in different blocks of ViM [17]. It can be observed that the external state correlations of (2) 8n vary across different blocks; some blocks exhibit correlations with foreground regions (such as the 1st layer), while others correlate with background regions (such as the 2nd layer). We consider correlations with the foreground to be interpretable, whereas correlations with the background are deemed uninterpretable. To quantify the interpretability of external state correlations, we propose the following definition of an external state correlation score based on commonly used interpretable evaluation methods such as perturbation test [29, 31] and segmentation test [32, 30, 33]:
Definition 1 (External Correlation Score). Given pre-trained Mamba model F()), input image i, foreground annotations of the imput image, and external state correlation ((.)) computed
through the proposed method, the external correlation score is defined as follows:
softmax @11)) softmax (F о i))
(CS(e(
IoU(elf.o)+,, m).
(8)
segmentation test
perturbation test
where e(.,))+ = 1(e((.0) 2 a) denotes the retention of important regions (those greater than the threshold a) in (ff.s), setting them to 1, and the removal of unimportant regions (those less than the threshold a). setting them to O, for use in negative perturbation tests. Conversely, (<fs)) = ((e(f.)) < ar) denotes the removal of important regions and the retention of unimportant regions, for use in positive perturbation tests. Additionally, (CS(e().))) E [0, +00), where higher value indicates higher correlation score.

---

samples
10
143
238
285
333
380
dimensions
Figure 3: Visualization of the internal state correlations ((f,a) of the output states (2) from the linear mapping module w! ViM [17] for samples of the same class. The horizontal axis represents the state dimensions, and the vertical axis represents the samples.
in
3.2 Internal State Correlation Analysis
To comprehend the influence of internal states on prediction outcomes within the operational mech- anism of Mamba, we introduce the Grad-ISC method. This method is utilized for computing the degree of correlation between the internal states and model predictions, termed the state-internal correlation.
Specifically, given the predicted distribution p outputted by the Mamba model and the true class k of the input sample, let's take the example of the n-th state xn E RD, which the output of the linear mapping matrix W The computation process for the corresponding internal state corelation (f.2z) E RR is follows:
Öp
j((.z) (t.x)
(t.x) In
(8)
Orn
(f.z) E RD denotes the weights applied to the dimensions of state In using gradient where On information. Similarly, the Grad-ISC method can perform internal state correlation analysis on the output of any module within the Mamba block.
(6)
Continuing with the example of state (f) we visualize the internal state correlations i((..)) for samples belonging the same class in ViM [17], as shown in Figure 3. It can be observed that for the same class, the regions of internal state correlation are relatively consistent. We posit that the more consistent and simpler the internal state correlation regions are for samples the same class, the more interpretable they are. Conversely, the more inconsistent and complex the internal state correlation regions are for samples of the same class, the more difficult they are to interpret. To quantify the interpretability of internal state correlations, we propose novel definition for the Internal Correlation Score:
Definition 2 (Internal Correlation Score). Given J samples belonging to the same class and the internal state correlation if :(f,a) E RD computed using the proposed method for particular sample, the Internal Correlation Score is defined as follows:
1 E (L,)+ (f,2)+
ICS(1 = =F
Σ j=1
(ba)) x
(f,))+
n.
+
n.2
+ + ((.a))))
n.J
(10)
simplicity
homogeneity
where {((.2) represents the internal state correlation for the j-th sample, ((x)+
1(( 3)
denotes the binarized internal correlation (set to if greater than 3, otherwise set to 0), 2 denotes the XOR operation. and (-) denotes the average over all dimensions. Furthermore, ICS (C.()(f.))
[0, +oo), where higher value indicates higher correlation score.
3.3 Identifying Flaw through Correlation Analysis
To analyze flaws within the Mamba model, we examined the external and internal state correlation scores under different conditions using the ViM model and the ImageNet- 10 dataset. This analysis allowed us to observe variations in correlation relationships across different states.
5

---

3 (a)
(b)
0.075
2
0.050
5
15
20
10
20
X
Yu
(c)
(d)
200
20
0
10 15
20
10
20
Figure 4: Comparison of external and internal state correlation scores across different blocks in the Mamba model between simple and difficult samples. (a) and (b) show the external state correlation scores for simple and difficult samples, respectively. (c) and (d) present the internal state correlation scores for simple and difficult samples, respectively.
Flaws in External State Correlation. To uncover flaws in the external correlations of states, we compare the scores of external state correlations in the Mamba model between simple and difficult samples, as illustrated in Fig. 4(a) and (b). For simple samples, it can be observed that across different blocks, the external correlation scores of states x!! ( 511) and 2 all better than those of state yn especially in deeper blocks. Furthermore, by comparing simple samples (a) with difficult samples (b), it can be noted that the external correlation scores of all states in all blocks have decreased. This indicates that for difficult samples, the Mamba model tends associate certain incomprehensible regions in the external states.
Flaws in Internal State Correlation. To unveil patterns in internal state correlations and detect flaws within them, we also conducted comparative analysis of the internal correlation scores within the Mamba model, as depicted in Figure 4(c) and (d). For both simple and difficult samples, the internal correlation score of state _.(2) is generally better than that of other states. Furthermore, simultaneous comparison of simple and difficult samples reveals decrease the internal correlation
scores of all states, including state (") This suggests that for difficult samples, the Mamba model tends correlate with incomprehensible regions within the internal states.
4 How to Repair Flaws?
In principle, repairing flaws within the model's internals enhance Mamba's decision-making process and improve its performance. However, there has been limited research specifically addressing such issues in Mamba models, particularly when applied the domain of visual processing. Therefore. in this section, we investigate post-hoc flaw repair in the Mamba model from two perspectives: external state correlation and internal state correlation.
4.1 External State Correlation Repair
To repair flaws related to external state correlations, it is essential to identify the key components within the Mamba block that need fixing. flaw identification regarding external state correlations,
it is observed that the states Jn Cn (e) and (e) exhibit flaws when predicting difficult samples. For these states, the external correlation scores for simple samples are higher than those for difficult samples. This implies that the model primarily correlates with foreground regions in correct predic- tions and with background regions incorrect predictions. Furthermore, considering that external state interactions occur only within Conv and SSM modules (as indicated by Eqn.(2) and E3n. we empirically suggest that the Conv and SSM components are crucial for influencing the model's anomalous decisions. Theoretically, it is also feasible to apply external flaw repair to other states.

---

We focus repairing the external correlation flaws of the states on and (0) output by the Conv and SSM modules in the deeper blocks. Specifically, we first identify difficult samples from the training set and then constrain the external correlations of the hidden states ea and Sn using foreground annotations m
Losse (n
c m) + Frw (el(.e) o m).
(11)
where Erw denotes the averaging operation over the two-dimensional matrix. It is important to note that during backpropagation, each term in the computational graph of Loss, is differentiable, which involves second-order gradients as specified in Eqn. (7). The complete loss function, combining with the original task loss, is formulated as follows:
Loss LosSce + ALosse,
(12)
where Lossce denotes the cross-entropy loss for the image recognition task, and À serves to balance the magnitudes of the respective loss components.
4.2 Internal State Correlation Repair
Similarly, when identifying flaws related to internal state correlations, it is observed that states in different blocks exhibit flaws when predicting difficult samples. For instance, the internal correlation scores of state ( for simple samples are higher than those for difficult samples. This suggests that in predictions biased toward correctness within class, the model aligns with internally consistent regions of the state, characterized by lower overall complexity. Conversely, in predictions biased toward incorrectness within a class, the model tends focus on internally inconsistent regions of the state with higher overall complexity. In our experiments, we consider the linear mapping W within deeper blocks critical component influencing the model's predictions. However, theoretically, it is also feasible to apply internal flaw repair to other states.
We focus on repairing the internal correlation flaws of the states (f) output by the linear mapping w!() within deeper blocks. Specifically, we first select J simple samples for each class from the training set and create corresponding internal correlation templates in 2(f,z)+ = - (} Σ/= maj ((f,))+)
for each class. We then utilize these templates to constrain the internal correlations of xm (9)
Loss; = ED (1 i((.))).
(13)
Similarly, during backpropagation, each term in the computational graph of Loss; is also differentiable, which involves second-order gradients specified in Eqn. (9). Combined with the loss for the original task, the complete loss as follows:
Loss = Lossce + Loss],
(14)
where y is a balancing factor to adjust the scale between loss functions.
It is important to note that external and internal state corrections represent two distinct perspectives, allowing these two methods to be used orthogonally. To maximize flaw repair, we recommend sequentially repairing external flaws (Eqn. (11)) followed by internal flaws (Eqn. (13)).
4.3 Results of Flaw Repair
We evaluated the proposed flaw repair method on five state-of-the-art Mamba models for visual tasks, including ViM-T [17], VMamba-T [18], SiMBA-S [21]. Efficient Mamban-T (EMamba-T) [34], and LocalVim-T [35]. To ensure smooth operation of the Mamba models within limited computational resources, we modified certain model parameters, such as the number of blocks and the dimensionality of hidden states. Additionally, we conducted experiments on three scales of the ImageNet [36] dataset: Im0geNet-50, ImageNet-3000 and ImageNet- To obtain the image foreground annotations m as defined in Definition 1 from the ImageNet dataset, we utilized annotations from the ImageNet-S dataset [37]. These annotations correspond to those used during training, covering 50 classes in ImageNet-50, 300 classes in ImageNet-300, and 919 classes in ImageNet-IK, with an average of 10 foreground-laaeledd samples per class. Further details on experimental settings are provided in Appendix B.
7

---

(a)
(b)
40
0.075
0.050
20
20
1
(
X
2n
y
200
(d)
100
0
10
15
20
15
20
Figure 5: Comparison of external and internal state correlation scores across different blocks the Mamba model before and after flaw repair. (a) and (b) show the external and internal state correlation scores before flaw repair, respectively. (c) and (d) present the external and internal state correlation scores after flaw repair, respectively.
Table I: Comparison of model accuracy after external state flaw repair for different states. Taking 'I("'' an example, represents the external flaw repair of state =!() The results in the second row of the table correspond to external flaw repair, while the third row presents the outcomes of
internal flaw repair. The experiment was conducted within the last block of the ViM model, using the ImageNet-50 dataset.
Base
(E)
On
(2)
8n
(E)
2n
(E)
y/u
(E)
Zn
(C)
+Gn
An +8n
(E)
(0)
In
(E) (E)
+sh
76.44
78.40
78.44
78.28
77.96
78.28
78.56
78.04
78.04 78.44
78.48
76.44
78.20
79.16
78.88
78.32
78.32
77.72
78.76
78.42
Comparing State Correlation Scores Before and After Repair. To validate the effectiveness of the flaw repair methods, we compared the internal and external state correlation scores of difficult samples before and after flaw repair, as shown in Figure 5. It can be observed that the proposed flaw repair methods effectively enhance the correlation scores of the Mamba model, thereby improving the predicted correlation regions within both internal and external states. Specifically, for the ViM model, after external state flaw repair, the external correlation scores the states of are significantly higher than before repair. Similarly, after internal state flaw repair, the internal correlation scores of the states x!() are also notably higher than before repair.
Repairing External and Internal Flaws in Different States. As shown in Table (second row), we conducted experiments on external flaw repair for different states within the same block. The results indicate that simultaneously performing external flaw repair on states en and (!) leads to an improvement in model accuracy. This aligns with findings in the main text regarding flaw
detection, where states xn (e) (E)
8n and ( exhibit flaws when predicting challenging samples. Furthermore, since the Conv and SSM modules in each block facilitate external interactions, applying
external flaw repair to states en and sy is effective.
Similarly, as illustrated in Table (third row), we performed experiments on internal flaw repair for different states within the same block. The results demonstrate that individually addressing internal flaws state on also results in improved model accuracy. This supports our earlier findings regarding flaw detection, where state (e) shows flaws when predicting challenging samples, thus validating the effectiveness of internal flaw repair for state x
Repairing SOTA Vision Mamba Models. As shown in the Table 2, we validated the proposed flaw repair methods on various mainstream Vision Mamba models. It can be observed that regardless of whether external or internal state flaw repair is performed, the accuracy of the repaired models exceeds

---

Table 2:
Comparison of the accuracy of the SOTA Vision Mamba model after flaw repair. Base' denotes the original model, +Ext', +int', and +All' represent models after external flaw repair, internal flaw repair, and simultaneous external and internal flaw repair, respectively.
ViM-T
VMamba-T
SiMBA-S
EMamba-T
LocalVim-T
ImageNet-50
Base 76.44
79,96
81.32
81.44
75.08
+Ext
78.48+2.04
81.08+1.12
84.48+3.16
82.68+1.24
78.68+3.60
+Int
78.20+1.76
82.36+2.40
84.64+3.32
83.24+1.80
78.20+3.12
+All
79.68+3.24
82.28+2.32
86.52+5.20
83.32+1.88
80.56+5.48
ImageNet-300
Base
75.11
75.04
68.08
74.67
70.39
+Ext
77.87+2.76
76.01+0.97
69.73+1.65
75.03+0.36
73.09+2.70
+Int
77.76+2.65
76.31+1.27
69.93+1.85
75.55+0.88
72.7142.32
+All
79.53+4.42
76.75+1.71
70.58+2.50
75.66+0.99
74.84+4.45
ImageNet-IK
Base
71.64
67.54
51.06
67.35
57.70
+Ext
73.02+1.38
68.34+0.80
52.12+1.06
67.62+0.27
59.30+1.60
+Int 12.79+1.15
68.67+1.13
52.23+1.17
67.89+0.54
59.90+2.20
+All 73.30+1.66
68.68+1.114
52.24+1.188
67.84+0.49
60.83+3.13
that of the original models. Specifically, for external state flaw repair, ImageNet-50, the accuracy of VMamba increased by 1.12% after flaw repair, and the accuracy of SiMBA-S increased by 3. 166.. For internal state flaw repair, on ImageNet-1K: 1, the accuracy of ViM increased by 1.15% after flaw worth noting that external state flaw
repair, and the accuracy of LocalViM increased by 2.20%. repair and internal state flaw repair can be orthogonal. For example, ImageNet-50, simultancously performing external state flaw repair and internal state flaw repair on ViM increased the accuracy by 2.04% and 1.76%, respectively, compared to performing each repair method individually. This resulted in an overall improvement 3.24% compared to the original model.
Related Works
Model Optimization. Optimizing model performance through various prior designs or theoretical derivations has always been pursuit in the field of artificial intelligence. Existing methods for optimizing the Mamba model, especially when applied to visual tasks, mainly involve adjusting the network architecture [17-22]. For instance, improvements have been made through the bidirectional scanning mechanism [17], cross-scanning mechanism [18], continuous 2D scanning mechanism [19], and incorporating
channel modeling techniques [21]. However, these enhancements designed and require extensive trial and error. For a more detailed understanding of Mamba's applications in computer vision, recommend readers refer recent surveys by Zhang et [14] and Xu et al. [16]. So far, there has been no research focusing on post-hoc optimization to correct internal flaws within the Mamba model improve its performance. Optimizing models beyond Mamba, such as Convolutional Neural Networks [1, 38] and Transformers [3, 39), is also primarily achieved through architectural improvements [2, 40-42, 4], feature enhancement [43-46], or some post-hoo hoc debugging methods [47-50]. However, these approaches either cannot be directly applied to a brand-new Mamba model or do not involve post-hoc optimization. In summary, unlike the aforementioned studies, the proposed Vision Mamba Mender is the first framework dedicated to post-hoe analysis and optimization of the Mamba model. Its aim is to rectify flaws within Mamba models in the domain of vision, thereby further enhancing the performance of Mamba models.
pre-
We adjusted the experimental setup and retined methodological details to enhance the fairness and repro- ducibility of the comparisons. Detailed experimental settings can be found in Appendix B.

---

Model Explanation. Enhancing model transparency and trustworthiness through research on model explainability has been a major focus in the field of artificial intelligence. In the context of Mamba model explainability, most studies have primarily concentrated on the model's context learning capabilities [26], factual recall abilities [27], and comparisons of the explainability between Mamba and previous models [28]. These studies do not address the identification of flaws in the Mamba model's mechanisms and are focused on natural language processing (NLP), leaving a gap in their application to visual tasks. In the realm of computer vision, Ali et al.[29] established a connection between the selective SSM in the Mamba block and the self-attention mechanism in Transformers. They developed a feature attribution method for Mamba based Attention- Rollout[51] and ceasffomer-Attrbbutiinn
[31], which can be used for image feature attribution. However, this method is limited to the selective SSM module and does not consider apply to other modules within the Mamba model. In research on model explainability beyond the Mamba model, plethora of methods have emerged, including activation-based methods [52-56], gradient-based methods [30, 57, 33], LRP-based methods [58, 59, 31], and perturbation- -based methods [60-62]. However, most of these explainability methods are designed for specific network architectures and may not be directly applicable to Mamba. Inspired by Grad-CAM [30], this paper introduces Grad- ESC and Grad-ISC methods tailored for the Mamba's external and internal states, respectively. These methods effectively establish the correlation between the model's states and predictions, providing simple yet powerful analytical tools for identifying flaws in the Mamba model's mechanisms.
6 Discussion
As one of the most prominent models today, Mamba has found widespread application in the field of computer vision. While attitudes towards its utilization visual tasks may vary, does not impede the exploration of its novel internal operational mechanisms. On the contrary, it urges further research into optimizing the performance of Mamba models applied in visual tasks. Morcover, at this juncture, considering the potential changes in future architectures, devising set of architecture-agnostic post-optimization methods becomes crucial.
However, in this work, we only analyzed the overall operational mechanisms of each module within Mamba, without delving into the internal details of each module. For instance, we hypothesize that there might be flaws within the SSM module itself during prediction, and further optimization of its internal details may potentially enhance the model's performance. Additionally, similar post- optimization paradigms should be further explored in non-visual tasks within Mamba even applied in non-Mamba architectures.
Conclusion
Mamba incorporates intricate computational modules, making it non-trivial to delve into its opera- tional mechanisms. In this paper, we have proposed post-analysis and post-optimization approach to understand the workings of Mamba in visual recognition tasks and enhance its performance. Specifically, departing from existing pre-defined optimization methods, we have introduced predictive correlation analysis and correlation scoring definitions from both internal and external perspectives of Mamba's 's states. These methodologies aim to identify flaws within Mamba's 's operational mechanisms. Simultaneously, have presented corresponding approaches for repairing internal and external flaws to optimize model performance. Extensive and comprehensive experiments have demonstrated the effectiveness of Vision Mamba Mender on mainstream Mamba architectures.
10

---

Acknowledgments and Disclosure of Funding
This research was supported by the Joint Funds of the Zhejiang Provincial Natural Science Foundation of China under Grant No. LHZSD24F020001 and Ningbo Natural Science Foundation (2022J182)
References
[1] Alex Krizhevsky, llya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
[2] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[3] Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn. Xiaohua Zhai. Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 1929, 2020.
[4] Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision, pages 10012-10022, 2021.
[5] Wenjie Luo, Yujia Li, Raquel Urtasun, and Richard Zemel. Understanding the effective receptive field in deep convolutional neural networks. Advances in neural information processing systems, 29, 2016.
[6] Jifeng Dai, Haozhi Qi, Yuwen Xiong, Yi Li, Guodong Zhang, Han Hu, and Yichen Wei. Deformable convolutional networks. In Proceedings of the IEEE international conference on computer vision, pages 764-773, 2017.
[7] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1004: 2019.
[8] Sinong Wang, Belinda Z Li, Madian Khabsa. Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, arXiv preprint 0476660000..4688, 2020.
[9] Fisher Yu and Vladlen Koltun. Multi-scale context aggregation by dilated convolutions. arXiv preprint arXiv:1511.07122. 2015.
[10] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165. PMLR, 2020.
[11] Albert Gu, Karan Gael, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint 1.00%9211 2021.
[12] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Mattco Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint ar30v:2335. 1304. 2023.
[13] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.000 0075222.00022... 2023.
[14] Hanwei Zhang, Ying Zhu, Dan Wang, Lijun Zhang. Tianxiang Chen, and Ye. A survey on visual mamba. arXiv preprint arXiv:2404. 15956, 2024.
[15] Xiao Liu, Chenxu Zhang, and Lei Zhang. Vision mamba: A comprehensive survey and taxonomy. arXiv preprint 2005 04404, 2024.
[16] Rui Xu, Shu Yang, Yihui Wang, Bo Du, and Hao Chen. A survey vision mamba: Models, applications and challenges. arXiv preprint arXiv:2404. 18861, 2024.
[17] Lianghui Zhu, Beneheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidircctional state space model. arXiv preprint vrXiv:2401.09111 2024.
[18] Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. arXiv preprint arXiv:2401.10166 2024.
11

---

[19] Chenhongyi Yang. Zehui Chen, Miguel Espinosa, Linus Ericsson, Zhenyu Wang, Jiaming Liu, and Elliot J Crowley. Plainmamba: Improving non-hierarchical mamba in visual recognition. arXiv preprint urXiv:2403. 2024.
[20] Shufan Li, Harkanwar Singh, and Aditya Grover. Mamba-nd: Selective state space modeling for multi- dimensional data, arXiv preprint arXiv:2402.0589,, 2024.
[21] Badri N Patro and Vijay Agneeswaran. Simba: Simplified mamba- based architecture for vision and multivariate time series. arXiv preprint arXiv:2403. 2024.
[22] Ali Behrouz, Michele Santacatterina, and Ramin Zabih. Mambamixer: Efficient selective state space models with dual token and channel selection. arXiv preprint arXiv:2403.19888 2024.
[23] Weihao Yu and Xinchao Wang. Mambaout: Do we really need mamba for vision? arXiv preprint MmXiv:2405.07992 2024.
[24] Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. 14052, 2022.
[25] Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoe Le. Transformer quality in linear time. In International conference on machine learning, pages 9099-9117. PMLR, 2022.
[26] Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is mamba capable of in-context learning? arXiv preprint arXiv:2402.03170 2024.
[27] Amnb Sen Sharma, David Atkinson, and David Bau. Locating and editing factual associations in mamba. arXiv preprint 3Xiv:2404.036466, 2024.
[28] Gonçalo Paulo, Thomas Marshall, and Nora Belrose. Does transformer interpretability transfer to rnns? 2024.
arXiv preprint arXi::444400597) 2024.
[29] Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint arXiv:2403.0155 2024.
[30] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das. Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient based localization. In Proceedings of the IEFE international conference on computer vision, pages 618-626, 2017.
[31] Hila Chefer, Shir Gur, and Lior Wolf. Transformer interpretability beyond attention visualization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 782-791, 2021.
[32] Grégoire Montavon, Alexander Binder, Sebastian Lapuschkin, Wojeiech Samek, and Klaus-Robert Müller. Layer-wise relevance propagation: an overview. Explainable Al: interpreting, explaining and visualizing deep learning, pages 193-209. 2019.
[33] Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. International conference on machine learning, pages 3319 3328. PMLR, 2017.
[34] Xiaohuan Pei, Tao Huang, and Chang Xu. Efficientymamba: Atrous selective scan for light weight visual mamba. arXiv preprint arXiv:2403.09977 2024.
[35] Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, and Chang Xu. Localmamba: Visual state space model with windowed selective scan. arXiv preprint arXiv:2403.09338 2024.
[36] Jia Deng. Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248 -55.. Icce, 2009.
[37] Shanghua Gao, Zhong-Yu Yu Li, Ming-Hsuan Yang, Ming-Ming Cheng, Junwei Han, and Philip Torr. Large- scale unsupervised semantic segmentation. 2022.
[38] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recogni- tion. arXiv preprint 14Xiv: 2014.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
12

---

[40] Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and puttern recognition. pages 2818-2826, 2016.
[41] Shunyu Liu, Yihe Zhou, Jie Song, Tongya Zheng. Kaixuan Chen, Tongtian Zhu, Zunlei Feng, and Mingli Song. Contrastive identit-aware learning for multi-agent value decomposition. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 11595- 11603, 2023.
[42] Lin Chen, Zhijie Jia, Lechao Cheng. Yang Gao, Jie Lei, Yijun Bei, and Zunlei Feng. Vit-calibrator: Decision stream calibration for vision transformer. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 1 1147-1155, 2024.
[43] Hugo Touvron, Matthieu Cord, Matthjs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efticient image transformers & distillation through attention. In International conference on machine learning, pages 10347-10333 PMLR, 2021.
[44] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint (xXXx::503..22555 2015.
[45] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 9729-97799 973897788222 2020.
[46] Kaiming He, Xinlci Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 16000-16009, 2022.
[47] Michael Brooks, Saleema Amershi, Bongshin Lee, Steven M Drucker, Ashish Kapoor, and Patrice Simard. Featureinsight: Visual support for error-driven feature ideation text classification. In 2015 IEFE Conference on Visual Analytics Science and Technology (VAST), pages 105-112 IEEE, 2015.
[48] Gabriel Cadamuro, Ran Gilad-Bachradd and Xiaojin Zhu Debugging machine learning models. In ICML Workshop on Reliable Machine Learning in the Wild, volume 103, 2016
[49] Zunlei Feng, Jiacong Hu, Sai Wu, Xiaotian Yu, Jie Song, and Mingli Song. Model doctor: A simple gradient aggregation strategy for diagnosing and treating enn classifiers. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 616-624, 2022.
[50] Josua Krause, Adam Perer, and Kenney Ng. Interacting with predictions: Visual inspection of black-box machine learning models. In Proceedings of the 2016 CHI conference on human factors in computing systems, pages 5686-5697, 2016.
[51] Samira Abnar and Willem Zuidema. Quantifying attention flow in transformers. arXiv preprint arXiv22005009928, 2020.
[52] Bolci Zhou, Aditya Khosla, Agata Lapedriza, Aude Oliva, and Antonio Torralba. Learning deep features for discriminative localization. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2921 2929, 2016.
[53] Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy. Do vision transformers see like convolutional neural networks? Advances in neural information processing systems, 1411611112,,28, 2021.
[54] Shunyu Liu, Jie Song, Yihe Zhou, Na Yu, Kaixuan Chen, Zunlei Feng, and Mingli Song. Interaction pattern disentangling for multi-agent reinforcement learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
[55] Shunyu Liu, Wei Luo, Yanzhen Zhou Kaixuan Chen, Quan Zhang, Huating Xu, Qinglai Guo, and Mingli Song. Transmission interface power flow adjustment: A deep reinforcement learning approach based on multi-task attribution map. IEFE Transactions on Power Systems, 39(2):3324-3335, 2024.
[56] Amin Ghiasi, Hamid Kazemi, Eitan Borgnia, Steven Reich, Manli Shu, Micah Goldblum, Andrew Gordon Wilson, and Tom Goldstein. What do vision transformers leam? a visual exploration. arXiv preprint arXiv:2212.0677 2022.
[57] Misha Denil, Alban Demiraj, and Nando De Freitas. Extraction of salient sentences from labelled documents. arXiv preprint 112 2014.
13

---

[58] Sebastian Bach, Alexander Binder, Grégoire Montavon, Frederick Klauschen, Klaus-Robert Müller, and Wojciech Samek. On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation. PloS one, 10(7)=001 30140, 2015.
[59] Grégoire Montavon, Sebastian Lapuschkin, Alexander Binder. Wojciech Samek. and Klaus-Robert Müller. Explaining nonlinear classification decisions with deep taylor decomposition. Puttern recognition. 65:211 222, 2017.
[60] Ruth Fong and Andrea Vedaldi. Interpretable explanations of black boxes by meaningful perturbation. In Proceedings of the IEEE international conference on computer vision, pages 3429 3437, 2017.
[61] Ruth Fong, Mandela Patrick, and Andrea Vedaldi. Understanding deep networks extremal perturbations and smooth masks. In Proceedings of the IEEE/CVF international conference on computer vision, pages 2950-2958, 2019.
[62] Scott M Lundberg and Su-In Lee. A unified approach to interpreting model predictions. Advances in neural information processing systems, 30, 2017.
[63] Xiao Wang, Shiao Wang, Yuhe Ding, Yuehang Li, Wentao Wu, Yao Rong, Weizhe Kong, Ju Huang, Shihao Li, Haoxiang Yang, et al. State space model for new-generation network alternative to transformers: A survey, arXiv preprint arXiv:2404.095/1 095516 2024.
[64] Badri Narayana Patro and Vijay Srinivas Agneeswaran. Mamba-360: Survey of state space models as transformer alternative for long sequence modelling: Methods, applications, and challenges. arXiv preprint arXiv:2404. 121 2024.
[65] Πyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv1771..51001 2017.
arXiv preprint
14

---

Vision Mamba Mender. The magnifying glass held in one hand symbolizes flaw detection, while the wrench held in the other hand symbolizes flaw repair.
Appendix for Vision Mamba Mender
To facilitate a better understanding of the value and significance of this work, as well as to thoroughly demonstrate the effectiveness and applicability of the proposed method, we have provided the algorithm code in supplementary material. This code will be made publicly available. Additionally, in the appendix, we have supplemented more information about the Mamba architecture, related work on Mamba interpretability. detailed experimental settings, additional results on state flaw identification, additional results on state flaw repair, and detailed ablation experiments, follows.
A Origin of the Mamba Model
One of the core components of Mamba the State Space Model (SSM). SSMs were initially inspired by continuous systems, which map a one-dimensional input signal æ(t) E R to a onc-dimensional output signal y(t) E R through an N-dimensional latent state h(t) RN The general computational process can be defined as follows:
(((() = AM(() + Bx((), (t)) Ch(t),
(15)
where A E NNN, BERNII and C E R are the state matrix, input matrix, and output matrix, respectively.
Unlike the continuity in Equation 15, Structured State Space Models (S4) [11] employ transformation methods, such as the zero-order hold (ZOH), to discretize it.
h'(t) At(t) Br(t),
(16)
vt) = Ch(t),
where A exp (AA) and B (AA) (ex)p-((exp))) (AA) 1) AB are the discretized parameters, and denotes the step size.
Building this foundation, the Selective State Space Model (S6) creatively introduces selective scanning to overcome the limitations imposed by time-invariant parameters on context representation learning:
B Sx(x), C = (x), A (A (x())),
(17)
where Sp, Sc, and SA are different linear mappings. Additionally, for more detailed information on Mamba's specifics and applications, we recommend readers consult several recent and well-regarded review articles [63, 64, 14-16).
15

---

B Detailed Experimental Settings
To facilitate the reproducibility of our results, we provide detailed experimental settings as follows. Throughout the entire experiment, we utilized 8 NVIDIA A40 GPU cards and CPU with 24 cores and 500GB of memory.
Model Parameter Settings. To enable the Mamba model to train efficiently within limited compu- tational resources, we adjusted certain parameters across the Mamba models. For instance, in the case of VMamba-T, we set the patch size to 16x16. For SiMBA-S, the model depth was adjusted to [2. 3, we
2], and we introduced a class token for classification in the last block. For similarly set the patch size to 16x 16. In the case of ViM-T, we reduced the model depth from 20 to and the state dimensionality to 128. These parameter adjustments were made to validate the proposed methods under constrained resources. While this may somewhat reduce the baseline performance, it does not affect the fairness of comparing the Mamba model's ' performance before and after flaw identification and repair.
Model Training Settings. To ensure a fair comparison with limited resources, our training settings primarily followed the experimental setup of DeiT [43]. Specifically, we employed data augmentation techniques such as random cropping and random horizontal flipping. When training on 224x224 input images, we optimized the model using AdamW [65] with momentum of 0.9, a total batch size of 128, and a weight decay of 0.1, We utilized cosine learning rate schedule with an initial learning rate of 5e-4, training the Mamba model for 300 epochs. In particular during the baseline model training, these training strategies resulted in an exceedingly smooth training curve in the later stages, effectively optimizing the fit. During testing, we performed center cropping the validation set to extract 224x224 images.
State Flaw Identification. For external state correlation analysis. the threshold is set to 0.5 by default. For internal state correlation, the threshold is set to 0.3 by default. The impact of thresholds on flaw detection can be found in the Appendices C and D. It is worth noting that for multi-branch Mamba, when conducting external flaw identification and repair, we merge all branches corresponding to the same state together. However, for internal flaw identification and repair, we consider each branch of the state separately.
State Flaw Repair. In our experiments, the balance weight À for the loss function of external state flaw repair is set to le + 7 by default, and the balance weight y for the loss function of internal state flaw repair is also set to le + 7 by default. Furthermore, based on the conclusions about flaw identification in the main text, external state flaw repair is applied by default to the first Mamba block, while internal state flaw repair is applied to the last Mamba block. The impact of repairing flaws in different blocks on model performance can be found in Appendices E and
C Impact of Threshold a on External State Correlation Scores
As shown in Figure 6, we compared the impact of different values of threshold a external state correlation scores. By contrasting simple samples with difficult samples, we found that the value of does not significantly affect the relative magnitudes of the external correlation scores. Specifically, regardless of whether is set to 0.0, 0.2, or 0.8, the external correlation scores for states ( (()
Sn (t) and yn in difficult samples all decrease compared to those in simple samples. Additionally, when is set between 0.4 and 0.6, more details can be observed; in different blocks, the external
correlation scores of states xn en 88 and =(() are better than those of state y/( Therefore, we
(e) (6) (f)
recommend that the value of a should not be too extreme when calculating external state correlation
scores. can be set between 0.4 and 0.6; we use default setting of a = 0.5.
D Impact of Threshold 33 on Internal State Correlation Scores
We also examined the impact of different values of threshold 3 on internal state correlation scores, as shown in Fig. 7. Similarly, by comparing simple and difficult samples, can be observed that the value of 3 does not affect the relative magnitudes of the internal correlation scores. Specifically,
16

---

(
Easy samples
Diffeult samples
a
0.0
a
0.0
3.00
0.083
2.95
0.082
20
20
0.2
0.2
3.0
2.5
0.07
20
5
α 0,4
α
0.4
0.075
0,050
20
5
20
0.075
α 0.6
0.6
2
0.050
0.025
0
5
10
15
20
5
10
15
20
α
0.8
20
0.8
0.5
0.01
0.0
0
10 15 20
10
15
20
Figure 6: Comparison of external state correlation scores for simple and difficult samples under different threshold values of a. The left column of images shows the external state correlation scores for simple samples at various a values. The right column of images shows the external state correlation scores for difficult samples at various a values.
regardless of whether is set to 0.0, 0.2, or 0.8, the internal correlation scores for states x. (e) (() 2(4) and (() difficult samples all decrease compared to those in simple samples. Additionally, when 3 is set between 0.4 and 0.6, more details can be observed; in different blocks, the internal
correlation scores of state (f( are better than those of states (E) (E) (2) and yn
(2) Therefore, we recommend that the value of B should not be too extreme when calculating internal state correlation
Sn 2n
scores. can be set between 0.4 and 0.6; we use default setting of 0.5.
E External Flaw Repair in Different Blocks
As shown in Figure E(a), we conducted study on repairing external flaws in different blocks, revealing a clear enhancement model performance following external flaw repair. Specifically, when external flaws are repaired in the final block, the model's 's accuracy improves by 2.04% compared to the baseline. In contrast, while repairing external flaws in other blocks, such as the tenth block, also yields accuracy increase, the improvement is only 1.72%. This discrepancy may be attributed
to the presence external flaws in the states across nearly every block, as illustrated Figure including states such zn samples. However, due to structures like residual connections, the influence of external flaws on
(() (e) (E) (E) and 2n which exhibit significant external flaws on difficult
8n
model predictions varies across different blocks.
17

---

9
Easy samples
Diffeult samples
2.5
B
0.0
B
0.0
2.0
1.5
1.5
15
10
15
20
B 0.2
20
B 0,2
50
10
15
10
20
200
B = 0.4
40
B 0,4
100
20
0
5
10
15
20
0
10
15
20
400
B
0.6
β
0.6
75
200
25
15
10
20
400
100
B
0,8
B
0,8
200
15
10
20
Figure 7: Comparison of internal state correlation scores for simple and difficult samples under different threshold values of B. The left column of images shows the internal state correlation scores for simple samples at various 3 values. The right column of images shows the internal state correlation scores for difficult samples at various 3 values.
Mww
(b)
79
79
78
78
77
5
15
20
5
15
Figure 8: Comparison of model accuracy after state flaw repair in different blocks. (a) shows the
results of external state flaw repair, conducted on states e and (e) in each block of the ViM model, using the ImageNet-50 dataset. (b) shows the results of internal state flaw repair, conducted on state (t)
on in each block of the ViM model, also using the ImageNet-50 dataset.
18

---

Internal Flaw Repair in Different Blocks
As shown in Figure E(b), we conducted study on repairing internal flaws in different blocks, which also demonstrates noticeable improvement in model performance. Specifically, when internal flaws are repaired in the final block, the model's s accuracy increases by 1.76% compared the baseline. In contrast, repairing internal flaws in other blocks, such as the second-to-las block, results in an accuracy improvement of 2.4%. Furthermore, repairing internal correlations other blocks yields varying levels of accuracy enhancement. This variation may similarly stem from the model's internal structures, such as residual connections, which lead to different degrees of influence from internal flaws model performance across the various blocks.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Section 1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long 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: Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create 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 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 the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to 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 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, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in 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.
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: [NA]
20

---

Justification: [NA]
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 complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section B.
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 be 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 to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way 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.
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) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either 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).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, 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 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 to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [Yes]
Justification: The algorithm code is available in the supplementary material and will be publicly accessible.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https: /nips. ctipsscc/public//uu guwdee/dss///// CCodeSubeesiaceecceee sisiissiinn sssssnoliccyy 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, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Guidelines:
Section B.
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, 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: [No]
Justification: To prevent the occurrence of accidental experimental results, we conducted each experiment at least twice to ensure the accuracy of the findings. The data presented in this paper are the averages of two or more experimental runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at 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, or 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.)
22

---

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.
It is OK to report 1-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 of range (e.g. negative emrr rates).
If error bars are reported in tables or 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 reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the 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 well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS 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 (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: [NA]
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.
23

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses 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 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 best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.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: Please refer the citations in the paper.
Guidelines:
The answer NA means that the paper does not 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 (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/dataseta 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 it has changed) should be provided.

---

If 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: [Yes]
Justification:
We have provided the complete code.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was 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 details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or 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]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on 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.
We 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 as the institution conducting the review.
25


Reviews:
[{"review_id": "BxNLVX1qo5", "reviewer": "Reviewer_4vzy", "summary": "1. Addressing the main issue in the existing model:\n   Mamba, despite its success in long sequence tasks, faces mixed opinions and challenges in visual tasks due to inherent flaws and suboptimal performance. Understanding these flaws and optimizing Mamba's performance in the visual domain are critical research questions.\n\n2. Proposed solution for the issue:\n   The paper proposes Vision Mamba Mender, a systematic approach to enhance Mamba's performance in visual recognition tasks. This approach involves predictive correlation analysis of Mamba's hidden states, both internally and externally, to identify and address flaws. Tailored repair methods are then applied to optimize model performance.\n\n3. How did the author evaluate the proposed solution on various standard dataset:\n   The efficacy of Vision Mamba Mender is validated through extensive experiments on prevalent Mamba architectures. These experiments demonstrate significant improvements in model performance across various standard datasets, showcasing the practical impact and effectiveness of the proposed methods. The algorithm code is also provided for transparency and reproducibility.", "questions": "Please refer to the Weaknesses block", "limitations": "Please refer to the Weaknesses block", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Innovative Post-Hoc Optimization: Vision Mamba Mender introduces a novel approach to optimize Mamba models post-training, focusing on identifying and rectifying operational flaws rather than predefining architectures. This method is innovative in its systematic approach to improving model performance.\n\n2. Comprehensive State Analysis: The paper introduces a detailed state correlation analysis method that evaluates Mamba's hidden states from both external and internal perspectives. This comprehensive analysis helps in pinpointing specific flaws that affect prediction accuracy in visual recognition tasks.\n\n3. Tailored Repair Methods: Tailored repair methods are proposed for both external and internal state flaws identified through the analysis. By imposing constraints on state correlations within Mamba modules, these methods effectively enhance the model's ability to make accurate predictions.\n\n4. Applicability to Existing Architectures: Vision Mamba Mender is designed to be applicable across various state-of-the-art Vision Mamba architectures. This versatility ensures that the optimization approach can be seamlessly integrated into different implementations without significant modifications.\n\n5. Experimental Validation: Extensive experiments validate the effectiveness of Vision Mamba Mender across different benchmarks and datasets. The approach consistently demonstrates improvements in model accuracy, showcasing its practical utility in real-world applications.\n\n6. Transparency and Reproducibility: The availability of algorithm code in the supplementary material and commitment to making it publicly accessible enhance the transparency and reproducibility of the research findings. This openness facilitates further validation and adoption of the proposed methods by the research community.", "weaknesses": "1. Limitation in Methodology: The paper mentions the introduction of state correlation analysis and repair methods for Mamba, but it lacks clarity on the specific algorithms or mathematical formulations used. This could hinder reproducibility and transparency, as other researchers may find it challenging to implement or verify the proposed methods without detailed methodology descriptions.\n\n2. Assumptions on Flaws and Repair Methods: The paper suggests that Vision Mamba Mender identifies and rectifies anomalies in Mamba's mechanisms, focusing on external and internal state correlations. However, without empirical evidence or case studies illustrating the nature of these flaws across different datasets and scenarios, the validity and generality of these assumptions remain unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and the positive feedback. We have carefully reviewed each of your comments. Although some comments may not fully align with our research objectives, we are nonetheless very appreciative of your feedback. Below are our responses to each of your comments.\n\n---\n\n> **Q1**: The paper mentions the introduction of state correlation analysis and repair methods for Mamba, but it lacks clarity on the specific algorithms or mathematical formulations used. This could hinder reproducibility and transparency, as other researchers may find it challenging to implement or verify the proposed methods without detailed methodology descriptions.\n> \n\n**A1**: Thank you for your valuable comments. We understand your concerns. In response, we have detailed the methods proposed for the novel Mamba model in the visual domain to ensure clarity and reproducibility.\n\nIn Section 3 of the paper, we thoroughly describe the computational process of the Mamba model using numerous formulas. Subsequently, in Section 4, we provide a detailed mathematical formulation and definitions for the algorithms used to identify defects in the Mamba model, including how to establish external/internal state correlations and how to define correlation scores. Finally, in Section 5, we offer a comprehensive description, using mathematical formulas, of the algorithms employed to repair identified defects in the Mamba model, including specific methods for addressing external/internal state defects.\n\nAdditionally, we have included detailed experimental setups in Appendix D. We hope these measures will assist readers in implementing our proposed algorithms effectively and address your concerns.\n\n---\n\n> **Q2**: The paper suggests that Vision Mamba Mender identifies and rectifies anomalies in Mamba's mechanisms, focusing on external and internal state correlations. However, without empirical evidence or case studies illustrating the nature of these flaws across different datasets and scenarios, the validity and generality of these assumptions remain unclear.\n> \n\n**A2**: Thank you for your comments. We fully understand your concerns.\n\nThe approach proposed in our paper, which is based on external and internal state correlations, aims to explore the mechanisms of the Mamba model from these two perspectives. We have demonstrated the presence of external and internal state defects in the Mamba model through a comparison of state correlation scores for simple and challenging samples. Extensive experiments and analyses are detailed in Section 3.3 of the paper, as well as in Appendices E and F.\n\nFurthermore, in the field of computer vision, other researchers have also attempted to understand the workings of the Mamba model, similar to our approach. For example, Ali et al. [1] have explored the mechanisms of Mamba in image recognition tasks and their findings reveal some external state anomalies. However, their work primarily focuses on external state correlations and lacks an analysis of internal state correlations. Additionally, their methods are applicable only to state space models and cannot be used for other computations within the Mamba block (as discussed in lines 113 to 119 of our paper).\n\nIn summary, this paper is the first to propose a post-hoc optimization approach for the Mamba model by analyzing both external and internal state correlations to identify and repair defects. This not only improves model performance but also aids readers in understanding the mechanisms of the novel Mamba model. We appreciate your feedback and hope the above response addresses your concerns.\n\n[1] Ali, Ameen, Itamar Zimerman, and Lior Wolf. \"The hidden attention of mamba models.\" *arXiv preprint 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": "1pz4a9OWZD", "reviewer": "Reviewer_nvv2", "summary": "This paper analysis Mamba model from a post-perspective. It introduce a state correlation analysis method to establish the correlation between hidden states and predicted results, and analysis the external state flaws and internel state flaws.  Furthermore, this manuscript propose repair method to handle these flaws. Extensive experiments show its advantanges.", "questions": "see weakness", "limitations": "see weakness", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is well written and easy to follow.\n\n2. The topic is interesting and novel, which propose a new analysis perspective for the recent model Mamba.\n\n3. Since the Mamba model is very popular recently, analysis this model is well motivate.\n\n4. The experiment results are abundant and solid, as well as the detailed experiment setting.", "weaknesses": "1. One importent point is the classification task is not necessary only rely on the foreground pixels. Background regions may provide useful information some times. So the proposed external flaws may not be totally reasonable.\n\n2. Only classification on ImageNet results are provides. It it suggented to provide more diverse experiments to full evaluation the proposed repair method.\n\n3. Minors:\n\n(1) Not all symbols are explained from Eq.(1)~(6).\n\n(2) Missing punctuation (comma or full stop) in Eq.(1)~(6).\n\n(3) In Fig.5, there are only (a)(b)(c)(d) in subfigures, while authors use (1)(2)(3)(4) in the caption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and comments. We are pleased that you find our research topic both interesting and novel, and that you believe the analysis of the Mamba model is well-motivated. We are also gratified that the experimental results have met your approval. Below are our responses to each of your comments.\n\n---\n\n> **Q1**: One important point is the classification task is not necessary only rely on the foreground pixels. Background regions may provide useful information some times. So the proposed external flaws may not be totally reasonable.\n> \n\n**A1**: Thank you for your valuable feedback. We agree that background regions can sometimes provide useful information. However, it is important to clarify that external flaws are identified through the analysis of external state correlation scores (Section 3.1). These scores reflect the interpretability of external correlations, specifically the degree to which the predicted results are associated with the foreground object regions (lines 146 to 159 of the paper). By comparing external state correlation scores between simple and difficult samples, we observed a significant drop in the scores for difficult samples (lines 196 to 203). This indicates that external flaws do not imply a complete disregard for background information, but rather that the Mamba model tends to focus more on background information in difficult samples, potentially neglecting important foreground details.\n\nFurthermore, during the defect repair process, we only constrained a small number of samples to focus on the foreground (lines 254 to 256 of the paper). This means that most samples are free to learn useful information from the background if it is beneficial, and our results demonstrate that this approach indeed improves the performance of the Mamba model.\n\n---\n\n> **Q2**: Only classification on ImageNet results are provides. It it suggested to provide more diverse experiments to full evaluation the proposed repair method.\n> \n\n**A2**: Thank you for your valuable feedback. ImageNet is a widely used large-scale dataset for classification experiments. To thoroughly evaluate the proposed repair method within the limited time available for the rebuttal, we have also tested the ViM model on the smaller CIFAR-10 dataset. The results of these experiments are shown bellow:\n\nCIFAR-10 | Base: 84.25 | Flaw Repair: 85.77,\n\nwhere “Base” refers to the original model, and “Flaw Repair” refers to the proposed defect repair method.\n\nAs shown above, the proposed defect repair method is also applicable to smaller datasets. For instance, the accuracy of the ViM model on CIFAR-10 improved by 1.50% after applying defect repair. It is worth noting that due to the small image size of CIFAR-10, the repair method in this experiment only addresses internal state defects.\n\nAdditionally, to further validate the effectiveness of the proposed repair method, we conducted comparative experiments on the external and internal state correlations in the Mamba model before and after defect repair. These comparisons are illustrated in Figure 5 and described in lines 268 to 275 of the paper. Figure 5 shows that prior to defect repair, both external and internal state correlations of the Mamba model had low scores, indicating that the model associated incorrect regions during predictions. After defect repair, the scores for both external and internal state correlations improved, demonstrating that the model associated more accurate regions during predictions.\n\nIn summary, the proposed repair method has broad applicability. We hope this response addresses your concerns, and we will update the paper with additional experiments to further validate the proposed repair method.\n\n---\n\n> **Q3**: Minors: (1) Not all symbols are explained from Eq.(1)~(6).\n> \n\n**A3**: Thank you for your careful review and observation. We indeed overlooked providing explanations for some symbols. We have now added the following explanations after Eq. (1)~(6):\n\n“Here, SiLU$(.)$, causal-Conv1D$(.)$, and selective-SSM$(.)$ denote the activation function, the casual 1D convolution, and the selective state model, respectively. $W_x^{(\\ell)}$, $W_z^{(\\ell)}$, and $W_y^{(\\ell)}$ are the projection matrices for the linear operations in the Mamba block, while $x$, $z$, $c$, $s$, and $y$ represent the intermediate states within the Mamba block.”\n\nThank you once again for pointing out this detail.\n\n---\n\n> **Q4**: *Minors: (2) Missing punctuation (comma or full stop) in Eq.(1)~(6).*\n> \n\n**A4**: Thank you for your meticulous review and observation. We have added the missing commas and full stops in *Eq.(1)~(6)* and have rechecked the entire manuscript.\n\n---\n\n> **Q5**: Minors: (3) In Fig.5, there are only (a)(b)(c)(d) in subfigures, while authors use (1)(2)(3)(4) in the caption.\n> \n\n**A5**: We apologize for the typographical errors. We have corrected these discrepancies in the manuscript and have carefully reviewed the entire paper to ensure consistency.", "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": "QyIzZxohYv", "reviewer": "Reviewer_3tmX", "summary": "The papers addresses the limitations of Mamba based models in vision through a post-hoc optimization scheme that address external state flaws and internal state flaws, identified through respective internal and external state correlation analysis. The introduced corrective measures improve performance on image classification tasks across various mamba-based backbones on the imagenet dataset.", "questions": "Please consider the Weaknesses, particularly point 2. It is important to have results for detection and segmentation to ensure the backbone is learning features that are rich enough for dense prediction tasks and not just over-fitting to the image classification task on imagenet.", "limitations": "There is some discussion on limitations in section 6. However, the authors should include additional limitations regarding their method, particularly about the fact that it requires additional annotations that may not be available across all datasets.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is very well-motivated and addresses a specific topic very relevant to the vision community.\n2. The presented correlation analysis is novel and detailed, and the resulting flaws identified along with presented mitigations are consistent and show consistent improvements across various backbones.\n3. The paper is well-written and easy to follow.", "weaknesses": "There are two major concerns here:\n\n1. The method is limited by the required annotations (for example foreground annotation) which limits its usefulness across other classification datasets and other tasks.\n2. The paper does not evaluate on any detection/segmentation benchmarks. This is required to understand whether the features learned by the backbone are rich enough to generalize to these more complex tasks. All recent vision backbones, including the various mamba-based ones, report detection results on MS coco and segmentation results on ade20k, which the current work lacks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and comments. We are pleased to hear that you find the motivation of this paper clear and that it holds significant importance for the vision community. We are also glad that you consider the correlation analysis proposed in the paper to be novel and detailed, and that the writing is clear and easy to understand. Below are our responses to each of your comments.\n\n---\n\n> **Q1**: The method is limited by the required annotations (for example foreground annotation) which limits its usefulness across other classification datasets and other tasks.\n> \n\n**A1**: We fully understand your concerns. While it is true that not all datasets or tasks provide foreground annotations, it is important to note that the amount of foreground annotation required for external state defect repair is actually quite minimal. In our experiments, ImageNet-S includes foreground annotations for 9,190 images, with 10 images per class from 919 classes in the ImageNet-1K training set, while approximately 1,200,000 images in the ImageNet-1K training set are not annotated. This means that the amount of required foreground annotation represents a very small fraction of the total (lines 254 to 256 of the paper).\n\nAdditionally, our exploratory experiments show that selectively annotating misclassified images in the training set is more effective than random annotation. Therefore, in practical scenarios, manually annotating approximately 10 challenging samples per class can sufficiently enable external state defect repair for Mamba. Furthermore, for internal state defect repair, the proposed method does not require any additional annotations (lines 242 to 235 of the paper).\n\nIn summary, the proposed method is highly practical and applicable to a wide range of classification datasets and tasks.\n\n---\n\n> **Q2**: The paper does not evaluate on any detection/segmentation benchmarks. This is required to understand whether the features learned by the backbone are rich enough to generalize to these more complex tasks. All recent vision backbones, including the various mamba-based ones, report detection results on MS coco and segmentation results on ade20k, which the current work lacks.\n> \n\n**A2**: Thank you for your valuable feedback. We appreciate your suggestion to evaluate the generalization of the learned features on detection and segmentation benchmarks.\n\nIt is important to clarify that the goal of this work is not to develop a general-purpose backbone. Instead, one of our primary motivations is to investigate the internal mechanisms of the Mamba model from a post-hoc perspective, identify operational flaws in specific tasks, and repair them. This approach focuses on optimizing the model based on interpretability (lines 109 to 119 of the paper). Similar to many interpretability studies [1,2], we chose to validate our methods using widely adopted classification tasks. The results demonstrate that the proposed method effectively identifies and repairs defects in the Mamba model, thereby improving model accuracy.\n\nHowever, we acknowledge that adapting post-hoc interpretability methods to more complex tasks like detection and segmentation presents challenges. Most interpretability research has not yet been validated on detection/segmentation tasks, but this does not imply that our method lacks applicability in these areas. Indeed, the paradigm of our method is broadly applicable to Mamba-like architectures and can analyze defects from both external and internal state perspectives, regardless of the specific task.\n\nFor detection and segmentation tasks, the main challenge is constructing correlations between external/internal states and prediction results. In classification tasks, these correlations are established using the predicted probabilities and gradients between the states and the true classes (Eq. (7) and Eq. (9) in the paper). For tasks like segmentation, where predictions involve probabilities for each pixel in the image, it is necessary to associate external/internal states with all pixels. One potential approach is to aggregate the predicted probabilities for different segmented regions from the ground truth, and calculate the gradients between the total prediction probabilities for each region and the states to construct these correlations. This can be used to analyze and constrain the defects in external/internal state correlations when segmentation results are poor (similar to Eq. (8) and Eq. (9) in the paper). Similarly, for detection tasks, it would be necessary to establish correlations between external/internal states and both the detected categories and spatial coordinates, analyzing these correlations to identify and repair defects.\n\nIn summary, while the proposed post-hoc defect identification and repair paradigm for Mamba-like architectures is applicable to complex detection and segmentation tasks, specific modifications are needed due to the unique nature of task outputs. This remains a focus of our current and future work. We appreciate the reviewer's concerns regarding the validation on classification tasks and hope this response addresses your concerns.\n\n[1] Ali, Ameen, Itamar Zimerman, and Lior Wolf. \"The hidden attention of mamba models.\" *arXiv preprint arXiv:2403.01590* (2024).\n\n[2] Jafari, Farnoush Rezaei, et al. \"MambaLRP: Explaining Selective State Space Sequence Models.\" *arXiv preprint arXiv:2406.07592* (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": "BJUl1GA5iv", "reviewer": "Reviewer_JDBm", "summary": "The paper introduces a novel post-hoc optimization strategy for existing Vision Mamba architectures, termed Vision Mamba Mender, aimed at enhancing the performance of Mamba models in visual recognition tasks. The authors seek to identify and rectify flaws in the Mamba model’s mechanisms from both external and internal hidden state perspectives, proposing a state correlation analysis method. Through this analysis, the authors pinpoint critical modules requiring repair and subsequently design a state correlation constraint method. The authors conduct extensive experiments on several Vision Mamba architectures to validate the efficacy of the proposed optimization strategy.", "questions": "1. In Eqns. (8) and (10), the two terms are multiplied together rather than averaged. What is the rationale behind this design choice?\n2. In Figure 4, why are both the ViM and VMamba structures adopted to obtain these observations?\n3. The loss function defined in Eq. (11) seems incorrect. Could the authors clarify this?\n4. In Table 1, it appears that the performance of \"Internal Flaw Repair\" is better than that of \"External Flaw Repair.\" Are the authors aware of this, and do they have any thoughts on why this is the case?", "limitations": "The limitations appear to have been discussed in Section 6 by the authors. Moreover, there is no societal impact on the work performed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper proposes a novel and intriguing method to enhance state space models from a post-hoc perspective for Mamba. The state correlation analysis used to identify flaws in Mamba, and the subsequent state correlation constraint employed to rectify these flaws, are both new explorations for this kind of technique.\n2. The proposed method for identifying and rectifying flaws in Mamba is systematic to some degree. Moreover, the experiments are clearly designed and well-executed, demonstrating that the proposed method can be applied to mainstream Mamba architectures, further improving their performance.\n3. The paper is well-written, with clear language and a well-structured presentation. The background and motivation are clearly articulated and convincing. Each part of the method is presented step-by-step and in detail, allowing readers to effectively grasp the core ideas.\n4. The paper also provides some insights into the working mechanisms of Mamba, having the potential to inspire future research on the explanation of Mamba models.", "weaknesses": "1. Before reading the methodology section, the description of \"certain modules of Mamba\" (Lines 65-67) may be confusing to readers. The authors should avoid using such vague terms.\n2. The results in Table 1 show that the proposed method improves the performance of Mamba by 1% to 5%. For the sake of rigor, the description in Line 71, stating that the performance is \"enhanced by 3% to 5%,\" does not seem very precise.\n3. In Lines 224-229, the authors claim that the states $c_n^{(l)}$ and $s_n^{(l)}$ in the shallowest blocks are crucial for influencing the model’s anomalous decisions. However, this cannot be observed from Fig. 4, as $c_n^{(l)}$ and $s_n^{(l)}$ do not appear to vary most significantly from the simple sample (Fig. 4a) to the hard sample (Fig. 4b).\n4. In addition to comparing the state correlation scores before and after flaw repair (Fig. 5), the authors should provide more visual comparisons of the state correlation (like Fig. 2 and Fig. 3), which would be more intuitive for the readers.\n5. Only analyzing and enhancing state space models from a post-hoc perspective for Mamba may not be enough.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s recognition of the novelty and interest of our proposed method and the acknowledgment of its contribution to the interpretability research of the Mamba model. Below are our responses to each of your comments.\n\n---\n\n> **Q1**: Before reading the methodology section, the description of \"certain modules of Mamba\" (Lines 65-67) may be confusing to readers. The authors should avoid using such vague terms.\n> \n\n**A1**: We apologize for the confusion caused by this description. Lines 65-67 refer to \"certain modules of Mamba,\" which specifically denotes modules within a Mamba block that exhibit defects after state correlation analysis. For example, the module identified with defects through internal state correlation analysis is $x_{n}^{(l)}$. We will update this section in the paper to clarify this description and avoid any confusion.\n\n---\n\n> **Q2**: The results in Table 1 show that the proposed method improves the performance of Mamba by 1% to 5%. For the sake of rigor, the description in Line 71, stating that the performance is \"enhanced by 3% to 5%,\" does not seem very precise.\n> \n\n**A2**: We apologize for this typographical error and have corrected the description in Line 71.\n\n---\n\n> **Q3**: In Lines 224-229, the authors claim that the states  $c_n^{(l)}$ and $s_n^{(l)}$  in the shallowest blocks are crucial for influencing the model’s anomalous decisions. However, this cannot be observed from Fig. 4, as $c_n^{(l)}$ and $s_n^{(l)}$  do not appear to vary most significantly from the simple sample (Fig. 4a) to the hard sample (Fig. 4b).\n> \n\n**A3**: Yes, you are correct. Figure 4a and 4b only show that for hard samples, the external state correlation scores decrease across all states. However, since only Conv and SSM participate in external interactions with states within a Mamba block (lines 121-123 of the paper), we believe that states $c_n^{(l)}$ and $s_n^{(l)}$ have the most significant impact on the model’s defective decisions. It is important to note that subsequent ablation experiments show that defect repair on states $c_n^{(l)}$ and $s_n^{(l)}$ leads to improved model performance (Appendix E).\n\n---\n\n> **Q4**: In addition to comparing the state correlation scores before and after flaw repair (Fig. 5), the authors should provide more visual comparisons of the state correlation (like Fig. 2 and Fig. 3), which would be more intuitive for the readers.\n> \n\n**A4**: Thank you for your suggestion. We found that, in addition to the improvement in state correlation scores after flaw repair (as shown in Figure 5), the visual morphology of the states, similar to Figures 2 and 3, also changes. Specifically, after defect repair, the difficult samples in Figure 2 exhibit a greater focus on the foreground compared to before repair. We will include these additional visual comparisons in the appendix of the paper.\n\n---\n\n> **Q5**: Only analyzing and enhancing state space models from a post-hoc perspective for Mamba may not be enough.\n> \n\n**A5**: Thank you for your insightful comment; we understand your concerns. In addition to post-hoc analysis and enhancement of the Mamba model, there are indeed preemptive approaches to strengthen the Mamba model's architecture, such as improving its scanning mechanisms. However, these methods often require strong priors and extensive trial-and-error (lines 32-38 of the paper).\n\nThe primary distinction of our work, compared to existing studies, lies in its post-hoc analysis and optimization of the Mamba model. This approach not only provides a clear understanding of Mamba's operational mechanisms but also allows for targeted and precise solutions to the issues encountered during inference. We believe that the insights gained from post-hoc analysis will better assist researchers in designing more refined Mamba architectures, which is one of the key contributions of this work.\n\nMoreover, our paper goes beyond analyzing and enhancing the state space model operations within Mamba blocks by also addressing various linear mapping operations, including general 1D convolutions and gating mechanisms within the Mamba blocks. This analysis and enhancement are comprehensive and thorough.\n\nWe hope this response addresses your concerns. Exploring additional optimization methods for the Mamba model and enhancing the applicability of our proposed methods to more complex tasks will be a continued focus of our future research.", "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:
Dear Reviewers,

Thank you for the time and effort you have invested in reviewing our paper. We are particularly grateful for your recognition of the novelty and originality of our work. We are also pleased that our approach to analyzing and optimizing the Mamba model from a post-hoc perspective has contributed to a better understanding of this new model and its performance improvements.

In addressing each of your comments and concerns, we believe that the paper has been significantly improved. We have individually responded to each comment and collected the following common issues raised by reviewers. If you find that our replies address your concerns, we would be grateful if you could consider raising the score.  Should you have any further questions, we are more than happy to engage in additional discussions.

---

>  Reviewer JDBm expressed concerns about the details of the flaw identification method for the Mamba model and the limitations of analyzing and enhancing the Mamba model solely from a post-hoc perspective.
> 

In response to these concerns, we have provided a detailed analysis and explanation of the defect identification methodology in our responses to the reviewer’s comments. Additionally, we have highlighted the significant advantages and importance of analyzing and enhancing the Mamba model from a post-hoc perspective compared to other approaches. We believe these responses address the reviewer’s concerns effectively.

---

> Reviewers nvv2 and 3tmX both raised concerns regarding the relevance of foreground information in the identification and repair of external state flaws. Reviewer nvv2 pointed out that classification tasks do not solely rely on foreground information, as background regions can sometimes provide useful data. Reviewer 3tmX noted that the requirement for foreground annotations limits the applicability of the proposed method to other classification datasets and tasks.
> 

Thank you for the valuable feedback. We apologize for any confusion caused. We have provided detailed responses to each reviewer’s comments below and believe these responses address their concerns. It is important to emphasize again that the proposed method for identifying and repairing external state flaws does not restrict the Mamba model from learning background information. Furthermore, the amount of foreground annotation required for repairing external state flaws is minimal, making the proposed method both reasonable and practical.

---

> Reviewer 3tmX expressed concerns about the lack of evaluation of backbone generalization on more complex tasks.
> 

We understand the reviewer’s concerns. To address this, we have provided detailed responses regarding why the proposed method was validated only on classification tasks, how the method can generalize to more complex tasks, and how it can be adapted for such tasks. In summary, the goal of this work is not to develop a universally applicable backbone for various complex tasks. Instead, one of the motivations for this work is to analyze the internal workings of the Mamba model from a post-hoc perspective, identify specific operational flaws in certain tasks, and repair them. The paradigm of identifying and repairing flaws from both external and internal state perspectives is applicable and valuable.

---

> Reviewers 4vzy, nvv2, and JDBm pointed out issues with the explanations of formulas and some typographical errors in the paper.
> 

We appreciate the reviewers' attention to these details. We have corrected these issues in the paper and thoroughly checked the entire text to ensure accuracy.

---

Besides the issues mentioned above, more detailed responses to the reviewers' comments can be found below each reviewer's specific feedback. We sincerely thank all the reviewers for their diligent review and valuable suggestions, which have greatly improved the paper.

Meta-review: This paper proposes to explore the internal learning mechanism of mamba, identifying mamba flows and thus optimizing the design. It received 3 positive reviews and 1 negative review initially. The raised concerns include method limitations, flaw assumption validation, only classification results provided. In the rebuttal phase, the authors have addressed several issues by responding in details. Finally, the remaining issue is not experimenting object detection / segmentation scenarios. The AC has checked all files, and agree to reviewers that the current form shall raise reporting. The authors are suggested to provide det/seg results in the camera-ready version. Congratulations.
Decision: Accept (poster)