动态调整视觉Mamba模型的计算量,提升效率且精度损失小。
Dynamic Vision Mamba
- 针对Mamba结构定制令牌剪枝,避免训练推理不一致
- 按图像需求动态选择状态空间块,减少35.2%计算量
- 适配多种模型架构与任务,适合部署优化场景
基于Mamba的视觉模型因计算效率高于注意力模型而受到广泛关注。然而,这些模型仍存在空间冗余问题,表现为令牌冗余和模块冗余。针对令牌冗余,我们发现早期剪枝方法会导致训练与推理不一致或引入额外推理开销。为此,我们为Mamba结构定制了令牌剪枝策略:在输入下一Mamba模块前重新排列剪枝后的序列。针对模块冗余,我们基于观察——Mamba视觉模型的推理速度主要受状态空间块数量影响——允许每张图像动态选择S个状态空间块(SSM blocks)。所提出的Dynamic Vision Mamba(DyVM)方法在仅损失1.7%准确率的前提下,实现Vim-S上35.2%的浮点运算量(FLOPs)减少。该方法在不同Mamba视觉模型架构和多种视觉任务中均具有良好泛化能力。代码将公开。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。