Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training
Authors: Zizheng Huang
Haoxing Chen
Jiaqi Li
jun lan
Huijia Zhu
Weiqiang Wang
Limin Wang
Venue: ICML 2025 poster (2025)
Keywords: Vision Mamba
Image Modeling
Training Regularization
Pre-training
Primary area: deep_learning->other_representation_learning
TL;DR: A plug-and-play training regularization for improving non-hierarchical Vision Mamba training and will be deactivated in inference.

Abstract: Recent Vision Mamba (Vim) models exhibit nearly linear complexity in sequence length, making them highly attractive for processing visual data. However, the training methodologies and their potential are still not sufficiently explored. In this paper, we investigate strategies for Vim and propose Stochastic Layer-Wise Shuffle (SLWS), a novel regularization method that can effectively improve the Vim training. Without architectural modifications, this approach enables the non-hierarchical Vim to get leading performance on ImageNet-1K compared with the similar type counterparts. Our method operates through four simple steps per layer: probability allocation to assign layer-dependent shuffle rates, operation sampling via Bernoulli trials, sequence shuffling of input tokens, and order restoration of outputs. SLWS distinguishes itself through three principles: \textit{(1) Plug-and-play:} No architectural modifications are needed, and it is deactivated during inference. \textit{(2) Simple but effective:} The four-step process introduces only random permutations and negligible overhead. \textit{(3) Intuitive design:} Shuffling probabilities grow linearly with layer depth, aligning with the hierarchical semantic abstraction in vision models. Our work underscores the importance of tailored training strategies for Vim models and provides a helpful way to explore their scalability. Code and models are available at https://github.com/huangzizheng01/ShuffleMamba

Stochastic Layer-Wist Wise Shuffle for Improving Vision Mamba Training
Zizheng Huang" 12
Haoxing Chen 3
Jiaqi Li
Jun Lan
Hujjia Zhu
Weiqiang Wang" Limin Wang"
Abstract
visual data, yet their quadratic complexity w.r.t. sequence length (Katharopoulos et al., 2020) remains significant drawback. In contrast, recent advances in State Space Mod- els (SSMs) (Kalman, 1960; Gu et al., 2021a;b; Smith et al., 2023) offer potentially more efficient sequence-based vision encoders (Zhu et al., 2024; Smith et al., 2023; Liang et al., 2024; Zhang al., 2024b; Li et al., 2024b). Among these, Mamba (Gu & Dao, 2023; Dao & Gu, 2024) stands out for its hardware-friendee design and selective scan computation, enabling near-linear complexity for longer sequences and prompting adoption in various vision tasks (Zhu et al., 2024; Liu et al., 2024e; Wang et al., 2024; Yang et al., 2024a). Extensions that incorporate 2-D scanning paths and visual priors (Zhu et al., 2024; Li et al., 2024a; Huang et al., 2024; Zhang et al., 2024a; Li et al., 2025; Tang al., 2024) have demonstrated competitive or compared to ViTs (Liang et al., 2024; Wu et al., 2024; Yue & Li, 2024). Such improvements, observed across su- pervised pre-training and diverse downstream applications (Chen et al., 2024; Patro & Agneeswaran, 2024; Phung et al., 2024), highlight Mamba's 's potential as an efficient, scalable foundation for visual processing (Yang et al., 2024b; Liu et al., 2024a; Xiao al., 2024).
Recent Vision Mamba (Vim) models exhibit nearly linear complexity in sequence length, mak- ing them highly attractive for processing visual data. However, the training methodologies and their potential are still not sufficiently explored. In this paper, we investigate strategies for Vim and propose Stochastic Layer-Wiss Shuffle (SLWS), a novel regularization method that can effec- tively improve the Vim training. Without archi- tectural modifications, this approach enables the non-hierarchical Vim to get leading performance on ImageNet-19 K compared with the similar type counterparts. Our method operates through four simple steps per layer: probability allocation to assign layer-dependent shuffle rates, operation sampling via Bernoulli trials, sequence shuffling of input tokens, and order restoration of outputs. SLWS distinguishes itself through three princi- ples: (I) Hlug-and-play No architectural modi- fications are needed, and it is deactivated during inference. (2) Simple but effective: The four-step process introduces only random permutations and negligible overhead. (3) Intuitive design: Shuf- fling probabilities grow linearly with layer depth, aligning with the hierarchical semantic abstrac- tion in vision models. Our work underscores the importance of tailored training strategies for Vim models and provides a helpful way to explore their scalability. Code and models are available at the open source URL.
superior performance
Initial efforts to scale up Vision Mamba (Vim) models were hindered by overfitting issues (Zhu et al., 2024; Ren et al., 2024; Li et al., 2025), causing performance degradation and even model collapse. In addition, the non-hierre -on-hierarc Vim architecture further complicates the pursuit of higher accu- racy (Li et al., 2024a; Tang et al., 2024). Although a limited number of supervised and unsupervised strategies (Wang et al., 2024; Liu & Yi, 2024) have successfully trained and scaled certain Mamba- Mamba-based models to huge sizes (Ren et al., 2024), recent research has moved beyond mere model enlargement toward broader, more robust improvements. Nevertheless, more efficient training methodologies are still urgently needed to overcome challenges like overfitting and narrow the performance gap with leading architectures such as VIT on ImageNet-Ik Ik (He et al., 2022; Wei et al., 2022a; Hou et al., 2022; Peng et al., 2022), where MambaMLP- L (Ren et al., 2024) (84.5%) still trails MAE-L (He et al., 2022) (85.9%).
1. Introduction
Vision Transformers (ViT') (Dosovitskiy et al., 2021; Liu et al., 2021; Dong et al., 2022; He et al., 2022; Bao et al., 2022) have achieved remarkable performance in modeling
Sttaee Key Lab of Novel Software Technology, Nanjing Uni- versity "Shanghai Innovation Institute 'Independent Researcher " "China Mobile Research Institute 'Shanghai Al Lab. Correspon- dence to: Limin Wang <lmwang nju.edu.en>>.
In this paper, we focus on training methods for Vim models and propose a Stochastic Layer-Wiss Shuffle regularization algorithm that effectively mitigates overfitting and boosts
Proceedings of the nnnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training
performance in large-scale Vim architectures. Concretely, the algorithm unfolds in four steps each layer's forward pass: (1) probability allocation to assign layer-dependent shuffle rates, (2) operation sampling via a Bernoulli trial, (3) shuffling the input token sequence, and (4) restoring the output sequence to its original order. The underlying ratio- nale is that deeper layers, which need higher-level semantic representations, can tolerate greater perturbations in token positions, whereas shallower layers should remain sensi- tive to low-level information. Restoring the sequence order prevents recursive effects for later layers. We verify the ef- fectiveness of this method in both supervised classification settings and a pre-training plus fine-tuning paradigm. The main contributions of this paper are summarized as follows:
ily focus on exploiting spatial locality, making assumptions about feature locality, translation, and scale invariance.
The introduction of ViT (Dosovitskiy et al., 2021) marks a turning point. Adapted from the NLP community (Vaswani et al., 2017), ViTs treat images as sequences of flattened 2D patches to capture global relationships (Liu et al., 2022a; Wang et al., 2021). As ViTs evolved, models like DeiT addressed optimization challenges (Touvron et al., 2021; He et al., 2022), while others introduced hierarchical structures and convolution operations to incorporate inductive biases of visual perception (Liu et al., 2021; Wang et al., 2021; 2022). These modifications allow for better performance across diverse visual tasks, although at the cost of added complexity in the models. Recently, there has been a trend of reverting to the original, plain ViT architecture due to its simplicity and flexibility in pre-training and fine-tuning across tasks (Bao et al., 2022; Xia et al., 2022; Carion et al., 2020; Cheng et al., 2022). However, one of the major challenges is the quadratic complexity of the self-attention mechanism (Katharopoulos et al., 2020; Zhu et al., 2023) limits the number of visual tokens that can be processed thereby impacting efficiency.
(I) We present a Stochastic Layer-Wiss Shuffle regular- ization algorithm for non-hiene non-hierarchical Vision Mamba models. This plug-and-play method effectively mit- igates overfitting, introduces minimal overhead, and requires no changes the underlying architecture.
(2) In a supervised setting, we show that the algorithm successfully addresses overfitting in large-scale mod- els, boosting performance in visual classification and downstream dense prediction tasks ADE20K seg- mentation and COCO detection).
State Space Vision Models Early state space transforma- tions (Gu et al., 202 Smith al., 2023; Gu al., 2023), inspired by continuous state models and bolstered by HiPPO initialization (Gu et al., 2020), showcased the potential for handling extensive dependency problems (Nguyen et al., 2023; Tallec & Ollivier, 2018). To overcome computational and memory issues, S4 (Gu et al., 2021a) enforced diagonal structure on the state matrix, while S5 (Smith et al., 2023) introduced parallel scanning to enhance efficiency further. The Mamba model (Gu & Dao, 2023; Dao & Gu, 2024) stands out for its novel approach to SSMs. By parameter- izing the state space matrices as projections of input data, Mamba proposes the more flexible selective scanning.
(3)
We further integrate masked feature distillation into the Vim pretraining process, demonstrating Vision Mamba can also beneficial from a semantic-rich frozen tok- enizer. Notably, incorporating SLWS achieves 87.6% accuracy on ImageNet-IK, K, establishing a new state-of- the-art for Vision Mamba models on this benchmark.
2. Related Work
Vision Backbones In the field of computer vision, the exploration of efficient and scalable backbone architec- tures has led to significant advancements (He et al., 2016; Krizhevsky et al., 2017: Dosovitskiy et al., 2021; Zhu et al., 2024), primarily driven by CNNs (Simonyan & Zisserman, 2015; Li et al., 2019; Liu al., 2022b) and ViTs (Dosovit- skiy et al.. 2021; Liu et al., 2021; Wang et al., 2021) recently. Initially, CNNs serve as the foundation and have evolved into deeper architectures, such as AlexNet (Krizhevsky et al., 2017), VGG (Simonyan& & Zisserman, 2015), and ResNet (He et al., 2016). Various studies have introduced advanced operators, architectures, and attention mechanisms to improve the effectiveness of models such as SENet (Hu et al., 2018 8) and SKNet (Li et al., 2019). The continuous refinement of convolutional layers has resulted in architec- tures like RepLKNet (Ding et al., 2022) and ConvNeXt (Liu et al., 2022b), which offer improved scalability and accuracy. Despite significant advancements, CNNs primar-
While ViTs and CNNs have laid a robust foundation for various visual tasks, Mamba offers a unique potential due to the ability to scale linearly with sequence length (Patro & Agneeswaran, 2024; Zhu al., 2024; Nguyen al., 2022; Lieber et al., 2024). S4ND (Nguyen et al., 2022) is the pioneering effort to integrate SSM into visual applications. However, the straightforward expansion did not efficiently capture image information. This gap led to further innova- tions in hybrid CNN-SSM hierarchical architecture, such as U-Mamba (Liu et al.. 2024b), VMnmba (Liu et al., 2024c) and MambaMixer(( (Behrouz et al., 2024). Recent efforts have sought to build generic vision backbones purely based on SSMs without relying on attention mechanisms (Zhu et al., 2024; Li et al., 2024a; Wang et al., 2024; Ren et al., 2024). Vision Mamba model, built by sequentially stacking Mamba blocks, has been shown to outperform ViT or perform on par in small model sizes (Wang et al., 2024; Liu & Yi, 2024;
2

---

Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training
Yang et al., 2024a). There are also some work exploring to refine the scanning method in Vim for visual data (Yang et al., 2024a; Li et al., 2024a; Huang et al., 2024; Chen et al., 2024; Tang et al., 2024; Pei et al., 2024). Nevertheless, Vims are stuck into issues like overfitring and have a noticeable performance gap compared to VIT in large sizes.
and Small Vim models. Mamba-Reg (Wang et al., 2024) introduces "registers" (a group of extra [CLS] tokens) to mitigate high-norm outliers, enabling Mamba-Reg to out- perform ViTs under supervised classification. Meanwhile, ARM (Ren et al., 2024) and MAP (Liu & Yi, 2024) adopt autoregressive pipelines to further scale up Vim models. Despite these advances, a noticeable performance gap re- mains between Vim and ViT. highlighting the urgent need for continued exploration of Vim's 's capabilities.
Training Methodologies To improve the training and generalization of deep models, various regularization tech- niques have been developed over the past years. Normal- izations (Ioffe & Szegedy, 2015; Ulyanov et al., 2016; Wu & He, 2018) are proven be effective for speeding up the convergence, in which the Layer Normalization (Ba et al., 2016) and RMSNorm (Zhang & Sennrich, 2019) are popular in training of large models. The family of data augmentations (Cubuk et al., 2020; Hoffer et al., 2020; Yun et al., 2019; Zhang et al., 2018) helps to produce more ro- bust representations and enhance performance. Stochastic depth and drop path (Huang et al., 2016; Larsson et al., 2016) drop the connection in the block level, which can not only overcome overfitting but also decrease the training cost. Weight decay (Krogh & Hertz, 1991; Loshchilov & Hutter, 2019) is commonly adopted for mitigating overfitting as well in a weight-penalizing manner. Besides, the earlier Dropout approach (Srivastava et al., 2014) introduces dis- turbance by dropping hidden units. They have played roles in various network training scenarios.
3. Methodology
In this section, we propose Stochastic Layer- Wise Shuf- fle Regularization (SLWS) for supervised training of non- hierarchical Vim models, along with a brief introduction of masked distillation strategy employed for pre-training. We first present the preliminaries in the following subsections to establish foundational concepts.
3.1. Preliminaries
State Space Model (SSM) (Gu al., 2021a;b) is origi- nally designed for modeling continuous-tine systems by projecting I-D input stimulation xxt) to the output signal v(t) via hidden state h(t) E RR. Formally, SSM is ex- pressed with the subsequent ordinary differential equation (ODE) as follows:
(() = Ah(t) + Br((), y(t) = Ch(t) + Dz((t,
When it comes to vision models, numerous training strategies have been proposed beyond supervised classification. Early self-supervised methods relied on surrogate tasks such as jigsaw puzzles (Noroozi & Favaro, 2016) predicting spatial context (Doersch et al., 2015), while subsequent contrastive approaches like SimCLR (Chen et al., 2020a), MoCo (He et al., 2020; Chen et al., 2020c; 2021)), and iBoT(Zhou et al., 2022) effectively trained both CNNs and ViTs by leveraging instance discrimination. More recently, masked pre-training techniques begin from MAE (He et al., 2022; Tong et al., 2022) and BEiT (Bao et al., 2022) have shown remarkable potential for scaling ViT models. These kinds of methods reconstruct raw pixels or discrete tokens to learn semantic-rich embeddings (Xie et al., 2022; Chen et al., 2020b). Additionally, with a Self-EMA or frozen tokenizer, masked feature distillation methods (Peng et al.. 2022: Hou et al., 2022; Fang et al., 2023; Baevski et al., 2022) can fur- ther elevate their generalization and performance of ViTs. In this strategy, the student model processes remaining patches after masking and is trained with the teacher target, which showcases superior efficiency and performance (Fang et al., 2023; Li et al., 2023; Peng et al., 2023).
(1)
where A E RR denotes the system's evolutionary ma- trix, with B E R
C E Rix" and D are projection parameters. a discrete system scemario, the above SSM is discretized by a timescale parameter A, transforming the expressions of A and B into their discrete equivalents Á and B. In Mamba models, such conversion is implemented with the Fedo-Order Hold (ZOH) rule, which is expressed as follows:
A = exp( AA),
B AA (exp( AA II)- AB.
Then, sequential input { is mapped via this dis- cretized system to its output {w:} J as:
h44 Ah; + Bag
y/: = Ch²y + Dr,,
Mamba (Gu & Dao, 2023) designs the B, C, and A to be input-dependent to improve the intrinsic capacity for con- textual sensitivity and adaptive weight modulation. Besides, a Selective Scan Mechanism is ensembled in for efficient computation. To this end, for a Vim (Zhu et al., 2024) block (or layer) se, it includes an SSM branch, whose output
For non-hierarchical Vim models, several training methods extend beyond scanning-basee approaches. Vim-F (Zhang et al., 2024c) explores frequency-doman training to enhance the global receptive field, showing improvements for Tiny
3

---

Stochastic Layer-Wiss Shuffle for Improving Vision Mamba Training
Pen
0.1
0.2
0.3
0.4
0.5
Layer Index
Sampled bg
2
5
0
1
1
0
1
5
Block0
Block1
Block2
Block3
Block4
Blocks
2
2
3
1
4
m
SSS
Figure 1: Stochastic Layer-Wiss Wise Shuffle Regularization (SLWS). Deeper layers are assigned larger probabilities for shuffle regularization to enhance positional transformation invariance. The variable by is sampled based on these probabilities to determine whether to execute regularization. SLWS only involves sequence permutation and restoration, and is not applied during inference. The snake icon indicates where regularization performed.
multiplied by the result of another gated branch to produce the final output sequence X, E TDD' Thus, the corre- sponding forward process of non-hierarchical Vim (without downsampling) is expressed in the following form:
sequences non-causal. Several previous studies have integrated different scanning strategies into Mamba layers to better capture spatial context (Zhu et al., 2024: Liu et al., 2024c; Yang et al., 2024a; et al., 2024a; Tang et al., 2024). Nevertheless, these methods remain reliant on simple 1-D comnerto-ccrner scanning and often suffer from overfitting. To address these limitations, propose Stochastic Layer- Wise Shuffle (SLWS), regularization technique guided by the following insights:
X (Xi-1).
(4)
Masked Feature Distillation (MFD) techniques enhance pre-training by masking significant portion of image patches and subsequently reconstructing the targets using the unmasked regions input. Methods such as MAE (He et al., 2022) have been proven effective in training foun- dational Vision Transformers (ViTs) without relying on labeled data. Further research has shown that employing feature-level targets can lead additional improvements, including the use of HOG features (Wei al., 2022a), Self- EMA (Baevski al., 2022), CLIP embeddings (Radford et al., 2021; Wei et al., 2022b; Hou et al., 2022), and dis- crete tokens (Peng et al., 2022). The MFD process can be formulated as follows:
(1) Fixed corner-to-on to-cree sequential or regional scanning Vim does not naturally align with the need capture both local and global spatial correlations.
(2) Deeper layers of a vision encoder should learn higher- level semantic representations, while shallower layers primarily encode low-level information.
(3) Achieving stronger semantic perception in deeper lay- requires transformation invariance for patch po- sitions, whereas shallower layers must preserve posi- tional sensitivity.
min E dist
(5)
where T represents the teacher tokenizer, f denotes the student model, and X" refers to the remained visible parts. dist [, the selected distance function.
(4) Introducing stochastic perturbations into the sequen- tial structure can increase task complexity, potentially mitigating overfitting, but also contributes simulate diverse spatial contexts.
3.2. Stochastic Layer- Wise Shuffle
(5) Besides designing layer-dependent for differing seman- tic requirements, a sequence restoration step ensures that subsequent layers receive inputs in the original order, thus avoiding unnecessary disruptions.
As formulated above, the SSM-based Mamba was origi- nally proposed for sequence modeling but does not natu- rally adapt to woo-dimensional image data, where patch

---

Mamba Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training
Random Shuffle Forward and Restoration Inspired
Algorithm Layer-Wise Shuffle forward
by stochastic depth (Huang al., 2016), we introduce Bernoulli random variable by E {0, J to determine whether the f-th layer will apply shuffle-based regularization. If by = 1, the input token sequence X will be randomly shuffled into _ thereby encouraging positional trans-
Require: token sequence X Rxxxx
layer ss, probability pr. training flag F
Ensure: token sequence X;
I: this layer is trained with regularization
2: if and rand () < pe then
3: shuffle_indices = -randperm()eexpanddee exxpandd 1, D)
formation invariance. Otherwise, X,-1 remains unchanged.
4: restore indices argsort(shuffle sinufffeiidiies...... dim=1)
We denote this operation by T-( be). and its inverse
5: X, gather((X,-1. 1, shuffle indices)
(. by) restores the shuffled output X₂ to the original
6: X; se(X
order to avoid recursive effects on later layers:
7: X; gather restore.indices)
8: else
X - TT (x be))).
(6)
inference or trained without regularization 10: x
9:
11: end if
Layer-Wise Probability Assignment Additionally, each Vim layer is assigned a distinct probability of applying SLWS, reflecting the intuition that deeper layers should ex- hibit greater transformation invariance. In this work, we use a linear scheduling function starting with l = O. Specifi- cally, the probability Pe of applying shuffle regularization at the f-th layer is:
12: Return: X
modeling of 2D visual data by encouraging position invari- ance in deeper layers. (3) By increasing task complexity, it helps mitigate overfitting without incurring heavy computa- tional overhead in training.
P(be = 1) Pt.
3.3. Masked Pre-training for Vim
The fundamental idea of visual masked modeling is to re- construct the complete target by leveraging relationships between unmasked image patches, thereby capturing com- plex semantic dependencies. We establish a simple masked pre-training pipeline for the Vim encoder, as formulated in Eq. (5) and illustrated in Fig. 2. Alongside the Vim student encoder, we employ frozen CLIP vision encoders (Radford et al., 2021) as the teacher tokenizers T, which provide feature targets. Inspired by MAE (He et al., 2022), our ap- proach adopts a auto-encoder design, featuring a lightweight self-attention decoder d that reconstructs the Vim features f(X") to match the teacher outputs. To enhance training stability, apply normalization layers to encoded features, decoder outputs, and teacher targets. We further employ the smooth-f, loss for the distance metric dist [.
where Pr. is a hyperparameter. Since we shuffle tokens according to a discrete uniform distribution, the probability that the i-th token moves the j-th position is:
1 P p(ak= - x;) L+1
P =
(I+I)) PL.
Notably, hierarchical Mamba architectures with spatial downsampling operations are incompatible with SLWS, as token sequence length reduction prevents output order restoration. Additionally, SLWS fundamentally differs from random scanning methods. Because our layer-dependent probability allocation imposes progressive regularization intensity that aligns with hierarchical semantic abstraction.
4. Experiments
3.2. EFFICIENCY ANALYSIS
We conducted extensive experiments evaluate Vim train- ing, exploring non-hierarchical models trained via super- vised classification and pre- vre-trannin paradigms, assessing their downstream task performance, and performing detailed algorithm analysis through ablation studies. We conduct both horizontal and vertical comparisons to analyze our model and approach.
Fig. and Algorithm 1 illustrate SLWS for Vim training with PyTorch pseudo-code. Random index generation in- curs OL)) complexity, while sorting for restoration adds Oo L). Because we apply the same random index to the entire batch, the batch size does not inflate these costs. Consequently, SLWS introduces only (L log L) additional maximal overhead, and our ablation results in Section 4.3 confirm the minimal impact on overall training efficiency.
4.1. Implementation Settings
Overall, SLWS offers several key advantages: (1) It is easy to implement and does not alter the model architecture, adding no extra cost at inference time. (2) It fosters stronger
We evaluate various sizes of non-hierarchical Vision models and details of settings are listed in Appendix A. Con- figurations of non-hierarchical models with different sizes
5

---

Stochastic Layer-Wiss Shuffle for Improving Vision Mamba Training
Table 2:
Vim training comparisons, where "S." indicates SLWS, "sup." indicates supervised classification, "reg." refers to token registers (Wang et al., 2024), and "cont." denotes contrastive training. All models are evaluated on the ImageNet- benchmark.
Teacher model
Loss
Model
Training tech.
#Params
Epoch
Acc.
supervised
Concal
Vim-M
sup.
74M
300
80,9
Vim model
Decoder
Vim-B
sup.
98M
300
79.8
Vim-B stride]
sup.
98M
300
81.2
Vim-L
sup.
284M
300
collapsed
Vim-M
sup.. S.
74M
300
82.8 (+1.9)
Vim-B
sup.. S.
98M
300
82.7 (+2.9)
Vim-L
sup.. S.
284M
200
82.9
Masking tokens
Vim-L (384 res.
sup.. S.
284M
220
84.5
MambuR-B
sup., reg.
99M
220
83.0
Figure 2: Masked feature distillation pipeline. A frozen semantic-rich teacher as tokenizer produces target for the student branch, which is in auto-encoder style.
MumbaMLP-L
sup.
297M
300
81.4
MambaR-B
sup.. reg. S.
99M
220
83.1 (+0.I)
MambuMLP-L
sup., S.
297M
300
82.9 (+1.5
pre-training
MambaMLP-B
cont.
85M
300
81.4
involved in experiments are listed in the Table 1, in which MambaR (Wang et al., 2024) models add a group of ex- tra tokens based on Vim (Zhu et al., 2024). We use the A optimizer (Loshchilov & Hutter, 2019) with a co- sine learning rate schedule and employ BFloat16 precision enhance training stability. Additionally, report results using Exponential Moving Average.
ManbaMLP-B
MAE
85M
300
81.6
MambaMLP-B
ARM
85M
300
82.5
MumbaMLP-B
ARM
85M
1600
83.2
MunbaMLP-L
ARM
297M
1600
84.5
MambaMLP-H
ARM
662M
800
85.0
MambaMLP-B
MAE, S.
85M
300
82.0 (+0.4)
MumbaMLP-B
MFD. S.
85M
300
84.3 (+1.I)
MumbaMLP-L
MFD
297M
300
86.4 (+1.9)
MunbaMLP-L
MFD. S.
297M
300
86.7 (+2.2)
MambaMLP-II
MFD, S.
662M
300
87.6 07.666)
Table I: Model configurations.
normalization the output features improve training bility. During pre-training,
Model
Block Config.
Width
Depth
#Param. (M)
image sizes of 192 and 224 for the MAE and MFD pipelines, respectively. The shuffle rate is set to 0.4 for large and huge models. For the masking strategy MFD, we follow existing studies (Peng et al., 2023; Hou et al., 2022) by setting the masking ratio to 0.5 and 0.6 with utilizing attentive masking.
VIT-B
Attention+MLP
768
12
Vim-B
Mamba
768
24
98
Vim-M
Mamba
576
32
74
MambaR-B
Mamba
768
24
99
MambaMLP-B
Mamba+MLP
768
12
85
VIT-L
Attention+MI.P
1024
24
309
Vim-L
Mamba
1024
284
ManbaR-L
Mamba
1024
48
341
ManbaMLP-L
MaP
1024
4.2. Main Results
VIT-H
Attention+MLP
1280
32
Vim Training Comparison To evaluate the SLWS regu- larization and MFD training pipeline, we compare it with
MambaMLP-H
Mambu+MLP
1536
24
662
state-of-the-art (SOTA) training methods in both supervised and pre- pre-training settings. Table presents the results, includ- ing Vim (Zhu et al., 2024), MambaR (Wang et al., 2024), and MambaMLP (Ren et al., 2024). Notably, ARM and MAE pre-train models with an input resolution of 192 x 192 and subsequently fine-tune with 224 x 224. We utilize a CLIP-B for MambaMLP-B with a CLIP-L for MambaMI.P-L and MambaMLP-H teacher tokenizers, respectively. Based these results, we draw the following observations:
For supervised training, train from scratch ImageNet- IK (Deng al., 2009), which contains 1.28 million samples for the classification task. Middle and base-size models are trained for 300 epochs with a batch size of 2048, while large models are trained for 200 epochs with a batch size of 1024. The shuffle rate Pr. is set to 0.5 for middle and base-size models and 0.6 for large models. Following the VideoMamba (Li et al., 2024a) setup, [CLS] token is prepended to the token sequences for classification.
For MFD pre-training we use frozen CLIP vision encoders
(1) SLWS significantly improves supervised Vim train- ing across model scales. For the middle-sized Vim-M, SLWS boosts accuracy by 1.9%, and for the base-sized Vim-B, the gain reaches 2.9% (from 79,8% to 82.7%).
as tokenizers. Inspired by MAE (He et al., 2022), our de- coder a lightweight self-attention transformer with four blocks and a hidden dimension of 512. We apply layer
6

---

Stochastic Layer- Wise Shuffle for Improving Vision Mamba Training
Table 3: ImageNet- - classification comparison among
3.5
Eveluation
Evaluation loss w/o SUWS
5.5
different backbone and training methods.
Training loss w/o SUVS
Training loss w/ SUWS
Model
Training
iPaaams
FLOPs
Acc.( %)
Hierarchical
4.5
RegNetY-4G
sup.
21M
4G
80.0
Training loss
RegNetY-8G
sup.
39M
8G
81.7
Evaluation
4.0
RegNetY- 16G
sup.
84M
16G
82.9
ConvNeXi-T
sup.
29M
4.5G
82.1
ConvNeXt-S
sup.
50M
8.7G
83.1
3.5
ConvNcXt-B
sup.
89M
15.4G
83.8
Swin-T
sup.
28M
4.6G
81.3
Swin-S
sup.
50M
8.7G
83,0
3.0
Swin-B
sup.
88M
15.4G
83,5
Swin-B
SimMIM
88M
15.4G
84.0
1.0
Swin-L
SimMIM
197M
35.8G
85.4
VMamba-T
sup.
31M
4.9G
82.5
50
100
150
200
250
300
VMamba-S
sup.
50M
8.7G
83.6
Epoch
VMamba-B
sup.
89M
15.4G
83.9
Non-Hierarchical
Figure 3: Training and evaluation loss curves for 300 epochs middle-size Vims.
ConvNeXt-S
sup.
22M
4.3G
79.7
ConvNcXt-B
sup.
87M
16.9G
82.0
DeIT-S
sup.
22M
4.6G
79.8
DeTT-B
Distill.
87M
17.6G
81.9
VIT-B JMAE
sup.
87M
17.60
82.3
Notably, SLWS enables stable training of the previ- ously collapsing Vim-L (284M parameters), achieving 82.9% accuracy and 84.5% with 384 x384 fine-tuning.
VIT-L. sup.]
sup.
309M
61.60I
82.6
ViT-B
MAE
87M
17.6G
83.6
ViT-L
MAE
309M
61.6G
85,9
VIT-H (14 stride]
MAE
632M
167G
86.9
VIT-B
MaskDistill
87M
17.60
85.0
(2) MFD pre-training substantially enhances Vim capa- bilities. When combined with SLWS, MambaMLP-B achieves 84.3% accuracy (+1. % over the ARM base- line), while MambaMLP-L reaches 86.7%, surpass- ing ARM's 1600-epoch result (84.5%) within just 300 epochs. This demonstrates clear training efficiency advantage over previous methods and highlights the importance of leveraging semantic-rich tokenizer.
VIT-L
MaskDistill
309M
61.6G
87.6
ViT-B
BETTv2
87M
17.6G
85.0
ViT-L
BEITv2
309M
61.6G
87.3
Vim-S
sup.
26M
4.3G
80,5
VideoMamba-S
sup.
26M
4.3G
81.2
VideoMamba-M
sup.
74M
12.7G
80.9
VidcoMamba-M
sclf-Distill.
74M
12.7G
82.8
LocalViM-S
sup.
28M
4.8G
81.2
Plain-amba--L2
sup.
25M
8.IG
81.6
PluinMamba-L3
sup.
50M
14.4G
82.3
MambaR-S
sup..
28M
4.5G
81.1
(3) SLWS provides complementary benefits across train- ing paradigms. In MAE pre-training. SLWS improves MambaMI.PRB by 0.4% (8.66% to 82.0%). For MFD, we observe a 0.3% improvement over MambaMLP-L. (86.4% to 86.7%) and SLWS enables MambaMLP-H to achieve 87.6%, i.e., the new state-of-the-art result for Vision Mamba on ImageNet-IK.
MambaR-B B
sup., reg.
99M
17.8G
83.0
MambaR-1. -
sup.,
341M
64.2G
83.6
MambaR I. |384 res.
sup., reg.
341M
179G
84.5
MambaMLP-B
ARM
85M
15.5G
83.2
MambaMLP-L
ARM
297M
54.7G
84.5
MambaMLP-H
ARM
662M
123G
85.0
MambaMLP.B
MFD. S.
85M
15.5G
84.3
MambaMI
MFD, S.
297M
54.7G
86.7
MambaMI.P4
MFD,
662M
123G
87.6
Consequently, SLWS not only prevents collapse in super- vised learning of large models through stochastic regular- ization but also enhances pros--prrdiigm generalization without any architectural changes. It is also worth noting that combining MFD with SLWS is particularly effective for non-hi an-hierarchical Vim training. Beyond the above accuracy evidence for mitigating overfitting, we plot the training and evaluation curves in Fig. for further demonstration. We observe that the model trained with SLWS stabilizes at higher training loss yet achieves a lower evaluation loss. By contrast, the ablated model tends to overfit, showing a lower training loss but a higher error rate on evaluation. This confirms that SLWS effectively adds perturbation to
sequential perception, raising task complexity and reducing the overfitting risk for Vim.
Comparison to Various Backbones. Table 3 reports ImageNet-IK classification results across range of back- bones. We include CNN-based methods (RegNetY (Ra- dosavovic et al., 2020), ConvNeXt (Liu et al., 2022b)), hi- erarchical Transformers (Swin (Liu et al., 2021) trained with SimMIM (Xie et al., 2022)), and ViT variants trained with DeiT (Touvron et al., 2021), MAE, MaskDistill (Peng et al., 2023), or BEITv2 (Peng et al., 2022). We also list

---

Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training
SSM-based approaches (VMamba (Liu et al., 2024c), Vim, VideoMamba (Li al., 2024a), LocalViM (Huang al., 2024), PlainMamba (Yang et al., 2024a), MambaR (Wang et al., 2024), ARM). Under purely supervised training, sev- eral SSM-based models match or exceed the performance of their CNN and hierarchical Transformer counterparts at comparable model sizes. When pre-training is introduced, masked modeling generally boosts performance across ar- chitectures. However, there remains noticeable gap be- tween the previous best SSM-based results and advanced ViT models trained via masked image modeling (e.g., ARM 85. 55,9% vs MAE 86.9%). Moreover, the introduction of our
Table 4: Semantic segmentation results on ADE20K Val. Computation FLOPs are measured under 512x2048 imput resolution. "MS" means multi-scale test.
Model
Parram.
FLOPs
mloU
+MS
ResNet-50
67M
953G
42.1
42.8
ResNet-100
85M
1030G
42.9
44.0
ConvNeXt-B
122M
1170G
49.1
49.9
DeiT-B+MLN
144M
2007G
45.5
47.2
VIT-B
127M
46.1
47.1
ViT-Adapter-B
134M
632G
48.8
49.7
VIT-L VIT-L/MAE]
127M
53.6
Swin-B
121M
1170G
48.1
49.7
ViM-S
46M
44.9
MambaMLP P w/ ARM
ViM-B
131M
477G
45.2
ImageNeet-T Accuracy (96
MambaMLP w/ MAP
MambaR-B
132M
47.7
VIT w/ MAE
km
MambaR-L.
377M
49.1
MambaMLP P w/ outs
Vim-M [S.]
106M
384G
47.2
48.2
26
Vim-B [S.]
131M
477G
47.0
48.3
MambaR-B [S.]
131M
477G
48.2
48.9
MambaR-AdappeerBB [S.]
145M
1428G
49.3
50.1
04
64.2%
MambaMLPPL. /MFD /MPD, S.J
324M
1270G
53.8
Adapter-B outperforms ViT-Adapter-B by 0.5%. Addition- ally, our MFD+SLWS framework enables MambaMLP-L to match MAE ViT-L's performance.
Base
Large
Huge
Figure 4: Pre-training comparison anchored by MAE on ImageNet-IK.
Table Object detection and instance segmentation results. FLOPs are calculated with size x280x800. Gray fonts indicate the models pre-trained ImageNet-21K.
SLWS + MFD pipeline narrows this gap considerably. By further unlocking the potential of non-hierarchical Vim-like models, enables them to outperform ViT models trained with MAE, demonstrated by Fig. 4. This substantial im- provement underscores the value of our approach in improv- ing the performance of Vim models.
AP'
AF
AP
Ap
AP
App
Com-NeXI-B
108M
186G
47.0
69.4
51.7
12.7
66.3
46.0
46.9
VIT-B WT-I
114M
42.9
65.7
46.8
39.4
62.6
42,0
337M
45.7
68.9
49.4
65.6
44.6
WT--Adapee-B B
120M
47.0
68.2
51.4
41.8
65.1
44,9
WT-Addpter-L
348M
48.7
53.2
43.3
67.0
46.9
PlainMamba--33
79M
696G
46.8
68
51.1
41.2
64.7
43.9
Dense Prediction Downstream Tasks To evaluate model capabilities, we conduct semantic segmentation experi- ments on ADE20K, detection and instance segmentation on COC02017 benchmark.
Vin [S.]
103M
564G
46.8
68.8
50.7
41.8
65.6
44,8
MambuR-BB B /S.]
131M
726G
47.7
69.7
51.8
42.6
66.7
45.8
La
383M
1734G
48.9
70.8
53.4
43.6
67.4
47.0
For downstream object detection and instance segmentation tasks, follow previous work to evaluate method. The Mask R-CNN (He et al., 2017) structure is adopted with schedule for 12-epoch fine-tuning. We utilize the com- monly adopted settings in previous work (Liu et al., 2021) and compare to different-type backbones. To compute the multi-scale features to fit the FPN network structure, we use the Adapter setup following (Yang et al., 2024a; Chen et al., 2023). The results are reported in Table 5. It be seen that our middle-size model is on par with the corresponding CNN and Transformer model, while the base-size model traine with registers and SLWS outperforms ViT-Adapter-B and ConvNext-B by 0.7 points AP¹, MambaR-L demon- strates higher AP"/AP" and even outperforms ImageNet- 21K pretrained ViT-Adapter-L. and VIT-L.
For segmentation experiment, adopt the UPerNet (Xiao al., 2018) head on ImageNet-IK 1K trained models. All the models are trained for 160K iterations with batch size 16. Following the common settings (Chen et al., 2023; Yang et al., 2024a; Wang et al., 2024), we use an Adam opti- mizer with 0.01 weight decay and polynomial learning rate schedule. The learning rates of the base and large-size models are as 6c-5 and 3e-5, respectively. The [CLS] and register tokens are discarded in the segmentation task. As shown in Table 4, our -rggulaaarriee MambaR-B surpasses both ViT-B and its non-SLWS counterpart, which consistently demonstrates the superiority brought by the proposed SLWS regularization. When integrating the multi- scale adapter configuration (Chen et al., 2023), MambaR-
8

---

Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training
4.3. Ablation Studies
Table 8: Ablation studies on [CLS] token shuffling with different size of models.
In this subsection, we perform ablation studies by varying the settings of the SLWS regularization investigate its effects and provide in in-depth analysis. We use middle- size vanilla Vim models as the default for experiments.
[CLS] token in Shuffling
Middle
Base
Large
X
82.6
82.6
82.8
82.7
82.6
82.9
Table 6: Ablation study on training throughput. Higher throughput (images/second) is better under the same setting.
Setting
128
224
384
512
768
supervised classification training except for MambaR con- figuration, we investigate whether including it in the shuf- fling process affects performance. The ablation results for different model sizes on ImageNet-IK are shown in Table 8. We observe that including the [CTS]] token in shuffling yields slightly better performance under the same settings for middle and large models. Consequently, for code sim- plicity, we shuffle the entire sequence by default, and the same approach applies when using registers.
w/ SLWS.
315.7
167.9
56.8
29.0
13.97
w/ SLWS.
311.4
164.8
55.7
28.6
13.72
Loss (%)
1.36
1.85
1.94
1.38
1.79
SLWS has a Negligible Impact on Training Through- put. SLWS operates on both input and output sequences of Mamba blocks, with efficiency analysis detailed in Sec- tion 3.2. To empirically evaluate its computational overhead, we conduct throughput measurements using standard image resolutions ranging from 28 to 768x768. Results in Table 6 demonstrate consistent throughput reduction below 2% across all resolutions. This negligible overhead con- firms SLWS as a simple but effective training regularization technique for vision mamba models.
5. Conclusion
We present Stochastic Layer-Wiss Shuffle, a specialized regularization method for non-hierarchical Vision Mamba training that addresses overfitting through layer-dependent sequence perturbations. By progressively increasing shuf- fle probabilities across layers, SLWS enhances positional transformation invariance in deeper semantic abstractions while preserving low-level spatial sensitivity. This approach achieves significant improvements for supervised training of Vision Mamba. When integrated with masked feature distillation, our Vim models establish new state-of-the-ae results on ImageNet- K and dense prediction tasks among the same type models. The method does not introduce archi- tecture modification and has negligible overhead, effectively unlocking the potential of Vision Mamba models.
Table 7: Ablation studies on probability settings. and "E" denote decreased linear, and exponential probability assignments, respectively.
Type
Py Acc.
Type
Pr Acc.
0.4 82.3
Layer-depend. (D) 0.5 81.2
0.5 82.7
Layer-depend. (E) 0.5 82.2
Layer-depend.
0.6 82.4
0.1 81.5
Constant
0.7 82.4
0.4 81.1
Layer-Wise Probability Assignment is Necessary. The layer-wise probability is a crucial component of the SLWS design, introducing semantic level prior across different layers. Table 7 presents the results under various probabil- ity assignment settings. Since shallower blocks are more sensitive to patch positions, the layer-dependent cases gen- erally outperform the constant settings. We also provide decreased linear probability assignment comparison, which takes larger shuffle probabilities for shallower layers. The result further demonstrates the correctness of the seman- tie level prior. Besides the linear setting, we experiment with a exponential one, i.c. a modification of Eq. (7) to P(b, = 1) = pp-- p( which has performance drop of 0.5 points compared to the vanilla linear one.
Acknowledgement
Thanks Di Yang from SII for his help. This work is sup- ported by the National Key R&D Program of China (No. 2022ZD0160900, Jiangsu Frontier Technology Research and Development Program (No. BF2024076), and the Col- laborative Innovation Center of Novel Software Technol- ogy and Industrialization. This work is funded by Nanjing University- -hhina Mobile Communications Group Lo.,Ltd. Joint Institute.
Impact Statement
This paper presents work whose goal is to advance the field of Deep Learning. There are many potential societal con- sequences of our work, none of which we feel must be specifically highlighted here.
Including the [CLS] Token in Shuffling Slightly Im- proves Performance. As the [CLS] token is used for

---

Stochastic Layer-Wise, Shuffle for Improving Vision Mamba Training
References
Deng, Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database.
Ba, J.J.L. L., Kiros, J.R and Hinton, G. E. Layer normalization. arXiv preprint
In CVPR, 2009.
2016.
Ding, X., Zhang, X., Han, J., and Ding, Scaling up your kernels 31x31: Revisiting large kernel design in cnns. In CVPR, pp. 1 11963-11975, 2022.
Baevski, A., Hsu, W.-N., Xu, Q.. Babu, A., Gu, J., and Auli, M. Data2vec: A general framework for self-supervised learning in speech, vision and language. ICML, pp. 1298-1312, 2022.
Doersch, C., Gupta, A., and Efros, A. Unsupervised visual representation learning by context prediction. In ICCV, pp. 1422-1430, 2015.
Bao, H., Dong, L., Piao, S., and Wei, F. Beit: Bert pre- training of image transformers. ICLR, 2022.
Dong, X., Bao, J., Chen, D., Zhang, W., Yu, N., Yuan, I.,, Chen, D., and Guo, B. Cswin transformer: general vision transformer backbone with cross-shaped windows. In CVPR, 2022.
Behrouz, A., Santacatterina, M., and Zabih, R. Mam- bamixer: Efficient selective state space models with dual token and channel selection. arXiv preprint (XXi::2403.1988: 2024.
Dosovitskiy, A., Beyer, L.., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICIR, 2021.
Carion, N., Massa, F., Synnaeve, G., Usunier, N., Kirillov, A., and Zagoruyko, S. End-to-end object detection with transformers. In ECCV, 2020.
Chen, G., Huang, Y., Xu, J., Pei, B., Chen, Z., Li, Z., Wang, J., Li, K., Lu, T., and Wang, L. Video mamba suite: State space model as versatile alternative for video understanding. arXiv preprint FrxXiv2403309622 2024.
Fang, Y., Wang, W., Xie, B., Sun, Q., Wu, L., Wang, X., Huang, T. Wang, X., and Cao, Y. Eva: Exploring the limits of masked visual representation learning at scale. In CVPR, pp. 99358-19369 2023.
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual rep- resentations. In ICML, 2020a.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
Chen, T., Kornblith, S., Swersky, K., Norouzi, M., and Hinton, G. Big self-supervised models are strong semi- supervised learners. In NeurlPS, 2020b.
Gu, A., Dao, T., Ermon, S., Rudra, A., and Ré, C. Hippo: Recurrent memory with optimal polynomial projections. In NeurlPS, 2020.
Chen, X., Fan, H., Girshick, R., and Hc, K. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.04297 2020c.
Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.0033 2021a.
Chen, X., Xie, S., and He, An empirical study of training self-supervised vision transformers. In ICCV, 2021.
Gu, A., Johnson, I., Goel, K., Saab, K.. Dao, T., Rudra, A., and Ré, C. Combining recurrent, convolutional, and continuous-time models with linear state space layers. NeurIPS, 2021b.
Chen, Z.. Duan, Y., Wang, W., He, J., Lu, T., Dai, J., and Qiao, Y. Vision transformer adapter for dense predictions. In ICLR, 2023.
Gu, A., Johnson, I., Timalsina, A., Rudra, A., and Ré, How train your hippo: State space models with gener- alized orthogonal basis projections. In ICLR, 2023.
Cheng, B., Misra, I., Schwing, A. G., Kirillov, A., and Gird- har, R. Masked-attention mask transformer for universal image segmentation. In CVPR, 2022.
He, K., Zhang, X., Ren, S., and Sun, Deep residual learning for image recognition. CVPR, pp. 770-778, 2016.
Cubuk, D., Zoph, B., Shlens, J., and Le, Q. Ran- daugment: Practical automated data augmentation with reduced search space. In CVPR Workshops, pp. 702-703, 2020.
He, K., Gkioxari, G., Dollár, P., and Girshick, R. Mask r-cnn. In ICCV, pp. 2961-2969, 2017.
Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. ICML,, 2024.
He, K. Fan, H., Wu, Y., Xie, S., and Girshick, R. Mo- mentum contrast for unsupervised visual representation learning. In CVPR, 2020.
10

---

Stochastic Layer- Shuffle for Improving Vision Mamba Training
He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, R. Masked autoencoders scalable vision learners. In CVPR, pp. 16000-16009, 2022.
Li, X., Zhong, B., Liang, Li, G., Mo, Z., and Song, S. Mambalet: Boosting tracking via long-term context state space model. arXiv preprint
2024b.
Hoffer, E., Ben-Nun, T., Hubara, I., Giladi, N., Hoefler, T., and Soudry, D. Augment your batch: Improving generalization through instance repetition. In CVPR, pp. 8129-8138, 2020.
Liang, D., Zhou, X., Wang, X., Zhu, X., Xu, W., Zou, Z.,
Ye, X., and Bai, Pointmamba: simple state space model for point cloud analysis. NIPS, 2024.
Lieber, O., Lenz. B., Bata, H., Cohen, G., Osin, J., Dalmedi- gos, I., Safahi, E., Meirom, S., Belinkov, Y. Shalev- Shwartz, S., et al. Jamba: hybrid transformer-mabbe language model. arXiv preprint arXiv:2403. 19887, 2024.
Hou, Z., Sun, F., Chen, Y.-K., Xie, Y.. and Kung, S.-Y. Milan: Masked image pretraining language assisted representation. arXiv preprint arXiv:2208 2004492008.0044,, 2022.
Hu, J., Shen, L.,, and Sun, G. auneerzeand-excitttion net- works. In CVPR, pp. 7132-7141, 2018.
Liu, J., Liu, M., Wang, Z., An, P., Li, X., Zhou, K., Yang. S., Zhang, R., Guo, Y., and Zhang, S. Robomamba: Efficient laoon-language-actte model for robotic reasoning and manipulation. In NIPS, 2024a.
Huang, G., Sun, Y., Liu, Z., Sedra, D., and Weinberger,
K. Q. Deep networks with stochastic depth. In ECCV, pp. 646-661, 2016.
Liu, J., Yang, H., Zhou, H.-Y., Xi, Y. Yu, L., Yu, Y., Liang, Y., Shi, G., Zhang, S., Zheng, H., et al. Swin- umamba: Mamba-based unet with imagenet-basee pre- training. arXiv preprint arXiv:2402. 033002 2024b.
Huang, T., Pei, X., You, S., Wang, F., Qian, C., and Xu, C. Localmamba: Visual state space model with windowed selective scan. arXiv preprint
2024.
loffe. S. and Szegedy,
C. Batch normalization: Accelerating deep network training by reducing internal covariate shift.
Liu, Y. and Yi, L. Map: Unleashing hybrid vision backbone's potential with masked autoregressive pretraining. arXiv preprint arXiv:2410.0087,, 2024.
In ICML, pp. 448-456, 2015.
Kalman, R. E. A new approach to linear filtering and pre- diction problems. 1960.
Liu, Y., Tian, Y., Zhao, Y., Yu, H., Xie, L.,, Wang, Y., Ye, Q., and Liu, Y. Vmamba: Visual state space model. arXiv preprint arXiv:2401.10166 2024c.
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, Transformers are rnns: Fast autoregressive transformers with linear attention. In ICML., 2020.
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. ICCV, pp. 10012- 10022, 2021.
Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. Commun. ACM, 60:84-99 90, 2017.
Liu, Z., Hu, H., Lin, Y., Yao, Z., Xie, Z., Wei, Y., Ning, J., Cao, Y., Zhang, 7.,, Dong, I.,, et al. Swin transformer v2: Scaling up capacity and resolution. In CVPR, 2022a.
Krogh, and Hertz, A simple weight decay can improve generalization. In NeurIPS, 1991.
Larsson, G., Maire, M., and Shakhnarovich, G. Fractalnet: Ultra-deep neural networks without residuals. In ICLR, 2016.
Liu, Z., Mao, H., Wu, C.-Y., Feichtenhofer, C., Darrell, T., and Xie, S. A convnet for the 2020s. In CVPR, pp. 11976- 2022b.
Li, K., Wang, Y., Li, Y., Wang, Y., He, Y., Wang, L., and Qiao, Y. Unmasked teacher: Towards training-efficies video foundation models. ICCV, 2023.
Loshchilov, and Hutter, Decoupled weight decay regu- larization. In ICLR, 2019.
Nguyen, E., Goel, K., Gu, A., Downs, G., Shah, P., Dao, T., Baccus, S., and Ré, C. S4nd: Modeling images and videos as multidimennional signals with state spaces. In NeurIPS, 2022.
Li, K., Li, X., Wang. Y., Hc, Y., Wang. Y., Wang, L., and Qiao, Y. Videomamba: State space model for efficient video understanding. In ECCV, 2024a.
Li, S., Singh, H., and Grover, A. Mamba-nd: Selective state space modeling for multi multi--iimensonal data. In ECCV, pp. 75-92, 2025.
Nguyen, E., Poli, M., Faizi, M., Thomas, A., Wornow, M., Birch-Sykes, C., Massaroli, S., Patel, A., Rabideau, C., Bengio, Y., et al. Hyenadna: Long-range genomic se- quence modeling at single nucleotide resolution. NeurIPS, 2023.
Li, X., Wang, W. Hu, X., and Yang, J. Selective kernel networks. In CVPR, pp. 510-519, 2019.

---

Stochastic Layer-Wiss Shuffle for Improving Vision Mamba Training
Noroozi, M. and Favaro, Unsupervised learning of visual representations by solving jigsaw puzzles. In ECCV, pp- 69-84, 2016.
Tong, Z., Song, Y., Wang, J., and Wang, L.. Videomae: Masked autoencoders are data-efficient learners for self- supervised video pre-training. NeurIPS, pp. 10078- 10093, 2022.
Patro, B. N. and Agneeswaran, S. Mamba-360: Survey of state space models as transformer alternative for long se- quence modelling: Methods, applications, and challenges. arXiv preprint NiNi::2404.16112. 2024.
Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. Training data-efficient image trans- formers & distillation through attention. In ICMI,, pp. 10347-10357, 2021.
Pei, X., Huang, T., and Xu, C. Efffiiiettvmmmb: Atrous se- lective scan for light weight visual mamba. arXiv preprint arXiv:2403.09977, 2024.
Ulyanov, D., Vedaldi, A., and Lempitsky, V. Instance nor- malization: The missing ingredient for fast stylization. arXiv preprint
2016.
Peng, Z., Dong. L.,, Bao, H., Ye, Q., and Wei, F. Beit v2: Masked image modeling with vector-quantized visual tokenizers. arXiv preprint htti:2208:06666 2022.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In NeurlPS, 2017.
Peng, Z., Dong, L., Bao, H., Ye, Q. and Wei, F. A unified view of masked image modeling. TMLR, 2023.
Wang. F., Wang, J., Ren, S., Wei, G., Mei, J., Shao, W., Zhou, Y., Yuille, A., and Xie, C. Mamba-r: Vision mamba also needs registers. arXiv preprint arXiv:2405. 14858. 2024.
Phung, H., Dao, Q., Dao, T., Phan, H., Metaxas, D., and Tran, A. Dimsum: Diffusion mamba a scalable and unified spatial-frequency method for image generation. In NIPS, 2024.
Wang, W., Xie, E., Li, X., Fan, D.-P., Song, K., Liang, D., Lu, T., Luo. P., and Shao, L. Pyramid vision transformer: A versatile backbone for dense prediction without convo- lutions. In ICCV, pp. 568-578, 2021.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G. Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In ICML, pp. 8748-8763, 2021.
Wang. W., Xie, E.,, Li, X., Fan, D.-P., Song, K., Liang, D., Lu, T., Luo, P., and Shao, L. Pvt v2: Improved baselines with pyramid vision transformer. Comput. Vis. Media, 8: 415-424, 2022.
Radosavovic, I., Kosaraju, R. P., Girshick, R., He, K., and Dollár, P. Designing network design spaces. In CVPR, pp. 10428-10436, 2020.
Wei, C., Fan, H., Xie, S., Wu, C.-Y,, Yuille, A., and Feicht- enhofer, C. Masked feature prediction for self-supervised visual pre-training. In CVPR, 2022a.
Ren, S., Li, X., Tu, H., Wang, F., Shu, F., Zhang, L., Mei, J., Yang, I.,, Wang, P., Wang, H., et al. Autoregres- sive pretraining with mamba in vision. arXiv preprint arXiv:2406.07537, 2024.
Wei, L., Xie, L., Zhou, W., Li, H., and Tian, Q. Mvp: Multimodality-gun -uided visual pre-traine pre-training. In ECCV, pp. 337-353, 2022b.
Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. ICLR, 2015.
Wu, H., Yang. Y., Xu. H., Wang, W., Zhou, J., and Zhu, Rainmamba: Enhanced locality learning with state space models for video deraining. In ACM MM, pp. 7881-7890, 2024.
Smith, J. T., Warrington, A., and Linderman, S. W. Simpli- fied state space layers for sequence modeling. In ICLR, 2023.
Wu, Y. and He, K. Group normalization. In ECCV, pp. 3-19, 2018.
Srivastava, N.. Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Dropout: simple way prevent neural networks from overfitting. J. Mach. Learn. Res., 15:1929-1958, 2014.
Xia, Z., Pan, X., Song, S., L.. E.,, and Huang, G. Vision transformer with deformable attention. In CVPR, 2022.
Xiao, C., Cao, Q., Luo, Z., and Lan, L. Mambatrack: a simple baseline for multiple object tracking with state space model. In ACM MM, pp. 4082 00924099,11 2024.
Tallec, C. and Ollivier, Y. Can recurrent neural networks warp time? In ICLR, 2018.
Tang, L., Xiao, H., Jiang, P.-T., Zhang, H., Chen, J., and Li, B. Scalable visual state space model with fractal scanning. arXiv preprint ar055:2405. 14480, 2024.
Xiao, T., Liu, Y., Zhou, B., Jiang, Y., and Sun, J. Unified perceptual parsing for scene understanding. In ECCV, 2018.
12

---

Stochastic Layer- Wise Shuffle for Improving Vision Mamba Training
Xie, Z., Zhang, Z.,, Cao, Y., Lin, Y., Bao, Yao, Z., Dai, Q.. and Hu, H. Simmim: A simple framework for masked image modeling. In CVPR, pp. 9653-9663, 2022.
Yang, C., Chen, Z.., Espinosa, M., Ericsson, L., Wang. Z.,, Liu, J., and Crowley, E. J. Plainmamba: Improving non- hierarchical mamba in visual recognition. In BMVC, 2024a.
Yang, S., Wang. Y., and Chen, H. Mambamil: Enhanc- ing long sequence modeling with sequence reordering in computational pathology. In MICCAI, pp. 296-306, 2024b.
Yue, and Li, Z. Medmamba: Vision mamba for medical image classification. arXiv preprint vrXiv:2403.03844 2024.
Yun, S., Han, D., Oh, S. J., Chun, S., Choe, J., and Yoo, Y.
Cutmix: Regularization strategy to train strong classifiers with localizable features. In ICCV, pp. 6023-6032, 2019.
Zhang. B. and Sennrich, R. Root mean square layer normal- ization. In NeurlPS, 2019.
Zhang. G., Fan, L., He, C.,, Lei, Z., Zhang, Z., and Zhang, L. Voxel mamba: Group-free state space models for point cloud based 3d object detection. In NIPS, 2024a.
Zhang, G., Liu, C., Cui, Y., Zhao, X., Ma, K., and Wang, L. Vfimamba: Video frame interpolation with state space models. In NIPS, 2024b.
Zhang, H., Cisse, M., Dauphin, N., and Lopez-Paz, D. mixup: Beyond empirical risk minimization. In ICLR, 2018.
Zhang. J., Bian, K., Cheng, P., An, W., Liu, J., and Zhou, Vim-f: Visual state space model benefiting from learning in the frequency domain. arXiv preprint a40iv:2405. 24679 2024c.
Zhou, J., Wei, C., Wang, H., Shen, W., Xie, C., Yuille, A., and Kong, ibot: Image bert pre-training with online tokenizer. ICLR, 2022.
Zhu, L., Wang, X., Ke, Z., Zhang, W., and Lau, R. W. Bi- former: Vision transformer with bi-level routing attention. In CVPR, 2023.
7hu, I.,, Liao, B., Zhang, Q. Wang. X., Liu, W., and Wang, X. Vision mamba: Efficient visual representation learning with bidirectional state space model. In /CML, 2024.
13

---

Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training
A. Implementation Details
A.I. Supervised Training Settings
In supervised classification training, we build our pipeline with the codebase of DeiT (Touvron et al., 2021) and the setups are listed in the following Table A.1. For MambaMLP-L training, the layer-wise shuffle rate and drop path rate are 0.5. For MambaR, follow the training setup from Wang al. (2024), which employs three-stage strategy equivalent approximately 220 epochs of training input resolution of 224 except setting layer-wise shuffle rate to 0.1.
Table A.1: Supervised training implementation settings.
Config
Base Middle
Large
optimizer
AdamW
base learning rate
5e-4
weight decay
0.1
0.15
layer-wise Ir decay
0
learning rate schedule
cosine decay
batch size
2048
1024
warmup epochs
30
training epochs
300
200
augmentation
RandAug (9, 0.5)
label smoothing
0.1
mixup
0.8
cutmix
1
reprob
0.25
drop path rate
0.5
0.7
layer-wise shuffle rate
0.5
0.6
EMA decay
0.99992
0.99992
A.2. Pre-training Settings
We provide the configurations of models used in MFD pre-training in the following Table A.2. During pre-training, we use image sizes of 192 and 224 for the MAE and MFD pipelines, respectively, and the masking ratio for MAE is 0.7. The 224 resolution is a common training setting adopted by works like Vim(Zhu et al., 2024), ViT((osovitskii et al., 2021), unique design requirement of dividing images
and MAB(He al., 2022). ARM(Ren et al., 2024) uses 192 due into multiple 64x64 patch groups. To ensure fair comparison under the same training epochs, their MAE experiment also followed this resolution. For MFD pre- training, used the standard 224 resolution but with only 18.75%-31 37.5% of ARM's training epochs, significantly reducing computational costs.
Table A.2: Pre-train Pre-training training implementation settings.
Config
Base
Large
Huge
optimizer
AdamW
base learning rate
1.5e-4
weight decay
0.05
learning rate schedule
cosine decay
batch size
2048
1024
1024
warmup epochs
30
training epochs
300
300
300
shuffle rate
0.1
0.4
0.6
masking ratio
0.5
0.6
0.6
augmentation
RandomResizedCrop
14


Reviews:
[{"review_id": "7weGGrSv0t", "reviewer": "Reviewer_mf21", "summary": "This paper proposes a stochastic layer-wise shuffle regularization (SLWS) method for efficient vision mamba training. As a plug-and-play method, SLWS mitigates the overfitting problem with introducing minimal overhead. The achieved results are impressive and downstream tasks also verified the effectiveness. Overall, the idea is novel and interesting and the method only has minor weaknesses.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Q1 & Q2: Hierarchical VMamba training and non-hierarchical plain mamba selection**.\n\n**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.  \n\n**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.\n\n\n\n> **Q3: 220 training epoch setting on some models**.\n\n**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.\n\n\n\n> **Q4: Training configuration of our models**.\n\n**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.\n\n\n\nThank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed all of your remaining concerns.", "claims_and_evidence": "YES", "methods_and_evaluation_criteria": "YES", "theoretical_claims": "YES", "experimental_designs_or_analyses": "YES", "relation_to_broader_scientific_literature": "Vision mamba has achieved great success in visual tasks and training process indeed a problem. This paper proposes a interesting stochastic layer-wise shuffle strategy to help the training process and achieves good results without increased overhead.", "essential_references_not_discussed": "NO", "other_strengths_and_weaknesses": "1. I think some vision mamba models have not the training issue, such as VMamba [NeurIPS'24];\n2. Why the method focuses on the plain mamba models?\n3. Why training for 220 epochs for some models, that is very strange. Did the author meet the overfitting problem?\n4. Are there any other tricks or modules or technique during the training of your models?", "other_comments_or_suggestions": "See above", "questions_for_authors": "See above", "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": "Mmta56lPny", "reviewer": "Reviewer_8tDk", "summary": "The paper introduces Stochastic Layer-Wise Shuffle (SLWS), a method designed to enhance the training of Vision Mamba models (ViM). This approach involves applying stochastic shuffling to input tokens at each layer, with the probability of shuffling systematically increasing as the layer depth progresses. Though conceptually simple, SLWS demonstrates significant benefits: it mitigates overfitting, promotes positional invariance across successive blocks, and improves model robustness. These advantages translate to measurable performance gains in both supervised and unsupervised pre-training paradigms.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **Q1-1: Supporting evidence of  learning curves for mitigating overfitting of SLWS**.\n\n**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].\n\n> **Q1-2: Qualitative examples for SLWS**.\n\n**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.\n\n> **Q2: Direct segmentation comparisons of SLWS vs. non-SLWS**\n\n**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.\n\n|     metric\\model| SLWS | non-SLWS |\n| --------- | ---- | -------- |\n| Cls. acc. | 82.7 | 81.2     |\n| Seg. mIoU | 47.0 | 45.2    |\n\n> **Q3 (review content in Experimental Designs Or Analyses): SLWS’s gains and relationship to training paradigm**\n\n**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.\n\n> **Q4 Table 4: Clarify the units, indicate if less/more is better**\n\n**A4:** Thanks for your suggestion and we have added such information to the revised version.\n\n\n\n------\n\nThank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed your remaining concerns.\n\n[1] Deep networks with stochastic depth, ECCV'16.", "claims_and_evidence": "1. Potential for Stronger Empirical Support on Overfitting Mitigation\nSLWS is presented as an effective method for mitigating overfitting in supervised training while maintaining computational efficiency and architectural simplicity. While the paper references extensive experiments to support this claim, the direct empirical evidence—limited to two learning curves—shows only modest reductions in the training-validation accuracy gap. To further substantiate the connection between SLWS and improved generalization, the authors could enhance this section with additional metrics (e.g., validation loss trends across epochs, comparisons of parameter sensitivity) or targeted ablation studies. Such additions would help clarify how the stochastic shuffling mechanism specifically contributes to reducing overfitting.\n\n2. Opportunities for Qualitative Insights in Downstream Tasks\nThe improved performance of SLWS in semantic segmentation and other downstream tasks is numerically compelling. However, the analysis could be enriched by qualitative examples illustrating how the method fosters robust or positionally invariant features. For instance, visualizations of feature activation patterns (e.g., attention maps or segmentation boundaries) in SLWS-trained models versus baselines could offer intuitive insights into why the method succeeds. While deeper mechanistic analysis might extend beyond the scope of this work, even simple visual comparisons would strengthen the paper’s narrative and provide readers with a clearer understanding of SLWS’s practical benefits.", "methods_and_evaluation_criteria": "yes.", "theoretical_claims": "The algorithm 1 for LWS forward is clear and correct.", "experimental_designs_or_analyses": "1. Supervised Classification Improvements and Comparison Scope\nThe paper demonstrates consistent, albeit modest, improvements in supervised classification tasks when applying SLWS, with accuracy gains ranging from 0.4% to 2.9% across Vision Mamba (ViM) and MambaMLP variants. These gains are further shown to generalize across different backbones and training methodologies. However, the analysis primarily evaluates Mamba-based models trained with Masked Feature Distillation (MFD), a specific regularization strategy. While the results are promising, it would be valuable to clarify whether SLWS’s benefits are intrinsic to the method itself or partially reliant on MFD’s regularization effects. Including comparisons with non-MFD-trained baselines (e.g., vanilla ViM or non-Mamba architectures) could help establish SLWS’s broader applicability and ensure equitable comparisons with other training regimes.\n\n2. Clarity in Segmentation Results\nIn semantic segmentation experiments, the performance gains attributed to SLWS are presented in a table that lacks direct counterparts for certain configurations. For example, the ViM-M models trained with SLWS do not appear to have equivalent baselines without SLWS. This omission complicates efforts to isolate the method’s contribution from architectural or training-specific factors. Providing explicit comparisons for all configurations—or transparently acknowledging these limitations—would enhance interpretability and strengthen confidence in SLWS’s role in the observed improvements.", "relation_to_broader_scientific_literature": "The paper introduces the shuffling for vim models. It shows yet another transformation to encourage model robust and invariant representation.  Scaling the models has been problematic due to overfitting and other constrains during training.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The proposed Stochastic Layer-Wise Shuffle (SLWS) stands out for its conceptual simplicity and computational efficiency, requiring no architectural modifications while delivering consistent performance gains. The method demonstrates measurable improvements across diverse tasks, including supervised classification (with accuracy increases of 0.4% to 2.9%) and semantic segmentation, suggesting broad applicability. Its ability to enhance positional invariance and reduce overfitting—even under varied backbones and training regimes—further underscores its potential as a versatile training aid.\n\nWeaknesses:\nWhile the empirical results are promising, the paper’s claims would benefit from additional supporting evidence to solidify mechanistic insights. For instance:\n\nData granularity: Including learning curves (e.g., training vs. validation loss/accuracy trends) would help visualize how SLWS mitigates overfitting in practice.\n\nQualitative examples: Visualizations of feature activations or segmentation outputs could clarify how SLWS fosters robustness or invariance.\n\nInterpretability of comparisons: Some results tables (e.g., semantic segmentation) lack direct SLWS vs. non-SLWS comparisons for key configurations (e.g., ViM-M), making it challenging to isolate the method’s impact. A brief discussion justifying the presentation format or addressing potential fairness concerns (e.g., MFD-focused evaluations) would enhance methodological transparency.", "other_comments_or_suggestions": "* Table 4: Clarify the units, indicate if less/more is better.", "questions_for_authors": "* Would it be possible to provide more evidence about the help of the method to address overfitting in Mamba models?", "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": "Lzt57AMBmS", "reviewer": "Reviewer_6YiU", "summary": "This paper introduces a method that addresses overfitting issues when scaling up vanilla Vision Mamba models to larger sizes. The key contribution is a Stochastic Layer-Wise Shuffle (SLWS) regularization technique that randomly shuffles token positions during training with layer-dependent probabilities. Experiments also improvements on classification, detection and segmentation tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **Q1: Performance improvements of our models compared to baselines especially Mamba-R**\n\n**A1:** Our proposed SLWS regularization demonstrates significant improvements across both supervised and self-supervised training paradigms. For example: \n\n- 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.  \n- The MambaMLP-H model trained with SLWS reaches 87.5% accuracy, setting a new state-of-the-art result for vision Mamba variants.\n\nThese advancements have been acknowledged by reviewers, who highlighted the \"impressive results,\" \"measurable improvements,\" and \"significant performance gains\" in their feedback.  \n\nRegarding 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.\n\n> **Q2 & Q3: Concerns about disrupting inherent locality and positional information of image data and dense prediction**\n\n**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.  \n\nFor 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.\n\n> **Q4 & Q5: Applying SLWS to hierarchical variants and ViT**\n\n**A4 & A5:**\n\n- Application to Hierarchical Variants:  \n  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).\n\n- Application to Vision Transformers (ViTs):  \n  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.\n\n> **Q6: Shuffling strategies with locality**\n\n**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.\n\n> **Q7 (review content in Claims And Evidence): \"The claim that vanilla Vision Mamba models couldn't previously be scaled up\"**\n\n**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.\"*  \n\n------\n\nThank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed your remaining concerns.", "claims_and_evidence": "The evidence partially supports the paper's claims but has several inconsistencies:\n\n1. The claim that it outperforms similarly-sized models is only selectively supported. When compared to Mamba-Reg models, the improvements are marginal (MambaR-B: 83.0% vs. MambaR-B with SLWS: 83.1%).\n2. The claim that vanilla Vision Mamba models couldn't previously be scaled up is contradicted by cited work, including Mamba-R and ARM, which have successfully scaled Vision Mamba to large and even huge sizes.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "No theoretical claim.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "1. Recent work in Vision Mamba scaling, particularly Mamba-R [1] and ARM [2], which have successfully scaled the Mamba up.\n1. The relationship to positional invariance in vision models and how token shuffling affects spatial understanding.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths:\n1. Clear motivation and well-written presentation\n2. Simple, plug-and-play approach with no inference overhead\n3. Comprehensive evaluation across multiple vision tasks\n\nWeaknesses:\n\n1. Limited performance improvements compared to baseline models especially Mamba-R\n2. Concerns about disrupting inherent locality of image data through token shuffling\n3. Potentially harmful to dense prediction tasks where positional information is crucial", "other_comments_or_suggestions": "None.", "questions_for_authors": "1. Can SLWS apply to hierarchical variants?\n2. Would other shuffling strategies with locality help?\n3. Can you apply SLWS on vision transformer if the claims hold?", "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": "p5M5p7fNZR", "reviewer": "Reviewer_knen", "summary": "This work introduces a stochastic hierarchical shuffle strategy SLWS for Vision Mamba (Vim) that successfully solves the overfitting issue of Mamba models in large-scale datasets without changing the model architecture, effectively improving the training of Vim. According to this paper, SLWS can help Vim achieve leading performance on the officially recognized ImageNet-1K without introducing significant additional computational overhead.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **Q1: SLWS's shuffle effects for visual connections between tokens**.\n\n**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.\n\n> **Q2: Training resolution setting and performance difference between training pipelines.**\n\n**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.\n\n>  **Q3: Hierarchical VMamba's effect in mitigating overfitting.**\n\n**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).\n\n> **Q4: Performance improvement based on MambaR.**\n\n**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.\n\n> **Q5: Horizontal and vertical comparisons of MAE, ARM and MaskDistill.**\n\n**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.\n\n> Q6: Reversed layer-wise probability experiment for SLWS's semantic awareness priori.\n\n**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.\n\n| probability config. | constant | layer-wise | reversed layer-wise |\n| ------------------- | -------- | ---------- | ------------------- |\n| **Acc.**            | 81.1     | 82.7       | 81.2                |\n\n> **Suggestions about Figure 2 and related work.**\n\n**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].\n\n------\n\nThank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed your remaining concerns.", "claims_and_evidence": "Yes, this paper provides the corresponding computational analysis and experimental proof for SLWS.", "methods_and_evaluation_criteria": "Yes, this work effectively improves the overfitting problem of Mamba models in large-scale datasets.", "theoretical_claims": "Yes, This work asserts that SLWS possesses negligible computational overhead, and the efficiency analysis in Section 3.2.1 and Table 4 provide the corresponding proof. This paper claims that SLWS can solve the overfitting problem of the Mamba model on large datasets, and Table 1 and Figure 2 provide the corresponding experimental proofs.", "experimental_designs_or_analyses": "The effectiveness of SLWS in solving the overfitting problem of Mamba model is demonstrated in Figure 2 by analyzing the training loss and evaluation loss with and without SLWS.", "relation_to_broader_scientific_literature": "[1] similarly discusses the impact of a training strategy for learning token location invariance on the model and demonstrates its effectiveness.\nReferences: [1] Efficient Training of Visual Transformers with Small Datasets", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths. \n1. The proposed SLWS significantly improves the performance of non-hierarchical Mamba models on large-scale datasets. \n2. SLWS does not incur significant computational overhead. 3. Numerous experiments have demonstrated the effectiveness of SLWS.\n\n\nWeakness. \n1. Token shuffle may break some connections between tokens that are supposed to be visually linked, and how SLWS avoids it. \n2. During pre-training, this work uses image sizes of 192 and 224 for the MAE and MFD pipelines, respectively. Why different pre-training resolution settings are used for ARM, MAE and MFD. Could this be the reason for the difference in performance between MAE, ARM and MFD. \n3. We observe that VMamba with similar parameter counts has higher performance compared to Vim paired with the SLWS strategy (83.9 vs 82.7). Whether hierarchical visual Mamba design is more effective in mitigating the overfitting problem of Mamba model in large-scale dataset. \n4. The improvement of SLWS for MambaR looks limited in Table 1, what is the reason for this. \n5. The MAE and ARM comparison setups pre-trained from scratch seem to be inadequate when compared to MFDs that have introduced pre-trained CLIP distillation. How does MFD compare to some distillation pre-training methods, such as MaskDistill paired with MambaMLP results. \n6. The ablation in Table 5 only looks to demonstrate that a dynamic shuffle strategy is better than a constant one, and is not sufficient to justify the proposed “realization of stronger semantic awareness in deeper layers requires translational invariance of patch locations, whereas shallower layers must remain location-sensitive”. What would be the effect on the results if a larger PL was used for shallower layers and a smaller PL for deeper layers.", "other_comments_or_suggestions": "1. The LWS in the legend of Figure 2 is not specified in the text as to whether it refers to a SLWS . \n2. Related work vision backbone example of the reference work is a bit out of date.", "questions_for_authors": "Although SLWS clearly improves the performance of non-hierarchical Mamba models on large-scale datasets, some key experiments and some explanations (see weakness) are needed. I suggest that the authors provide more experimental evidence to further support their work.", "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": "3IUmSBrBlE", "reviewer": "Reviewer_bypj", "summary": "This paper proposes a plug-and-play training strategy for Vision Mamba. It shuffers the sequence order of the input tokens layer by layer. The authors conduct mask feature distillation to pre-train the vision Mamba with the proposed layerwise shuffling strategy. Experiments on classification and dense prediction tasks show improvement. The overall writing is easy to understand, and the method is simple and clean.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **Q1: Applying SLWS to hierarchical structures.**\n\n**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 currentl*y 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.\n\nIn 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.\n\n> **Q2: Linear layer-wise probability design reasons**\n\n**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.\n\n> **Q3: Applying SLWS to diverse training paradigms and their horizontal and vertical comparisons.**\n\n**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.\n\n------\n\nThank you for all your comments, which we believe have strengthened our work; we hope our responses have addressed your remaining concerns.\n\n[1] Deep networks with stochastic depth, ECCV'16.", "claims_and_evidence": "Yes. The author conducts extensive experiments to validate the method. Several baseline methods， such as Vim and MambaMLP, are adopted to test the proposed method.", "methods_and_evaluation_criteria": "Yes, the IoU and accuracy on classification detection and segmentation tasks are reported.", "theoretical_claims": "Yes, the method is simple and easy to understand. A shuffling strategy is applied to Mamba for pertaining.", "experimental_designs_or_analyses": "Yes, the experimental design follows the masked feature distillation pipeline.", "relation_to_broader_scientific_literature": "The method could impact the pretraining of current vision mamba based method.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n\n1. The method is technically sound and well-motivated. Only shuffling the token order is simple and effective, which can impact more vision Mamba works. \n\n2. The experiments are extensive, including pretraining, classification, and dense prediction tasks. \n\n3. The performance improvement shows the effectiveness of the method. \n\nWeakness. \n\n1. The shuffling strategy is only applied to plain architecture. While most vision Mamba have hierarchical structures such as VMamba [1], the authors are suggested to apply it to hierarchical structures. Otherwise, the impact of this work is constrained. \n\n2. The layerwise probability design seems to be handcrafted. Intuitive reasons are not enough to explain it. Please provide more solid reasons or develop a more dedicated design for the layerwise probability. \n\n3. The effects of applying a shuffling strategy directly instead of masked feature distillation should be explored in more detail since it is the most straightforward way to apply the proposed shuffling strategy. For example, is it possible to apply the shuffling strategy to the direct pretraining of image classification or MAE-style [2] pretraining?  Since no teacher model weights are used, it can be more convenient for other followers. \n\n[1] VMamba: Visual State Space Model in NeurIPS24. https://arxiv.org/pdf/2401.10166\n\n[2] Masked Autoencoders Are Scalable Vision Learners in CVPR22. https://arxiv.org/abs/2111.06377", "other_comments_or_suggestions": "Please address the concerns in weakness.", "questions_for_authors": "None", "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:

Meta-review: All reviewers are in favor of acceptance for this paper, in particular they note that the proposed Stochastic Layer-Wise Shuffle (SLWS) distinguishes itself through its straightforward design and computational efficiency, achieving consistent performance improvements without the need for any architectural changes. So I m in favor of accepting this paper to ICML
Decision: Accept (poster)