Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training

Thanks for your constructive comments. We are glad that you found our work is technically sound and well-motivated, has extensive experiments with effectiveness. We provide our feedback as follows.

Q1: Applying SLWS to hierarchical structures.

A1: Hierarchical architectures like VMamba employs a more complex downsampling process, which demonstrates strong performance under Tiny, Small, and Base model sizes. However, there is currently no evidence to confirm its scalability to larger model sizes beyond the Base size. Furthermore, the hierarchical structure of VMamba is incompatible with our proposed straightforward shuffle-based regularization method. This is because downsampling at certain layers leads to inconsistent input-output sequence lengths, making it infeasible to implement the "order restoration" step in our method. This limitation is discussed in lines 253–260 of the submitted manuscript.

In contrast, plain vision Mamba models feature a non-hierarchical architecture. 1) This design is simpler, more fundamental, easier to stack, 2) and seamlessly compatible with a wide range of existing sequence modeling frameworks like MLLMs, as well as training paradigm like MAE and masked feature distillation. 3) Such architectures (e.g., Vim, Mamba-R, ARM, PlainMamba) have been widely adopted and benchmarked in prior research. Building upon this foundation, we propose SLWS, a plug-and-play regularization method that further enhances the scalability of these models. Our approach achieves state-of-the-art results on ImageNet among Vision Mamba variants, contributing to the exploration of efficient regularization strategies for plain Mamba architectures.

Q2: Linear layer-wise probability design reasons

A2: The layer-wise probability design is inspired by well-established practices in hierarchical feature learning. For instance, Stochastic Depth[1] adopts a linear probability schedule across layers, guided by the intuition that deeper layers handle higher-level semantics and are more robust to structural variations. Similarly, our linear probability design aligns with the inherent hierarchy of visual processing. To assuage reviewers' concern, we provide a further dedicated design for the layerwise probability experiment here (i.e., $p_{\ell}=P_L^{(L-\ell+1)}$, similar to the layer-wise learning rate decay form, where $P_L=0.5$) and the corresponding test accuracy is 82.2, which is lower than the original linear setting of 82.7 in Table 5. We think that it is due to the fact that the power function form of the design leads to too low a probability for the middle layers, leading to inadequate regularization.

Q3: Applying SLWS to diverse training paradigms and their horizontal and vertical comparisons.

A3: As demonstrated in Table 1, we conducted comprehensive experiments to validate the effectiveness of the SLWS across non-hierarchical Mamba architectures and diverse training paradigms, including: naive supervised classification training, self-supervised MAE pretraining, and Masked feature distillation (MFD). These results and common improvements confirm that SLWS can be compatible with the training paradigms. For example, SLWS even elevates the accuracy of the previously collapse-prone Vim-L model to a 84.5%, and applying the shuffling strategy to the MAE-style pretraining brings a 0.4 points gains.

---

Thank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed your remaining concerns.

[1] Deep networks with stochastic depth, ECCV'16.

Thanks for your constructive comments. We are glad that you found our work significantly improves the performance of non-hierarchical Mamba models and does not incur significant computational overhead. We provide our feedback as follows.

Q1: SLWS's shuffle effects for visual connections between tokens.

A1: Mamba's scanning mechanisms inherently restrict token interactions to 1-D adjacency within the sequence, which is misaligned with the 2-D structural priors of images. Our shuffle operation introduces randomness to enable non-local token interactions as regularization, effectively sampling from the full $O(n^2)$ dependency space while preserving efficiency. Critically, we implement three safeguards to maintain positional coherence: 1) Positional encodings explicitly retain the inherent locality of the original image data. 2) Layer-wise shuffle probability: Deeper layers (handling global features) adopt higher shuffle probabilities, while shallow layers (processing low-level information) is more likely to remain unshuffled. 3) Order restoration: The original sequence order is restored after each shuffled layer, preventing recursive disruption of positional relationships. 4) Experimental results demonstrate both the rationality and effectiveness of SLWS.

Q2: Training resolution setting and performance difference between training pipelines.

A2: The 224 resolution is a common training setting adopted by works like Vim, ViT, and MAE. ARM uses 192 due to its unique design requirement of dividing images into multiple 64×64 patch groups. To ensure fair comparison under the same training epochs, their MAE experiment also followed this resolution. For MFD pre-training, we used the standard 224 resolution but with only 18.75%-37.5% of ARM's training epochs, significantly reducing computational costs. Meanwhile, MAE was trained with the same epochs as that in ARM. Thus, the comparisons are fair, and our method consistently achieves improvements across different training strategies.

Q3: Hierarchical VMamba's effect in mitigating overfitting.

A3: VMamba employs a more sophisticated downsampling process and adopts a hierarchical structure, which demonstrates strong performance in smaller model scales (e.g., Tiny, Small, and Base). However, there is no evidence to suggest that this design inherently benefits training at larger scales, as VMamba has only been extended up to the Base size. Further investigation would be required to validate its effectiveness in mitigating overfitting for larger models. Please also see details in our response to Reviewer mf21 (response 1 & 2).

Q4: Performance improvement based on MambaR.

A4: Mamba-R successfully scales Mamba to the Large size under supervised training but requires the addition of extra register tokens to the original architecture, which brings overhead in training and inference. In contrast, SLWS achieves comparable results without any architectural modifications and overhead in inference. Our experiments on Mamba-R aim to demonstrate compatibility with prior techniques, and SLWS still outperforms it by 0.5 points in segmentation tasks, further validating its effectiveness.

Q5: Horizontal and vertical comparisons of MAE, ARM and MaskDistill.

A5: Our comparisons encompass both MAE and MFD training with and without SLWS regularization, and the results demonstrate that SLWS consistently improved performance under both training paradigms. When evaluating cross-strategy performance, MFD indeed benefits from CLIP’s rich semantic knowledge, outperforming MAE. Furthermore, as shown in Table 2, existing self-supervised methods (e.g., ARM’s MambaMLP-L at 84.5) lagged significantly behind MaskDistill’s ViT-L (87.6). However, our MFD-trained MambaMLP-L achieves 86.7, which narrows the gap with MaskDistill’s ViT-L (87.6), showcasing substantial progress.

Q6: Reversed layer-wise probability experiment for SLWS's semantic awareness priori.

A6: Following your suggestion, we conducted experiments with a reversed layer-wise probability strategy (assigning larger shuffle probabilities to shallower layers and smaller ones to deeper layers). The results showed a performance drop of 1.5 points on ImageNet as the table below. Thus, our original statement about semantic-awareness still holds.

Suggestions about Figure 2 and related work.

A: Thank for your suggestions and we have revised the legend of Figure 2, as well as included discussed some new related vision backbone literatures like Mamba Vision[CVPR'25], TransNeXt[CVPR'24].

---

Thank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed your remaining concerns.

Thanks for your constructive comments. We are glad that you found our work has clear motivation and comprehensive evaluation, and simple plug-and-play design. We provide our feedback as follows.

Q1: Performance improvements of our models compared to baselines especially Mamba-R

A1: Our proposed SLWS regularization demonstrates significant improvements across both supervised and self-supervised training paradigms. For example:

• It enables successful supervised training of the previously collapse-prone vanilla Vim-L model without architectural modifications, achieving a accuracy of 84.5% on ImageNet.
• The MambaMLP-H model trained with SLWS reaches 87.5% accuracy, setting a new state-of-the-art result for vision Mamba variants.

These advancements have been acknowledged by reviewers, who highlighted the "impressive results," "measurable improvements," and "significant performance gains" in their feedback.

Regarding Mamba-R, it is an existing method that successfully scales Mamba to Large size under supervised training but requires the addition of extra register tokens to the original architecture. In contrast, SLWS achieves comparable results without any architectural modifications. Our experiments on Mamba-R aim to demonstrate compatibility with prior techniques, and SLWS still outperforms it by 0.5 points in segmentation tasks, further validating its effectiveness.

Q2 & Q3: Concerns about disrupting inherent locality and positional information of image data and dense prediction

A2 & A3: Mamba's scanning mechanisms inherently restrict token interactions to 1-D adjacency within the sequence, which is misaligned with the 2-D structural priors of images. Our shuffle operation introduces randomness then enables some global token interactions, effectively sampling from the full $O(n^2)$ dependency space while preserving efficiency. Critically, we implement three safeguards to maintain positional coherence: 1) Positional encodings explicitly retain the inherent locality of the original image data. 2) Layer-wise shuffle probability: Deeper layers (handling global features) adopt higher shuffle probabilities, while shallow layers (processing low-level information) remain unshuffled. 3) Order restoration: The original sequence order is restored after each shuffled layer, preventing recursive disruption of positional relationships.

For dense prediction tasks (e.g., segmentation), SLWS is used for backbone pre-training and achieves improved performance as listed in Table 3, demonstrating both the rationality and effectiveness of SLWS.

Q4 & Q5: Applying SLWS to hierarchical variants and ViT

A4 & A5:

• Application to Hierarchical Variants:
  Hierarchical architectures (e.g., VMamba) incorporate downsampling layers, which alter feature map dimensions across stages. This disrupts the sequence length consistency required for SLWS's order restoration step, a critical process for SLWS. Please also see details in lines 253–260 of the manuscript and our response to Reviewer mf21 (response 1&2).

• Application to Vision Transformers (ViTs):
  ViTs inherently perform global $O(n^2)$ interactions via self-attention. Token shuffling in ViTs would not meaningfully alter their ability to model long-range dependencies. As self-attention is permutation-equivariant, shuffling input tokens does not change the attention output.

Q6: Shuffling strategies with locality

A6：Other shuffling strategies with localization might be useful, such as implementation within a certain window. However, this will significantly increase the complexity of the implementation, and we believe that layer-wise probability settings are also necessary, as this is important for conforming to the deep model semantic hierarchy prior.

Q7 (review content in Claims And Evidence): "The claim that vanilla Vision Mamba models couldn't previously be scaled up"

A7：We appreciate the reviewer’s attention to this point. But our manuscript does not have such claim. Instead, we explicitly acknowledge in the Introduction (lines 36–38) and Related Work sections that "a limited number of ... strategies [ARM, MambaR, MAP] have successfully trained and scaled certain Mamba-based models to Huge sizes."

---

Thank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed your remaining concerns.

Q1 & Q2: Hierarchical VMamba training and non-hierarchical plain mamba selection.

A1: VMamba employs a more complex downsampling process and adopts a hierarchical architecture, which demonstrates strong performance under Tiny, Small, and Base model sizes. However, there is currently no evidence to confirm its scalability to larger model sizes beyond the Base size.

A2: In contrast, plain vision Mamba models feature a non-hierarchical architecture. 1) This design is simpler, more fundamental, easier to stack, 2) and seamlessly compatible with a wide range of existing sequence modeling frameworks like MLLMs, as well as training paradigm like MAE and masked feature distillation. 3) Such architectures (e.g., Vim[ICML'24], Mamba-R[CVPR'25], ARM[ICLR'25], PlainMamba[ECCV'24]) have been widely adopted and benchmarked in prior research. Building upon this foundation, we propose SLWS, a plug-and-play regularization method that further enhances the scalability of these models. Our approach achieves state-of-the-art results on ImageNet among Vision Mamba variants, contributing to the exploration of efficient regularization strategies for plain Mamba architectures.

Q3: 220 training epoch setting on some models.

A3: Our training protocol for 220 epochs follows the settings of Mamba-R to ensure a fair comparison. Mamba-R adopts a three-stage training strategy, which its paper states is equivalent to approximately 220 epochs of training under a 224 input resolution. Therefore, we adhered to its protocol to demonstrate that our SLWS method is compatible with the register-based framework proposed in their work. This approach validates the seamless integration of our regularization technique with existing methodologies.

Q4: Training configuration of our models.

A4: Beyond SLWS, we did not employ any special additional tricks, modules, or techniques to alter the training pipeline. All training configurations strictly adhere to the essential settings previously used in vision Mamba models, ensuring the validity of comparisons. Detailed training protocols are provided in the Appendix tables.

Thank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed all of your remaining concerns.

Thanks for your constructive comments. We are glad that you found our work is simple but demonstrates significant benefits, and has measurable performance gains across diverse tasks and training paradigms. We provide our feedback as follows.

Q1-1: Supporting evidence of learning curves for mitigating overfitting of SLWS.

A1-1: We are glad to indicate that the Figure 2 in our submitted manuscript is exactly what you need, i.e., the training and validation loss curves with and without using SLWS. As the analysis stated in line 318-324 in the right column of the paper, the model trained with SLWS stabilizes at a higher training loss yet achieves a lower evaluation loss and a better final accuracy, which implies the effectiveness of mitigating overfitting. This type of curve are also adopted by some classical work like ResNet and Stochastic Depth[1].

Q1-2: Qualitative examples for SLWS.

A1-2: Qualitative visual analysis could help better understand the training strategy. However, since Mamba lacks attention scores like those in ViT, such visualizations are inherently challenging. To enable effective qualitative analysis and intuitive insights, we provide some examples of segmentation outputs comparisons in the anonymous link [https://postimg.cc/GHSPq9Py ]. These visualizations reveal that SLWS pre-trained model achieves more accurate segmentation boundaries when transferred to segmentation task, indicating that SLWS's higher semantic awareness, which is consistent with the quantitative results.

Q2: Direct segmentation comparisons of SLWS vs. non-SLWS

A2: The Vim-M model already exhibited significant disadvantages in classification tasks under non-SLWS training (80.9 vs 82.8), which led us to exclude it from downstream segmentation comparisons. In Table 3, the SLWS-trained MambaR-B achieves a 0.5 higher mIoU when directly compared to its non-SLWS baseline. To further address reviewer concerns, we conducted additional experiments on direct comparison of ViM-B in the table below, demonstrating that non-SLWS-trained models consistently underperform their SLWS-trained counterparts in segmentation tasks. This further validates the performance generalizability and robustness between SLWS classification training and its downstream tasks.

Q3 (review content in Experimental Designs Or Analyses): SLWS’s gains and relationship to training paradigm

A3: Our study provides comprehensive experimental results by applying SLWS regularization to various Mamba models. This includes evaluations under: Supervised classification training, Self-supervised MAE , and Masked Feature Distillation (MFD)-based comparisons. As the reviewer mentioned, SLWS consistently achieves measurable improvements across these frameworks. These results demonstrate that the proposed regularization method intrinsically enhances the training of Vision Mamba models while remaining compatible with diverse training paradigms. Notably, SLWS even elevates the accuracy of the previously collapse-prone Vim-L model to a 84.5% on ImageNet, demonstrating its effectiveness.

Q4 Table 4: Clarify the units, indicate if less/more is better

A4: Thanks for your suggestion and we have added such information to the revised version.

---

Thank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed your remaining concerns.

[1] Deep networks with stochastic depth, ECCV'16.