Dynamic Vision Mamba

📄 arXiv: 2504.04787v1 📥 PDF

作者: Mengxuan Wu, Zekai Li, Zhiyuan Liang, Moyang Li, Xuanlei Zhao, Samir Khaki, Zheng Zhu, Xiaojiang Peng, Konstantinos N. Plataniotis, Kai Wang, Wangbo Zhao, Yang You

分类: cs.CV, cs.AI

发布日期: 2025-04-07

💡 一句话要点

Dynamic Vision Mamba (DyVM)：通过动态token剪枝和块选择提升Mamba视觉模型的效率。

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

关键词: Mamba 视觉模型 token剪枝 动态块选择 计算效率

📋 核心要点

  1. 现有基于Mamba的视觉模型存在token和块冗余，导致计算效率瓶颈，限制了其在资源受限场景的应用。
  2. DyVM通过定制token剪枝策略和动态SSM块选择机制，在减少计算冗余的同时，尽量保持模型性能。
  3. 实验表明，DyVM在Vim-S上实现了35.2%的FLOPs减少，精度损失仅为1.7%，并具有良好的泛化性。

📝 摘要（中文）

基于Mamba的视觉模型因其比基于注意力机制的模型具有更高的计算效率而受到广泛关注。然而，这些模型中仍然存在空间冗余，表现为token冗余和块冗余。对于token冗余，我们分析发现，早期的token剪枝方法会导致训练和推理之间的不一致，或者为推理引入额外的计算。因此，我们通过在将剪枝后的序列输入到下一个Mamba块之前重新排列序列，定制了token剪枝以适应Mamba结构。对于块冗余，我们允许每个图像根据经验观察动态选择SSM块，即Mamba视觉模型的推理速度很大程度上受到SSM块数量的影响。我们提出的方法Dynamic Vision Mamba (DyVM)有效地减少了FLOPs，同时性能下降很小。在Vim-S上，我们实现了35.2%的FLOPs减少，而精度损失仅为1.7%。它还在不同的Mamba视觉模型架构和不同的视觉任务中表现出良好的泛化能力。我们的代码将会开源。

🔬 方法详解

问题定义：现有基于Mamba的视觉模型虽然在计算效率上优于基于注意力机制的模型，但仍然存在显著的空间冗余，具体表现为token冗余和块冗余。Token冗余指的是图像中存在大量不包含重要信息的token，这些token参与计算会浪费资源。块冗余指的是并非所有SSM块对于所有输入图像都是同等重要的，固定数量的SSM块导致计算资源的浪费。现有token剪枝方法要么导致训练和推理不一致，要么引入额外的推理计算开销。

核心思路：DyVM的核心思路是通过动态地减少token数量和SSM块的数量来降低计算复杂度，同时尽量保持模型的性能。针对token冗余，DyVM提出了一种定制的token剪枝策略，该策略在将剪枝后的序列输入到下一个Mamba块之前重新排列序列，从而避免了训练和推理的不一致性。针对块冗余，DyVM允许每个图像根据其自身的内容动态地选择SSM块。

技术框架：DyVM的整体框架包括两个主要部分：动态token剪枝和动态SSM块选择。首先，输入图像被分割成token序列。然后，动态token剪枝模块根据token的重要性对token进行剪枝，并重新排列序列。接下来，剪枝后的token序列被输入到一系列动态选择的SSM块中。最后，模型的输出被用于执行下游视觉任务。

关键创新：DyVM的关键创新在于其动态token剪枝策略和动态SSM块选择机制。动态token剪枝策略通过重新排列序列来避免训练和推理的不一致性，并减少了推理时的额外计算开销。动态SSM块选择机制允许每个图像根据其自身的内容选择最相关的SSM块，从而提高了计算效率。

关键设计：在动态token剪枝中，使用一个可学习的mask来指示哪些token应该被保留，哪些token应该被剪枝。mask的生成基于token的激活值。在动态SSM块选择中，使用一个门控机制来控制每个SSM块是否被激活。门控信号基于图像的全局特征。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DyVM在Vim-S数据集上实现了显著的性能提升。具体来说，DyVM在仅损失1.7%的精度的情况下，将FLOPs降低了35.2%。此外，DyVM在不同的Mamba视觉模型架构和不同的视觉任务中表现出良好的泛化能力。这些结果表明，DyVM是一种有效的降低Mamba视觉模型计算复杂度的技术。

🎯 应用场景

DyVM具有广泛的应用前景，尤其是在资源受限的场景中，例如移动设备、嵌入式系统和边缘计算。它可以用于各种视觉任务，例如图像分类、目标检测和语义分割。通过降低计算复杂度，DyVM可以使Mamba视觉模型在这些场景中更高效地运行，从而提高用户体验。

📄 摘要（原文）

Mamba-based vision models have gained extensive attention as a result of being computationally more efficient than attention-based models. However, spatial redundancy still exists in these models, represented by token and block redundancy. For token redundancy, we analytically find that early token pruning methods will result in inconsistency between training and inference or introduce extra computation for inference. Therefore, we customize token pruning to fit the Mamba structure by rearranging the pruned sequence before feeding it into the next Mamba block. For block redundancy, we allow each image to select SSM blocks dynamically based on an empirical observation that the inference speed of Mamba-based vision models is largely affected by the number of SSM blocks. Our proposed method, Dynamic Vision Mamba (DyVM), effectively reduces FLOPs with minor performance drops. We achieve a reduction of 35.2\% FLOPs with only a loss of accuracy of 1.7\% on Vim-S. It also generalizes well across different Mamba vision model architectures and different vision tasks. Our code will be made public.