Research map / Architectures and efficiency / Attention and transformers
Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Stochastic Layer-Wise Shuffle (SLWS), a training-only regularizer that randomly permutes token order in non-hierarchical Vision Mamba layers with depth-linear probability, and combines it with CLIP-teacher masked feature distillation to reach 87.6% ImageNet-1K for Vision Mamba.
Problem
Vision Mamba (Vim) models overfit and collapse when scaled up, and they trail ViTs on ImageNet-1K; training methodology for non-hierarchical Vim is under-explored.
Approach
Randomly shuffle the input token sequence of each Mamba layer with a Bernoulli trial whose probability increases linearly with depth, then restore the original order at the output; the method is disabled at inference and adds no architectural change.
Key results
- SLWS raises supervised Vim-B on ImageNet-1K from 79.8% to 82.7% and lets Vim-L train without collapse (82.9%).
- MFD + SLWS with MambaMLP-H reaches 87.6% top-1 on ImageNet-1K, reported as the new SOTA for Vision Mamba.
- SLWS reduces training throughput by under 2% across resolutions 128 to 768.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add SLWS to supervised training | Vim-M · ImageNet-1K | accuracy: 80.9 -> 82.8 | Helped |
| Add SLWS to supervised training | Vim-B · ImageNet-1K | accuracy: 79.8 -> 82.7 | Helped |
| Add SLWS to supervised training | Vim-L · ImageNet-1K | accuracy: collapsed -> 82.9 | Helped |
| Add SLWS to supervised training | MambaR-B · ImageNet-1K | accuracy: 83.0 -> 83.1 | No clear effect |
| Add SLWS to MAE pre-training | MambaMLP-B · ImageNet-1K | accuracy: 81.6 -> 82.0 | Helped |
| Add SLWS to MFD pre-training | MambaMLP-L · ImageNet-1K | accuracy: 86.4 -> 86.7 | Helped |
| Use MFD pre-training with SLWS | MambaMLP-H · ImageNet-1K | accuracy: 87.6 | Helped |
| Add SLWS to MambaMLP-L | MambaMLP-L · ImageNet-1K | accuracy: 81.4 -> 82.9 | Helped |
| Use layer-dependent linear probabilities with P_L=0.5 instead of decreased linear probabilities (ablation) | Vim-M · ImageNet-1K | accuracy: 81.2 -> 82.7 | Helped |
| Use layer-dependent linear probabilities with P_L=0.5 instead of exponential probabilities (ablation) | Vim-M · ImageNet-1K | accuracy: 82.2 -> 82.7 | Helped |
| Use layer-dependent linear probabilities with P_L=0.5 instead of constant probabilities (ablation) | Vim-M · ImageNet-1K | accuracy: 81.5 -> 82.7; 81.1 -> 82.7 | Helped |
| Set the layer-dependent linear shuffle rate to P_L=0.4, 0.6, or 0.7 instead of 0.5 (ablation) | Vim-M · ImageNet-1K | accuracy: 82.7 -> 82.3; 82.7 -> 82.4; 82.7 -> 82.4 | Mixed |
Surprising findings
- Applying the shuffle probability in decreasing order across layers (shallow layers shuffled more) drops accuracy to 81.2%, versus 82.7% for increasing linear schedule.
Limitations
- Hierarchical Mamba architectures with downsampling are incompatible with SLWS because order restoration fails under length reduction.
- Gains on MambaR and MambaMLP are small in some settings (e.g., +0.1 for MambaR-B) and the shuffle rate needs per-model tuning.
- Full pipeline relies on a frozen CLIP teacher, so the best results depend on extra pre-training compute.
Experimental setup
Datasets: ImageNet-1K ADE20K COCO2017
Benchmarks: ImageNet-1K ADE20K COCO2017
Metrics: top-1 accuracy mIoU AP box AP mask throughput
Baselines: Vim MambaR MambaMLP ARM MAE VMamba VideoMamba ViT-L MAE MaskDistill BEiT v2
Builds on: Mamba Vision Mamba stochastic depth masked feature distillation MambaR ARM
Closest papers by meaning
- MambaOut: Do We Really Need Mamba for Vision? (CVPR 2025 · similarity 0.88)
- Mamba-Adaptor: State Space Model Adaptor for Visual Recognition (CVPR 2025 · similarity 0.87)
- EfficientVMamba: Atrous Selective Scan for Light Weight Visual Mamba (AAAI 2025 · similarity 0.86)
- Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (ICML 2024 · similarity 0.86)
- Adventurer: Optimizing Vision Mamba Architecture Designs for Efficiency (CVPR 2025 · similarity 0.86)
- Multi-Scale VMamba: Hierarchy in Hierarchy Visual State Space Model (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.