针对视觉状态空间模型设计新剪枝方法,大幅降耗几乎不丢精度。
Exploring Token Pruning in Vision State Space Models
- 提出适配SSM特性的新剪枝方法,解决传统剪枝破坏序列位置的问题。
- ImageNet上实现41.6%计算量减少,准确率仍达81.7%。
- 适用于多种视觉任务,为SSM模型研究提供新视角。
状态空间模型(SSMs)相比Transformer的注意力模块具有线性计算复杂度优势,已被用作新型强大的视觉基础模型。受视觉Transformer中最终预测仅依赖少数关键令牌的启发,本文首次尝试通过令牌剪枝提升基于SSM的视觉模型效率。然而,直接套用现有针对ViT的剪枝技术效果不佳,即使经过大量微调也难以改善。我们重新审视SSM的独特计算特性,发现简单剪枝会破坏令牌的顺序位置关系。这一洞察促使我们设计一种专为基于SSM的视觉模型定制的新剪枝方法。首先引入剪枝感知的隐藏状态对齐机制,稳定保留令牌的邻域结构以提升性能;其次基于深入分析,提出适配SSM模型的令牌重要性评估方法,指导剪枝过程。结合高效实现与实用加速策略,该方法带来实际速度提升。大量实验表明,该方法可在不同任务上实现显著计算量降低且性能损失极小。特别地,在Pruned PlainMamba-L3上,ImageNet准确率达81.7%,同时减少41.6%的FLOPs。本工作还为未来理解基于SSM的视觉模型行为提供了更深层次见解。
原文摘要 · Abstract (English)
State Space Models (SSMs) have the advantage of keeping linear computational complexity compared to attention modules in transformers, and have been applied to vision tasks as a new type of powerful vision foundation model. Inspired by the observations that the final prediction in vision transformers (ViTs) is only based on a subset of most informative tokens, we take the novel step of enhancing the efficiency of SSM-based vision models through token-based pruning. However, direct applications of existing token pruning techniques designed for ViTs fail to deliver good performance, even with extensive fine-tuning. To address this issue, we revisit the unique computational characteristics of SSMs and discover that naive application disrupts the sequential token positions. This insight motivates us to design a novel and general token pruning method specifically for SSM-based vision models. We first introduce a pruning-aware hidden state alignment method to stabilize the neighborhood of remaining tokens for performance enhancement. Besides, based on our detailed analysis, we propose a token importance evaluation method adapted for SSM models, to guide the token pruning. With efficient implementation and practical acceleration methods, our method brings actual speedup. Extensive experiments demonstrate that our approach can achieve significant computation reduction with minimal impact on performance across different tasks. Notably, we achieve 81.7\% accuracy on ImageNet with a 41.6\% reduction in the FLOPs for pruned PlainMamba-L3. Furthermore, our work provides deeper insights into understanding the behavior of SSM-based vision models for future research.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。