Vision Mamba Mender

Authors: jiacong  hu
anda  cao
zunlei  feng
shengxuming  zhang
yi  wang
lingxiang  jia
mingli  song
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 18:11:33

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.

Main paper: 
 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 Transformer [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] [10] [11] [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 complexity while achieving a global receptive field, leading to outstanding performance. Consequently, it has been widely adopted in the field of computer vision [14] [15] [16] .(SSM). Unlike the previous time-invariant SSM [11] , it allows the parameters of the SSM to depend on the input, enhancing the model's expressive capacity. Furthermore, inspired by H3 [24] and Gated MLP [25] , each Mamba block also incorporates modules such as causal convolution and gating mechanisms. As illustrated in Figure 1 , given the hidden state h (ℓ)-1 n of the i-th token in the input of the ℓ-th Mamba block, the computation of the hidden state h (ℓ) n of the i-th token in the output of the ℓ-th Mamba block is as follows:x (ℓ) n = SiLU(h (ℓ)-1 EQUATIONEQUATIONEQUATIONwhere SiLU(.), causal-Conv1D(.), and selective-SSM(.) denote the activation function, the casual 1D convolution, and the selective state model, respectively. W (ℓ)x , Wz , andW (ℓ)y 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. Finally, by applying a residual connection, the hidden state h (ℓ) n is obtained as follows:EQUATIONIn Eqn. (5) , ⊙ denotes the Hadamard product. The term z (ℓ) n in Eqn. ( 4) is computed through a separate pathway, serving as a gating mechanism to regulate the information flow in the main pathway.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 requirements, such as ViM [17] , VMamba [18] , PMamba [19] , Mamba-ND [20] , and SiMBA [21] . Unlike the predefined optimizations, our approach analyzes the working mechanism of the Mamba model post-hoc, identifying flaws and making repairs to further enhance the model's performance. 3 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 of Mamba models in the NLP domain, such as 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 a significant challenge. Ali et al. [29] established a connection between the selective SSM within the Mamba block and the self-attention mechanism in Transformers, allowing the selective SSM to represent the interaction process between any two states by constructing a 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, as causal convolution, as 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 computational processes within Mamba as state interactions 1 . 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 a state interacts only with its internal information to form a new state, as shown in Eqn. ( 1), (4), and (5) . We explore the operational mechanisms of Mamba from both the external and internal state interaction perspectives to identify flaws. n from the selective-SSM module in different ViM [17] blocks. The depth of the ViM blocks increases from left to right. 
external state correlation analysis: 
 To understand the impact of external states on 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 decision-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 k of the input sample, the calculation process for external state correlation e (ℓ,s) ∈ R H×W (where H and W represent the height and width of the input image) of the states {s EQUATIONwhere g (ℓ,s) ∈ R D represents the weights used to weight the state sn along its dimensions based on gradient information, E D (•) denotes the average across all dimensions, ŝ(ℓ) n ∈ R signifies the degree of correlation between the n-th state and the prediction after integrating gradient information, and R 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 s (ℓ) n as an example, we visualize the external state correlations of s(ℓ)n in different blocks of ViM [17] . It can be observed that the external state correlations of s (ℓ) n 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 a pre-trained Mamba model F(•), an input image i, foreground annotations m of the input image, and external state correlation e (ℓ,s) computed through the proposed method, the external correlation score is defined as follows:ECS(e (ℓ,s) ) = softmax(F(e (ℓ,s)+ ⊙ i)) softmax(F(e (ℓ,s)-⊙ i))EQUATIONwhere e (ℓ,s)+ = 1(e (ℓ,s) > α) denotes the retention of important regions (those greater than the threshold α) in e (ℓ,s) , setting them to 1, and the removal of unimportant regions (those less than the threshold α), setting them to 0, for use in negative perturbation tests. Conversely, e (ℓ,s)-= 1(e (ℓ,s) < α) denotes the removal of important regions and the retention of unimportant regions, for use in positive perturbation tests. Additionally, ECS(e (ℓ,s) ) ∈ [0, +∞), where a higher value indicates a higher correlation score. n from the linear mapping module w ℓ s in ViM [17] for samples of the same class. The horizontal axis represents the state dimensions, and the vertical axis represents the samples.To comprehend the influence of internal states on prediction outcomes within the operational mechanism 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 as 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 x ℓ n ∈ R D , which is the output of the linear mapping matrix w ℓx . The computation process for the corresponding Internal State Prediction Correlation iEQUATIONwhere g (ℓ,x) n ∈ R D denotes the weights applied to the dimensions of state x (ℓ) n using gradient information. Similarly, the Grad-ISC method can perform internal state correlation analysis on the output of any module within the Mamba block.Continuing with the example of state x (ℓ) n , we visualize the internal state correlations i (ℓ,x) n for samples belonging to 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 of 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 a novel definition for the Internal Correlation Score: Definition 2 (Internal Correlation Score). Given J samples belonging to the same class and the internal state prediction correlation i (ℓ,x) n ∈ R D computed using the proposed method for a particular sample, the Internal Correlation Score is defined as follows:EQUATIONwhere i(ℓ,x)n,j represents the internal state correlation for the j-th sample, i(ℓ,x)+ n,j = 1(i (ℓ,x) n,j> β) denotes the binarized internal correlation (set to 1 if greater than β, otherwise set to 0), ⊕ denotes the XOR operation, and E D (•) denotes the average over all dimensions. Furthermore, ICS(i (ℓ,x) ) ∈ [0, +∞), where a higher value indicates a higher correlation score.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. n are all better than those of state y (ℓ) n , 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 to associate certain incomprehensible regions in the external states.Correlation. To unveil patterns in internal state correlations and detect flaws within them, we also conducted a 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 x (ℓ) n is generally better than that of other states. Furthermore, simultaneous comparison of simple and difficult samples reveals a decrease in the internal correlation scores of all states, including state x (ℓ) n . This suggests that for difficult samples, the Mamba model tends to correlate with some incomprehensible regions within the internal states. 
how to repair flaws?: 
 In principle, repairing flaws within the model's internals can 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 in 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.To repair flaws related to external state correlations, it is essential to identify the key components within the Mamba block that need fixing. In flaw identification regarding external state correlations, it is observed that the states x(ℓ) n , c (ℓ) n , s (ℓ)n , and z (ℓ) n 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 predictions and with background regions in incorrect predictions. Furthermore, considering that external state interactions occur only within Conv and SSM modules (as indicated by Eqn.(2) and Eqn.( 3)), 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 on repairing the external correlation flaws of the states c EQUATIONwhere E HW 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 e 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:EQUATIONwhere Loss ce denotes the cross-entropy loss for the image recognition task, and λ serves to balance the magnitudes of the respective loss components.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 x (ℓ) n for simple samples are higher than those for difficult samples. This suggests that in predictions biased toward correctness within a 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 to focus on internally inconsistent regions of the state with higher overall complexity. In our experiments, we consider the linear mapping w ℓ s within deeper blocks as a 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 x for each class.We then utilize these templates to constrain the internal correlations of x (ℓ)n :EQUATIONSimilarly, during backpropagation, each term in the computational graph of Loss i is also differentiable, which involves second-order gradients as specified in Eqn. (9) . Combined with the loss for the original task, the complete loss is as follows:EQUATION)where γ 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)).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] , EfficientVMamba-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: ImageNet-50, ImageNet-300, and ImageNet-1K. 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] . n . 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.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 EQUATION 
related works: 
 Model Optimization. Optimizing model performance through various prior designs or theoretical derivations has always been a 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] [18] [19] [20] [21] [22] . For instance, improvements have been made through the bidirectional scanning mechanism [17] , cross-scanning mechanism [18] , continuous 2D scanning mechanism [19] , and incorporating new channel modeling techniques [21] . However, these enhancements are predesigned and require extensive trial and error. For a more detailed understanding of Mamba's applications in computer vision, we recommend readers refer to recent surveys by Zhang et al. [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 to 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] [41] [42] 4] , feature enhancement [43] [44] [45] [46] , or some post-hoc debugging methods [47] [48] [49] [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-hoc 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.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 on Attention-Rollout [51] and Transformer-Attribution [31] , which can be used for image feature attribution. However, this method is limited to the selective SSM module and does not consider or apply to other modules within the Mamba model. In research on model explainability beyond the Mamba model, a plethora of methods have emerged, including activation-based methods [52] [53] [54] [55] [56] , gradient-based methods [30, 57, 33] , LRP-based methods [58, 59, 31] , and perturbation-based methods [60] [61] [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. 
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 in visual tasks may vary, it 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. Moreover, at this juncture, considering the potential changes in future architectures, devising a 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 postoptimization paradigms should be further explored in non-visual tasks within Mamba or even applied in non-Mamba architectures. 
conclusion: 
 Mamba incorporates intricate computational modules, making it non-trivial to delve into its operational mechanisms. In this paper, we have proposed a 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 states. These methodologies aim to identify flaws within Mamba's operational mechanisms. Simultaneously, we 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.A Vision Mamba Mender. The magnifying glass held in one hand symbolizes flaw detection, while the wrench held in the other hand symbolizes flaw repair. 
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 is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (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 to the citations in the paper.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
 
 Appendix: 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 a CPU with 24 cores and 500GB of memory.Model Parameter Settings. To enable the Mamba model to train efficiently within limited computational 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, 3, 2] , and we introduced a class token for classification in the last block. For EfficientVMamba-T, we similarly set the patch size to 16x16. In the case of LocalViM-T, we reduced the model depth from 20 to 9 and set 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 a momentum of 0.9, a total batch size of 128, and a weight decay of 0.1. We utilized a 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 on 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 1e + 7 by default, and the balance weight γ for the loss function of internal state flaw repair is also set to 1e + 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 F.As shown in Figure 6 , we compared the impact of different values of threshold α on 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 xEQUATIONn , and y (ℓ) n 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 x n . Therefore, we recommend that the value of α should not be too extreme when calculating external state correlation scores. It can be set between 0.4 and 0.6; we use a default setting of α = 0.5.We also examined the impact of different values of threshold β on internal state correlation scores, as shown in Fig. 7 . Similarly, by comparing simple and difficult samples, it can be observed that the value of β does not affect the relative magnitudes of the internal correlation scores. Specifically, regardless of whether β is set to 0.0, 0.2, or 0.8, the internal correlation scores for states xEQUATIONn , s(ℓ) n , z (ℓ)n , and y (ℓ) n 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 internal correlation scores of state x (ℓ) n are better than those of states cEQUATIONn , and y (ℓ) n . Therefore, we recommend that the value of β should not be too extreme when calculating internal state correlation scores. It can be set between 0.4 and 0.6; we use a default setting of β = 0.5.As shown in Figure E(a) , we conducted a study on repairing external flaws in different blocks, revealing a clear enhancement in model performance following external flaw repair. Specifically, when external flaws are repaired in the final block, the model'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 an accuracy increase, the improvement is only 1.72%. This discrepancy may be attributed to the presence of external flaws in the states across nearly every block, as illustrated in Figure 2 , including states such as x As shown in Figure E(b) , we conducted a study on repairing internal flaws in different blocks, which also demonstrates a noticeable improvement in model performance. Specifically, when internal flaws are repaired in the final block, the model's accuracy increases by 1.76% compared to the baseline. In contrast, repairing internal flaws in other blocks, such as the second-to-last block, results in an accuracy improvement of 2.4%. Furthermore, repairing internal correlations in 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 on model performance across the various blocks.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. 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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).In Mamba models used for visual recognition, each state corresponds to an image patch or a learnable classification token.We adjusted the experimental setup and refined methodological details to enhance the fairness and reproducibility of the comparisons. Detailed experimental settings can be found in Appendix B. 


Decision: 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.

Reviews: review: paper_summary: paper_summary: 1. Addressing the main issue in the existing model:
   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.

2. Proposed solution for the issue:
   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.

3. How did the author evaluate the proposed solution on various standard dataset:
   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.
strength_weakness: strength_weakness: 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.

2. 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.
questions: questions: Please refer to the Weaknesses block
limitations: limitations: Please refer to the Weaknesses block
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. Only classification on ImageNet results are provides. It it suggented to provide more diverse experiments to full evaluation the proposed repair method.

3. Minors:

(1) Not all symbols are explained from Eq.(1)~(6).

(2) Missing punctuation (comma or full stop) in Eq.(1)~(6).

(3) In Fig.5, there are only (a)(b)(c)(d) in subfigures, while authors use (1)(2)(3)(4) in the caption.
questions: questions: see weakness
limitations: limitations: see weakness
paper_summary: paper_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.
strength_weakness: strength_weakness: There are two major concerns here:

1. The method is limited by the required annotations (for example foreground annotation) which limits its usefulness across other classification datasets and other tasks.
2. 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.
questions: 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. 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).
4. 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.
5. Only analyzing and enhancing state space models from a post-hoc perspective for Mamba may not be enough.
questions: questions: 1. In Eqns. (8) and (10), the two terms are multiplied together rather than averaged. What is the rationale behind this design choice?
2. In Figure 4, why are both the ViM and VMamba structures adopted to obtain these observations?
3. The loss function defined in Eq. (11) seems incorrect. Could the authors clarify this?
4. 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: limitations: The limitations appear to have been discussed in Section 6 by the authors. Moreover, there is no societal impact on the work performed.

Comments: title: Vision Mamba Mender
Thanks for response
Rebuttal by Authors [Q6-Q9]
Rebuttal by Authors [Q3&Q4]
comment: Thank you for your diligent review and comments. We are pleased that our responses have addressed your concerns.  We appreciate your support for this work.
Thank you for your thorough review and valuable suggestions. We are pleased that our responses have addressed your concerns. We will incorporate these points into the final version of the manuscript. Once again, we appreciate your contribution to improving this work.
Dear Reviewer 3tmX and Reviewer nvv2,

Thank you both for your continued engagement and thoughtful feedback.

We appreciate  Reviewer 3tmX’s perspective on the importance of evaluating the generalizability of features learned by our method through detection and segmentation tasks. While we understand the significance of such experiments, we would like to reiterate that the primary focus of our study is to analyze the internal mechanisms of the Mamba model from a post-hoc interpretability standpoint, specifically within the context of classification tasks.

As highlighted by Reviewer nvv2, the primary objective of our work is not to develop or evaluate a general-purpose backbone, but rather to perform a detailed post-hoc analysis of the Mamba model. Our focus is on understanding and enhancing specific aspects of the model's internal mechanisms, particularly in the context of classification tasks, which is a widely accepted approach in related studies. Our choice to validate on classification tasks aligns with the goals of the study, as it allows us to systematically identify and repair flaws within the model, thereby improving its accuracy.

That said, we fully acknowledge that extending our evaluation to include detection and segmentation tasks is an important future direction. Furthermore, as mentioned in our response to your second question (Q2), the implementation of post-hoc analysis on the Mamba model in classification tasks differs somewhat from that in detection or segmentation tasks. We believe that our internal mechanism analysis and optimization paradigm hold potential for application in more complex tasks as well. We are actively working on this as part of our ongoing research.

Once again, we sincerely thank you both for your constructive feedback, and we hope that this clarification helps to contextualize our approach within the intended scope of the paper.

Best regards,

Authors
Dear Reviewer 3tmX and Reviewer nvv2,

Thank you both for your continued engagement and thoughtful feedback.

We appreciate  Reviewer 3tmX’s perspective on the importance of evaluating the generalizability of features learned by our method through detection and segmentation tasks. While we understand the significance of such experiments, we would like to reiterate that the primary focus of our study is to analyze the internal mechanisms of the Mamba model from a post-hoc interpretability standpoint, specifically within the context of classification tasks.

As highlighted by Reviewer nvv2, the primary objective of our work is not to develop or evaluate a general-purpose backbone, but rather to perform a detailed post-hoc analysis of the Mamba model. Our focus is on understanding and enhancing specific aspects of the model's internal mechanisms, particularly in the context of classification tasks, which is a widely accepted approach in related studies. Our choice to validate on classification tasks aligns with the goals of the study, as it allows us to systematically identify and repair flaws within the model, thereby improving its accuracy.

That said, we fully acknowledge that extending our evaluation to include detection and segmentation tasks is an important future direction. Furthermore, as mentioned in our response to your second question (Q2), the implementation of post-hoc analysis on the Mamba model in classification tasks differs somewhat from that in detection or segmentation tasks. We believe that our internal mechanism analysis and optimization paradigm hold potential for application in more complex tasks as well. We are actively working on this as part of our ongoing research.

Once again, we sincerely thank you both for your constructive feedback, and we hope that this clarification helps to contextualize our approach within the intended scope of the paper.

Best regards,

Authors
Dear Reviewer 3tmx,

Please allow me to join your discussion. I agree with the authors that this paper may not need to provide comprehensive experimental results in downstream tasks because this manuscript is more like a machine learning analysis paper. From my side, I think it is acceptable to only provide the results in classification tasks and analyze the inner mechanisms of the model architecture using this basic task. This is the style many machine learning scientists follow in their research.
Dear Authors,

Thank you for your response. My concerns have been addressed partially. However, the main point for the request for detection/segmentation is to understand whether features learned by your method are truly rich enough, or are just over-fitting to the classification task. Which is why it is common practice in literature to use the ImageNet pretrained backbone in detection and segmentation. I believe this experiment is fundamental to the analysis important, and therefore I will be keeping my rating.
The rebuttal addresses all my concerns. I decide to increase my rating from borderline accept to weak accept. The authors should add the points from rebuttal into the final version if this manuscript is finally accepted.
Thank you for your response. Most of my concerns have been addressed in the rebuttal.
Thank you for your valuable feedback and comments. We have provided our responses to your concerns above. Should you have any further questions or require additional clarifications, we would be more than happy to discuss them with you.
Thank you for your valuable feedback and comments. We have provided our responses to your concerns above. Should you have any further questions or require additional clarifications, we would be more than happy to discuss them with you.
Thank you for your valuable feedback and comments. We have provided our responses to your concerns above. Should you have any further questions or require additional clarifications, we would be more than happy to discuss them with you.
Thank you for your positive feedback. We are pleased that your concerns have been addressed. We greatly appreciate your support for our work.
Thank you for your response. My concerns have been mainly addressed. I have also reviewed the comments from other reviewers and the corresponding responses. I have increased my score to 6.
> **Q6**: In Eqns. (8) and (10), the two terms are multiplied together rather than averaged. What is the rationale behind this design choice?
> 

**A6**: The rationale behind this design choice is that the units and scales of the terms in these equations are different. Averaging them could lead to information loss and might not adequately reflect the importance of each term. Multiplying the terms amplifies the differences between them and highlights the significance of each term more effectively.

---

> **Q7**: In Figure 4, why are both the ViM and VMamba structures adopted to obtain these observations?
> 

**A7**: We sincerely apologize for the confusion regarding the description of Figure 4. In reality, all observations in Figure 4 are based on the ViM structure. As detailed in lines 196-198 and 204-206 of the paper, Figure 4(a,b) shows the external state correlation scores for different states within Mamba blocks for simple and hard samples, respectively. Similarly, Figure 4(c,d) displays the internal state correlation scores for these states in the same Mamba blocks for simple and hard samples. The x-axis of each subplot represents different Mamba blocks, and the y-axis indicates the magnitude of the correlation scores. We will correct the description of Figure 4 in the paper.

---

> **Q8**: The loss function defined in Eq. (11) seems incorrect. Could the authors clarify this?
> 

**A8**: We apologize for the typographical error. The correct loss function should be: $\text{Loss}\_{\mathbf{e}} = \mathbb{E}\_{HW}(\mathbf{e}^{(\ell,c)+} \odot m) + \mathbb{E}\_{HW}(\mathbf{e}^{(\ell,s)+} \odot m)$ We have corrected this in the paper. Thank you for pointing this out.

---

> **Q9**: 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?
> 

**A9**: Great question. The reason why "Internal Flaw Repair" shows better performance than "External Flaw Repair" is primarily due to the more pronounced patterns observed in internal state correlations compared to external state correlations. Specifically, as shown in Figure 4(c,d), the internal state correlation scores for states such as $x_{n}^{(\ell)}$ drop more significantly for difficult samples. Therefore, repairing internal state correlations for $x_{n}^{(\ell)}$ results in more substantial improvements.
> **Q3**: 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.
> 

**A3**: Thank you for your comment. We sincerely hope that the previous response addresses your concerns. It is important to reiterate that flaw identification and repair are post-hoc optimization methods applied to a trained model. Like many studies exploring internal mechanisms of models, the patterns discovered are specific to that model. Consequently, the backbone obtained from flaw identification and repair on a classification task with Mamba may not be directly applicable to detection or segmentation tasks. However, as mentioned earlier, the proposed flaw identification and repair methods are adaptable to detection and segmentation tasks.

---

> **Q4**: 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.
> 

**A4**: Thank you for your valuable feedback. We agree that discussing additional limitations would make the paper more comprehensive. We have added the following content to the paper to better reflect the limitations of our method:

“In the context of external state defect repair, while our method performs well with a small amount of annotations, it is undeniable that, with limited annotation resources, even manual labeling of a few samples may affect the usability of the proposed method. Similar to internal state defect repair, exploring defect repair methods that do not require additional annotations is a direction we will continue to investigate.”

We hope that this response addresses your concerns.