Exploring Token Pruning in Vision State Space Models

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.

Paper

References (50)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer g8vG6/10 · confidence 4/52024-07-03

Summary

The paper introduces a pruning-aware hidden state alignment method, stabilizing token neighborhoods and maintaining model performance during token pruning. It also proposes an adapted token importance evaluation method tailored for SSM-based models, effectively guiding the pruning process. Extensive experiments demonstrate the method's efficacy, achieving significant computation reduction with minimal performance impact. Notably, the approach achieves 81.7% accuracy on ImageNet with a 41.6% reduction in FLOPs for pruned PlainMamba-L3, highlighting its practical acceleration benefits and effectiveness in maintaining high accuracy in vision tasks.

Strengths

- the proposed method ensures the stability of token neighborhoods, maintaining model performance even after pruning. - The importance evaluation is tailored specifically for SSM-based models, this evaluation method effectively guides the token pruning process. - The approach achieves significant computation reduction with minimal impact on performance, exemplified by 81.7% accuracy on ImageNet with a 41.6% reduction in FLOPs for pruned PlainMamba-L3. - The paper provides extensive visualization to further validate the effectiveness

Weaknesses

- The proposed method is limited to plain, non-hierarchical SSM-based models.

Questions

Can the proposed method generalize to the hierarchical variants like VMamba?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors adequately addressed the limitations

Authorsrebuttal2024-08-12

Dear Reviewer, Thank you very much for spending time reviewing our paper and acknowledging our contributions. Since the discussion will end very soon, we sincerely hope that you have found time to check our detailed response to your previous questions/comments. If you have any further questions, please feel free to let us know. We will try our best to reply to you before the discussion deadline. Thank you very much, Authors

Reviewer TrFe5/10 · confidence 3/52024-07-10

Summary

This manuscript aims to propose an effective pruning method for SSMs to achieve great trade-off of computation overhead and accuracy. The authors find that directly applying pruning strategies designed for transformer structure would greatly impair the performance of SSMs and give related analysis, that is SSMs take traversal path to consturct the interaction between different tokens and simple pruning would disrupt this connection. To address this challenge, the author first propose a strategy to build the alignment between the remaining tokens with the pruned tokens by keeping the indices and part calculation process of the pruned tokens. Based on the output of SSM layers, the author also propose a token selection strategy.

Strengths

* The authors give a clear presentation of the motivation and the proposed method. The structure of the manuscript is well-organized. * The motivation for designing specific pruning strategy for SSMs that take traversal interaction path is reasonable. The experiments of directly applying pruning method of EViT makes sense. * Experiment results show that the proposed method reduces comparable FLOPs with directly applying EViT pruning strategy yet achieves better performance.

Weaknesses

* The authors only compare the proposed method with EViT. However, it seems that EViT is proposed in 2022. The authors should consider comparing with more recent pruning methods. * Lack of clear explanation and analysis of why the proposed pruning-aware hidden state alignment strategy is helpful. From Eq.6 we can see that the hidden state of h_{q+1} do not utilize any information of the pruned tokens, while the indices of evolution matrix A are changed. This part is confusing and I hope the author could give more descriptions about why only changing the indices brings such remakable improvements. * The proposed Token Pruning based on Importance Evaluation does not make sense. Why the output of SSM can directly reflect the token importance? Just because it shares the same length with input? I recommend the authors could give move explanations.

Questions

* I have an additional explanation about why directly applying pruning strategy of EViT greatly impair the performance of SSMs. Because transformer takes non-local operation, all the tokens have knowledge of the overall context, so that removing some tokens would not cause huge information loss. However, for SSMs, the traversal interaction path makes them greatly rely on the continuous context and token pruning would lead to huge information loss. If space available, I am happy to discuss with the authors about that.

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors have discussed the potential limitations.

Reviewer M2Lk6/10 · confidence 4/52024-07-12

Summary

This paper introduces a token pruning method for vision SSMs to improve computational efficiency while maintaining performance. The authors identify that direct application of existing token pruning techniques designed for ViTs fails in SSMs due to disruption of sequential token positions. To address this, they propose a pruning-aware hidden state alignment method to stabilize the neighborhood of remaining tokens.

Strengths

- The proposed method is well-motivated, with a clear analysis of why existing token pruning techniques fail for SSMs. - Comprehensive experiments demonstrate efficiency gains across multiple tasks. - The structure and presentation of the paper are clear and well-organized.

Weaknesses

See the questions below.

Questions

- In Eq.6, how do you handle the accumulation of errors when computing hidden states for multiple consecutive pruned tokens? Does this lead to any instability in longer sequences? - How sensitive is the performance to the number of pruned tokens? Is there an upper limit to K_i before performance degrades significantly? - Is there any specific analysis about the pruning rate?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

None

Reviewer kiGk5/10 · confidence 4/52024-07-12

Summary

This paper explores the token pruning in vision state space models. It gives the observations that utilizing the token pruning techniques designed for ViTs leads to significant performance drop in SSMs. The main reason is that naive application disrupts the sequential token positions. To solve this, this paper introduces a pruning-aware hidden state alignment method and a token importance evaluation method to guide the token pruning while maintain better performance. Extensive experiments on ImageNet classification and object detection demonstrates the effectiveness of the proposed methods.

Strengths

The main strengths of the paper are: 1、 This paper gives insights about the failure of traditional token-pruning method adopted in ViTs on SSMs. This observation can better guide the token-pruning in SSMs in the future research. 2、 Extensive ablation studies on all relevant components. Apart from the quantitative evaluations, its visual results are also clear. 3、 The paper is well written and easy to read and follow.

Weaknesses

This paper gives interesting observations about the token-pruning in vision state space models. However, I have some concerns about the analysis and the token-pruning methods designed for SSMs. 1、 Fig. 1 illustrates the token computation difference between ViT and ViM. This introduces the large performance gap when adopting the traditional token-pruning techniques. Is there any more quantitative analysis result here to prove the impact of token computation patterns on performance of SSMs, such as randomly exchanging the order of adjacent tokens. 2、 What inspires the design of token importance in Eq(9). I can’t understand the reason of this design. Why adopting value 0 as the clipped threshold? Why does a larger S correspond to the most contextually relevant tokens? The result in Table 3 doesn’t indicate much superiority of the algorithm. In the traditional token-pruning method, the importance of a token should be represented by its’ impact on the loss value. 3、 This paper also introduces the efficient implementation for the SSM scan. The FLOPs and Params are provided to demonstrates the computations reduction. I think we care more about the actual running latency. This information is necessary for us to evaluate the practicality of the methods. The improvement of Throughput in table 4 is not as obvious as that of FLOPs. What caused this? 4、 Some detailed information is missing. Such as what proportions are reduced in different layers? For the SSM model, does the ratio of token reduction in different layers need to follow any principles? I think this can make the article better.

Questions

The observation of this paper is interesting. However, I have major concerns about the novelty of the designed methods. I hope the author can solve my questions in the weaknesses part.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations have been discussed in this paper. The efficiency is still limited by the baseline model architecture design.

Authorsrebuttal2024-08-12

Dear Reviewer, Thank you very much for acknowledging our observation and other contributions. Since the discussion will end very soon, we sincerely hope that you have found time to check our detailed response to your previous questions/comments. If you have any further questions, please feel free to let us know. We will try our best to reply to you before the discussion deadline. Thank you very much, Authors

Reviewer 6myn4/10 · confidence 4/52024-07-13

Summary

This paper presents a token pruning method for vision state space models. The goal of this paper is to expand the token pruning methods for ViTs to recent SSM-based vision backbones. The authors observed that token pruning will change the computational characteristics of SSMs and lead to significant accuracy drop. To solve this problem, a hidden state alignment method is designed to explicitly adjust the hidden states of SSMs after token pruning. Experiments are conducted on on ImageNet, COCO and ADE20k to show the effectiveness of the method.

Strengths

- The paper is well organized. The analyses in Section 3.2 clearly show the problem that the paper want to solve and help the readers understand the background easily. - The proposed alignment method and importance metric look reasonable, which are new since they are designed for SSM-based models.

Weaknesses

- There are still noticeable accuracy drop even after finetuning. Although there is no previous work on token pruning in SSM-based vision models, many methods are proposed for ViT token pruning and achieved nearly no accuracy drop when reducing around 20% FLOPs. But I find the proposed method will lead to around 0.5 accuracy drop. The results are not that impressive. - The actual speed-up is not reported. One key advantage of ViT token pruning compared to traditional channel pruning is hardware friend structures after pruning to easily lead to actual speed-up. If the proposed method can also achieve similar results, the method will be much more useful.

Questions

- How to determine the pruning protocol mentioned in line 262-265? It seems that the models are not uniformly partitioned. So it is better to show more details behind the partition strategy used in the experiments. - How about the results with different pruning ratios? The paper only demonstrate one pruning ratio for each model and the ratios are different for different models. Overall, I find there are some interesting observations and results presented in the paper. But the experimental studies look a bit weak and the results are not that impressive/useful. Therefore, I would rate this paper as "Borderline reject".

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The limitations have been discussed in Section 6.

Authorsrebuttal2024-08-12

Dear Reviewer, Thank you very much for acknowledging our observation and other contributions. Since the discussion will end very soon, we sincerely hope that you have found time to check our detailed response to your previous questions/comments. If you have any further questions, please feel free to let us know. We will try our best to reply to you before the discussion deadline. Thank you very much, Authors

Reviewer 6myn2024-08-13

Thanks for the detailed feedback. It is good to see that the method can also lead to actual speed-up after pruning. However, my concerns are still not fully addressed: - The actual speed-up of the proposed method seems less ideal than that of the ViT pruning method like DynamicViT. I am still concerned about the value of studying token pruning for SSM-based models. - EViT is a token pruning baseline published in 2022. Recent methods like AViT [r1] and STViT [r2] can achieve no accuracy drop after pruning. The results presented in the paper are relatively weak. [r1] AdaViT: Adaptive Tokens for Efficient Vision Transformer, CVPR 2022. [r2] Making Vision Transformers Efficient from A Token Sparsification View, CVPR 2023. Overall, I would keep my original rating.

Authorsrebuttal2024-08-14

**We sincerely appreciate the feedback from the reviewer. We address the newly raised questions below.** --- **Q1. The actual speed-up of the proposed method seems less ideal than that of the ViT pruning method like DynamicViT. I am still concerned about the value of studying token pruning for SSM-based models.** Token pruning is crucial for enhancing efficiency, especially with long sequences (please refer to our answer to **Q1** of Reviewer **M2Lk**). By reducing the number of tokens processed, we can significantly decrease computational cost and memory usage, thereby improving the scalability of SSMs for real-world applications. In this work, we introduce a general token pruning method specifically designed for SSM-based vision models. This is the **first successfully** token pruning work to **effectively** handle this problem on SSMs with **huge performance improvement**. Given the limited time (less than one day), we conducted an acceleration comparison between our method and DynamicViT under the same pruning ratio on SSMs. The results are as follows: |Model| Method | FLOPs(G)| Throughput | |-|-|-|-| | ViM-S |Ours| 3.60 | 1.27×| | ViM-S |DynamicViT| 3.56 | 1.30×| | PlainMamba-L3 |Ours|8.44| 1.43×| | PlainMamba-L3 |DynamicViT| 8.37 | 1.46×| As shown, there is minimal difference in acceleration under the same pruning ratio when compared to methods like DynamicViT on SSMs. However, our method provides **much higher** accuracy compared to ViT methods. For example, recent works such as ToMe [1] and LTMP [2], which we discussed in our response to **Reviewer TrFe**, highlight that our method achieves **significant performance improvements** over token pruning methods in ViT. The results are as follows: |Methods | FLOPs(G) | Top-1 Acc. (%)| |------------|-----|-----| | ViM-T |1.50| 76.1 | | ViM-T-ToMe | 1.28|71.6| | ViM-T-LTMP | 1.29|72.2| | ViM-T-EViT | 1.28|71.3| | ViM-T-Ours | 1.29 |75.1| --- **Q2. EViT is a token pruning baseline published in 2022. Recent methods like AViT [r1] and STViT [r2] can achieve no accuracy drop after pruning. The results presented in the paper are relatively weak.** While different token pruning methods in ViT may prune tokens without reducing accuracy, they all face the **same challenges** when applied to SSMs. In SSM, the relationships between tokens are **different** from those in ViT. SSMs handle sequential data, where the dynamic characteristics and contextual dependencies of each token determine its importance. Simple pruning might disrupt the sequence's integrity or the transmission of crucial information, leading to **significant** performance **degradation**. We would like to highlight that our work is the first successful application of token pruning to SSMs, resulting in substantial performance improvements. Due to the limited time (less than one day for applying AViT and STViT), we summarize the more recent work like ToMe [1] and LTMP [2] from our answer to **reviewer TrFe**. We will add these results into our revised version. We implemented ToMe[1] as well as LTMP [2] for our SSM-based models to provide a comprehensive comparison with state-of-the-art techniques. As demonstrated in the following table, our method can outperform all baselines with non-marginal improvements. |Methods | FLOPs(G) | Top-1 Acc. (%)| |------------|-----|-----| | ViM-T |1.50| 76.1 | | ViM-T-ToMe | 1.28|71.6| | ViM-T-LTMP | 1.29|72.2| | ViM-T-EViT | 1.28|71.3| | ViM-T-Ours | 1.29 |75.1| **Overall, we would like to emphases that regardless of the ViT token pruning method, if it does not address the specific characteristics of SSMs, it will not achieve the same level of effectiveness in SSMs. Therefore, it is unreasonable to compare the accuracy drop level on ViT with it on SSMs.** [1] Token Merging: Your ViT But Faster, ICLR 2023 [2] Learned Thresholds Token Merging and Pruning for Vision Transformers, TMLR 2023

Reviewer 6myn2024-08-14

Thanks for your new and detailed results. I agree that the proposed method can clearly improve existing methods that are not designed for SSM-based models, as I mentioned in "Strengths". However, I think the authors may misunderstand my core concern. I think there is still no clear evidence that SSM-based models are better than ViTs on visual tasks, and ViTs is still the first choice for core application scenarios of visual backbones including visual understanding (like detection, segmentation), and multimodal modeling (like CLIP, MLLMs). Therefore, in my opinion, only if achieving results comparable to ViTs' efficiency/performance can make a solid and impressive impact, simply beating some weak baselines may not fully demonstrate the value of studying this problem.

Authorsrebuttal2024-08-14

Thank you for your detailed feedback and for acknowledging the strengths of our work. We understand your core concern regarding the comparative performance of SSM-based models and ViTs on vision tasks, and we'd like to address this more explicitly. 1. **SSMs as an Emerging Foundation Model:** While it is true that ViTs have demonstrated strong performance across a variety of visual tasks, SSMs represent a new and rapidly evolving class of models. Given their emergent nature, it may be premature to conclude that ViTs are definitively superior to SSMs. In fact, SSMs have already shown **remarkable performance** in several vision tasks, indicating their potential to be **strong competitors** or even **successors** to ViTs as research in this area progresses. 2. **Our Baselines are Not Weak:** We want to clarify that the baselines we compare against are not weak; they represent the state of the art in the relevant categories. Our method **significantly outperforms** these **SOTA** methods, demonstrating the practical value and effectiveness of our approach within the current landscape of vision SSM models. 3. **Contributions Beyond Baseline Comparisons:** Our work goes beyond simply applying token pruning to a vision SSM model. We provide critical insights into the design of future vision SSM models, particularly regarding the handling of **token continuity** and **positional information**. These contributions are **foundational** and offer **valuable guidance** for the continued development of vision SSM models. It is important to recognize that early-stage work in emergent fields often serves as a **foundation stone**, even if it does not immediately surpass established benchmarks in ViTs-based models. The insights we offer pave the way for future innovations and improvements in the design of SSMs for vision tasks.

Reviewer TrFe2024-08-10

Thanks for the authors' effort and rebuttal. Most of my concerns have been addressed. I'm leaning to accepting this manuscript.

Authorsrebuttal2024-08-10

We sincerely thank the reviewer for recognizing that most concerns have been addressed and for leaning to accepting this manuscript. Given the enhancements made, we hope that our work now meets the criteria for a higher rating. We are truly grateful for your efforts in helping us refine our paper, and we look forward to your final assessment.

Reviewer M2Lk2024-08-10

I appreciate the author's response. While most of my concerns have been resolved, I'll keep my current rating for now and continue to pay attention to other reviews and ongoing discussions.

Authorsrebuttal2024-08-12

Dear Reviewer, Thank you for your prompt response and thoughtful consideration of our rebuttal. We appreciate your acknowledgment of our efforts to provide extensive results and explanations to address your concerns. We sincerely thank you for the constructive comments and positive rating. We will add the constructive suggestions in the final version of our paper. Thanks again for your time! Warm regards, Authors

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC