摘要
状态空间模型 (SSM) 与 Transformer 中的注意力模块相比,具有线性计算复杂度的优势,已被用作一种新型强大的视觉基础模型应用于视觉任务。受视觉 Transformer (ViT) 中最终预测仅基于一小部分信息量最大的 token 的观察结果启发,我们采取了通过基于 token 的剪枝来提高 SSM 基于视觉模型效率的新方法。然而,现有的为 ViT 设计的 token 剪枝技术的直接应用无法提供良好的性能,即使经过广泛的微调。为了解决这个问题,我们重新审视了 SSM 的独特计算特性,并发现简单的应用会破坏顺序 token 位置。这一见解促使我们设计了一种专门针对基于 SSM 的视觉模型的新颖且通用的 token 剪枝方法。我们首先介绍了一种剪枝感知的隐藏状态对齐方法,以稳定剩余 token 的邻域,从而提高性能。此外,根据我们详细的分析,我们提出了一种适用于 SSM 模型的 token 重要性评估方法,以指导 token 剪枝。通过高效的实现和实用的加速方法,我们的方法带来了实际的加速。大量的实验表明,我们的方法可以在不同任务中实现显著的计算量减少,而性能影响最小。值得注意的是,我们在 ImageNet 上获得了 81.7% 的准确率,同时将剪枝后的 PlainMamba-L3 的 FLOPs 减少了 41.6%。此外,我们的工作提供了更深入的见解,有助于理解基于 SSM 的视觉模型的行为,从而促进未来的研究。