Stochastic Layer-Wise Shuffle for Improving Vision Mamba Training

📄 arXiv: 2408.17081v2 📥 PDF

作者: Zizheng Huang, Haoxing Chen, Jiaqi Li, Jun Lan, Huijia Zhu, Weiqiang Wang, Limin Wang

分类: cs.CV

发布日期: 2024-08-30 (更新: 2025-06-02)

备注: accpeted to ICML25

期刊: Proceedings of the 42nd International Conference on Machine Learning, 2025

💡 一句话要点

提出随机分层Shuffle方法，提升Vision Mamba在ImageNet上的训练效果

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: Vision Mamba 正则化 随机Shuffle 图像分类 深度学习

📋 核心要点

  1. Vision Mamba虽然具有线性复杂度优势，但其训练方法和潜力仍待挖掘，需要更有效的训练策略。
  2. 提出随机分层Shuffle (SLWS) 正则化方法，通过层相关的token shuffle来改善Vim的训练。
  3. 实验表明，SLWS无需修改Vim架构，即可在ImageNet-1K上取得领先性能，提升了模型效果。

📝 摘要（中文）

Vision Mamba (Vim) 模型在序列长度上表现出近乎线性的复杂度，使其在处理视觉数据方面极具吸引力。然而，针对Vim的训练方法及其潜力尚未得到充分探索。本文研究了Vim的训练策略，并提出了一种新的正则化方法——随机分层Shuffle (SLWS)，该方法可以有效改善Vim的训练。无需修改架构，该方法就能使非层级Vim在ImageNet-1K上获得与同类型模型相比领先的性能。SLWS通过每层四个简单的步骤进行操作：概率分配以指定层相关的shuffle率，通过伯努利试验进行操作采样，对输入token进行序列shuffle，以及恢复输出的顺序。SLWS具有三个特点：(1) 即插即用：无需架构修改，并且在推理期间停用。(2) 简单但有效：这四个步骤仅引入随机排列和可忽略的开销。(3) 直观设计：Shuffle概率随层深度线性增长，与视觉模型中的分层语义抽象相一致。这项工作强调了为Vim模型量身定制训练策略的重要性，并为探索其可扩展性提供了一种有用的方法。

🔬 方法详解

问题定义：现有的Vision Mamba训练方法可能存在过拟合或泛化能力不足的问题，无法充分发挥其在视觉任务中的潜力。痛点在于如何设计一种有效的正则化方法，在不增加过多计算负担的前提下，提升模型的训练效果和泛化能力。

核心思路：核心思路是引入随机的token shuffle，作为一种正则化手段，来增强模型的鲁棒性和泛化能力。通过在不同层应用不同程度的shuffle，模拟数据增强的效果，防止模型过度依赖于特定的token顺序。shuffle概率随层深度线性增长，符合视觉模型中语义抽象的层级特性。

技术框架：SLWS方法主要包含四个步骤：1) 概率分配：为每一层分配一个shuffle概率，概率值随层深度线性增长。2) 操作采样：使用伯努利试验，根据分配的概率决定是否对该层进行shuffle操作。3) 序列Shuffle：如果决定进行shuffle，则对输入token的序列进行随机排列。4) 顺序恢复：在shuffle操作后，将输出的token顺序恢复到原始顺序。该方法在训练过程中使用，推理阶段停用。

关键创新：关键创新在于提出了层相关的随机shuffle策略，并将其与视觉模型的层级语义抽象相结合。与传统的全局shuffle或dropout方法不同，SLWS允许不同层具有不同的shuffle强度，从而更好地适应不同层级的特征表示。此外，SLWS是一种即插即用的方法，无需修改模型架构，易于集成到现有的Vim训练流程中。

关键设计：SLWS的关键设计包括：1) shuffle概率的线性增长策略，确保深层网络具有更强的shuffle强度。2) 使用伯努利试验进行操作采样，引入随机性，避免模型过度依赖于shuffle操作。3) 在shuffle后恢复原始顺序，保证模型输出的一致性。没有引入额外的损失函数或复杂的网络结构。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SLWS方法在ImageNet-1K数据集上取得了显著的性能提升，在不修改模型架构的前提下，达到了与同类型模型相比领先的性能。具体的性能数据和对比基线在论文中给出，证明了SLWS作为一种有效的正则化方法，能够显著改善Vision Mamba的训练效果。

🎯 应用场景

该研究成果可广泛应用于各种视觉任务，如图像分类、目标检测、语义分割等。通过提升Vision Mamba的训练效果，可以提高模型在这些任务上的性能和泛化能力。此外，该方法还可以应用于其他基于序列模型的视觉任务，例如视频理解和动作识别。未来，该方法有望推动Vision Mamba在实际应用中的普及。

📄 摘要（原文）

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.