Exploring Token Pruning in Vision State Space Models

要約

状態空間モデル (SSM) には、トランスフォーマーのアテンション モジュールと比較して線形計算の複雑さを維持できるという利点があり、新しいタイプの強力なビジョン基盤モデルとしてビジョン タスクに適用されています。
ビジョン トランスフォーマー (ViT) の最終予測は最も有益なトークンのサブセットのみに基づいているという観察に触発され、私たちはトークンベースの枝刈りを通じて SSM ベースのビジョン モデルの効率を高めるという新しいステップを踏みました。
ただし、ViT 用に設計された既存のトークン プルーニング技術を直接適用すると、大規模な微調整を行ったとしても良好なパフォーマンスを実現できません。
この問題に対処するために、SSM の固有の計算特性を再検討し、単純なアプリケーションが連続したトークンの位置を混乱させることを発見しました。
この洞察は、SSM ベースのビジョン モデルに特化した、斬新で一般的なトークン プルーニング手法を設計する動機となります。
まず、パフォーマンスを向上させるために残りのトークンの近傍を安定させるために、枝刈りを意識した隠れ状態の調整方法を導入します。
さらに、詳細な分析に基づいて、トークンプルーニングをガイドするために、SSMモデルに適合したトークン重要度評価方法を提案します。
効率的な実装と実用的な高速化手法により、私たちの手法は実際の高速化をもたらします。
広範な実験により、私たちのアプローチがさまざまなタスク全体のパフォーマンスへの影響を最小限に抑えながら、大幅な計算量の削減を達成できることが実証されました。
特に、ImageNet では 81.7\% の精度を達成し、プルーニングされた PlainMamba-L3 の FLOP は 41.6\% 削減されました。
さらに、私たちの研究は、将来の研究のために SSM ベースのビジョン モデルの動作を理解するためのより深い洞察を提供します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google