Exploring Token Pruning in Vision State Space Models

We sincerely appreciate the reviewer for recognizing the strengths of our papers and providing valuable feedback. We are happy to address the raised questions as below.

---

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

We would like to kindly point out that ViM has proved to be the first fast and strong baseline for vision state space models and its design has been widely adopted by later vision state space models. That’s why we choose the ViM and its variants as our base models.

Moreover, all the token reduction methods face similar difficulties when adapting to hierarchical model architecture due to their conflicts with the nature of hierarchical model architecture. However, with analysis and specific designs, the efficiency can still be achieved on hierarchical model architecture with our method. We discuss more design details in the following Response to Q1 to show how we can tackle this problem with our method.

---

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

Thanks for pointing this out, following the settings from previous token pruning work on Swin[1], we apply the sparsification operations at stage 3 of Vmamba which contributes most of the complexity. This is because the Layer number inside each stage of VMamba-T is [2,2,8,2] and VMamba-S is [2,2,15,2], with major computations in Stage 3.

The toke pruning process cannot transfer across stages due to the downsampling process between stages. Therefore, we need to perform padding to restore the sequence at the end of the stage. But significant speedup inside stages can be achieved.

[1] Dynamic Spatial Sparsification for Efficient Vision Transformers and Convolutional Neural Networks. TPAMI

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

We sincerely appreciate the reviewer for recognizing the strengths of our papers and providing valuable feedback. We are happy to address the raised questions as below.

---

W1. More comprehensive comparison.

We agree that a more comprehensive comparison would help the audience better understand the contribution of our work. We would also like to point out that EViT is still considered a strong baseline for token pruning methods. As for other token reduction methods, there have been recent advancements on Token merging as well as pruning + merge. Therefore, 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.

[1] Token Merging: Your ViT But Faster, ICLR 2023

[2] Learned Thresholds Token Merging and Pruning for Vision Transformers, TMLR 2023

---

W2. Proposed pruning-aware hidden state alignment strategy explanation.

As discussed in Section 3.2, for ViT, the hidden state of the pruned token is removed, i.e., $h^\prime_{(q_{i})+1} = 0$ if the token $x_{(q_{i})+1}$ is pruned, which leads to substantial accuracy drop as this strategy dropping hidden states within an SSM block disrupts the original token positions in the SSM scan. Consequently, tokens that were not previously adjacent become neighbors during the scan in different directions or paths, leading to a distorted scan functionality and a significant accuracy degradation. Especially considering that the tokens are actually image patches in visual tasks with semantic information, disrupting their positions during the scan brings great difficulties to understand their relationship and the overall semantics.

To address this issue, our method adopts $h^\prime_{(q_{i})+1} = \mathbf{\overline{A}} h^\prime_{(q_{i})} + \mathbf{\overline{B}} x_{(q_{i})+1} = \mathbf{\overline{A}} h^\prime_{(q_{i})}$ where the token $x_{(q_{i})+1}$ is pruned. It leads to Equation (6). As shown in Equation (6), if a token is pruned, we do not simply remove its corresponding hidden state during the scan as it leads to substantial accuracy degradation. Instead, its hidden state in the scan can be obtained by using the previous state with one step forward, i.e., $h^\prime_{(q_{i})+1} = \mathbf{\overline{A}} h^\prime_{(q_{i})}$. In this way, the hidden states in SSM corresponding to pruned tokens are aligned with that of the original unpruned tokens to maintain the sequential positions of the original tokens without disrupting their neighbours. For the question “our method does not utilize any information of the pruned tokens”, our method is aware of the pruned tokens and can maintain the sequential information of the hidden states even if the tokens are pruned.

We would like to clarify that in SSM models, the SSM scan keep doing a sequential computation by multiplying the previous hidden state with evolution matrix A and add the next hidden to the current state, Therefore, for the concern that SSM did not change the indices of evolution matrix A, the changing number on the head of evolution matrix A is the exponent of A. We maintain the position information by maintaining the sequential positions of the original tokens without disrupting their sequential relation of evolution matrix A.

---

W3. Explanation of our design of Importance Evaluation.

Please refer to Global rebuttal A2 for more detailed explanation.

Mamba utilizes implicit attention within its layers. It processes information through SSM layers, allowing tokens to interact and influence each other. By the time tokens reach the output of the SSM layers, they have undergone multiple rounds of implicit interactions and transformations. Mamba 2[1], has shown that the SSM block is equivalent to a form of linear attention. Therefore, the output contains the cumulative effect of these interactions, reflecting how each token has contributed to and been influenced by the overall context of the input. We then aggregate the clipped values across all channels of the output as our token importance score, as discussed in Line 228-231.

In Figure 3 of our paper, we provide a visualization of the attention patterns derived from the output, which helps to illustrate the implicit attention mechanism in mamba. Moreover, deriving token importance directly from the model itself without additional token selection layers or algorithms can be beneficial for specific hardware optimization.

[1] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality.

---

Q1. Insight discussion.

We would like to first thank the reviewer for such an insightful discussion. In “Reason for the failure” discussion of section 3.2 of our paper, we try to find out the reason why token pruning can work on ViT while failing on ViM. An overall explanation could be that the non-local operation of transformers is due to the quadratic operation of self-attention’s token wise multiplication. This enables the model to gain comprehensive knowledge of the overall context.

Moreover, our explanation on SSM aligns with the reviewer's thought on vision state space models that the SSM relies on sequential information (continuous context), directly removing the tokens could lead to huge information loss.

Our results in Table 1 show that plainmamba could achieve larger FLOPs reduction (-41.4%) while maintaining good performance (-0.6%). This could be because plainmamba chose a continuous 2D scanning which scans from 4 directions, providing more overall context and robustness compared with 1D scanning in ViM, which lead to higher token pruning ratios.

We sincerely appreciate the reviewer for recognizing the strengths of our papers and providing valuable feedback. We are happy to address the raised questions as below.

---

Q1. How to handle the consecutive pruned tokens and longer sequences?

Thank the reviewer for raising this valuable question. Here is a more detailed explanation about how our method handles the accumulation of errors when computing hidden states for multiple consecutive pruned tokens. Our method adopts $h^\prime_{(q_{i})+1}= \mathbf{\overline{A}} h^\prime_{(q_{i})} + \mathbf{\overline{B}} x_{(q_{i})+1} = \mathbf{\overline{A}} h^\prime_{(q_{i})}$ where the token $x_{(q_{i})+1}$ is pruned, which leads to Equation (6). As shown in Equation 6, if a token is pruned, we do not simply remove its corresponding hidden state during the scan as it leads to substantial accuracy degradation. Instead, its hidden state in the scan can be obtained by using the previous state with one step forward. In this way, the hidden states in SSM corresponding to pruned tokens are aligned with that of the original unpruned tokens to maintain the sequential positions of the original tokens without disrupting their neighbors. Based on the visualization results of Figure A1 in appendix, our method can prune consecutive tokens without additional designs while maintaining good performance.

In our paper, we provide Table 2 and Table A1 to show longer sequence results on object detection with PlainMamba. Where the input size is 1280×800 and 512×2048 respectively. Demonstrating our stability on longer sequences.

---

Q2. How sensitive is the performance to the number of pruned tokens?

We observe that the performance drops faster after accuracy reaching below 70%. Setting this 70% accuracy as a threshold, when input length is 197:

For ViM-T, there are 60 remaining tokens after our token pruning, with around 31% FLOPs reduction. For ViM-S, there are 36 remaining tokens, which is around 56% reduction FLOPs. We can observe that the larger the model, the more the tokens can be pruned. Also, the upper limit is also closely related to the pruning locations, thus a better pruning location may lead to more tokens to be pruned. We will include this analysis in the revision.

---

Q3. Is there any specific analysis about the pruning rate?

Regarding pruning ratio, we set the same pruning ratio (e.g. 0.8) for each location, with a progress pruning manner to prune additional tokens following the pruning ratio for each location. We discuss the pruning locations in Line 263-265. The FLOPs reduction can vary due to different pruning locations associated with the number of layers.

We managed to finish evaluating more results with different pruning ratios on ViM-T and ViM-S, we will include more results in the revision.

Table A. Classification results of different ratios on ImageNet-1K.

We sincerely appreciate the feedback from the reviewer. We address the raised questions as below.

---

W1. More quantitative analysis result here to prove the impact of token computation patterns on performance of SSMs.

We would like to thank the reviewer for this valuable suggestion, we conduct additional experiments to randomly shuffle the token positions. For each pruning location of ViM (not all layers), we use the shuffle() function from the random package to randomly exchange the order of tokens. We observe the following performance. On ViM-T, the accuracy dropped to 26.35% for zero shot, and 69.2%(-6.9) after finetune. On ViM-S, the accuracy dropped to 25.69% for zero shot, and 74.1% (-6.4) after finetune. The results indicate that the random token positions lead to much worse performance.

---

W2. Reason behind our design of token importance.

Currently, SSM is a novel model architecture and to our best knowledge, there are no previous works studying the importance of tokens in SSM models, especially for vision SSM models. Therefore, we try to tackle this problem by leveraging previous experience from Transformers. In Transformers, the softmax operation ensures that importance scores are always positive. We aimed to maintain a similar property in our approach. That’s why we choose general metrics like $\ell_1$ norm, $\ell_2$ norm and clip at 0 that could provide positive importance score values.

Mamba utilizes implicit attention within its layers. It processes information through SSM layers, allowing tokens to interact and influence each other. By the time tokens reach the output of the SSM layers, they have undergone multiple rounds of implicit interactions and transformations. Mamba 2[1], has shown that the SSM block is equivalent to a form of linear attention. Therefore, the output contains the cumulative effect of these interactions, reflecting how each token has contributed to and been influenced by the overall context of the input. We then aggregate the clipped values across all channels of the output as our token importance score, as discussed in Line 228-231.

Furthermore, the choice of 0 as the clipping threshold serves a role similar to ReLU activation:

1. It introduces non-linearity, which can help in capturing more complex relationships.
2. It improves the stability of backpropagation by preventing negative gradients.
3. It could introduce sparsity which we think could benefit the token pruning fine-tuning process. Our experimental results also show that the choice of 0 as the clipping threshold could yield fairly good results.

In Figure 3 of our paper, we provide a visualization of the attention patterns derived from the output, which helps to illustrate the implicit attention mechanism in mamba.

Moreover, deriving token importance directly from the model itself without additional token selection layers or algorithms can be beneficial for specific hardware optimization.

---

Q1. Questions about throughput in table 4.

Thank the reviewer for pointing this out, the FLOPs of ViM-S was a typo. We have double checked that the accuracy and throughput results in our paper are correct, sorry for the confusion. The "1.21G" typo corresponds to the FLOPS of ViM-T (ratio=0.7), which has been presented in table of Global rebuttal A1. We have reported all the results for better clarity. The updated table is shown below:

Note that the purpose of this table is to do an ablation study on the effectiveness of the alignment approach, using our token importance pruning method. From the results, we can see that our method could improve the throughput by 1.30× with around 30% FLOPS reduction.

---

Q2. Detailed information about pruning ratio and pruned layers.

Thank the reviewer for pointing out this missing part. The pruning location is determined through zero-shot experiments. This allowed us to identify the most effective pruning locations for different model architectures.

The varying number of layers in different models necessitated a flexible approach to partitioning. We found that a uniform partitioning strategy was suboptimal due to these architectural differences.

We would like to clarify that existing works on ViT pruning also have different partitions. For example, Table 3 in DynamicViT paper (TPAMI version) and Section 5 of SPViT paper demonstrate that different scales of ViT models use different pruning locations. While the exact pruning locations vary between models, we maintained a consistent approach by ensuring fixed intervals between pruning stages for each model. This also aligns with methodologies in existing works.

Regarding pruning ratio, we set the same pruning ratio (e.g. 0.8) for each location, with a progress pruning manner, pruning an additional $r$ tokens for each location. We also add additional results of different ratios (e.g. 0.7, 0.9) in A1.

We managed to finish evaluating more results with different pruning ratios on ViM-T and ViM-S, we will include more results in the revision.

Table A. Classification results of different ratios on ImageNet-1K.

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

We sincerely appreciate the feedback from the reviewer. We address the raised questions as below.

---

W1. Accuracy drop after finetuning.

We appreciate the detailed feedback and would like to address the concern regarding the observed accuracy drop. Our work represents a novel advancement in enhancing the efficiency of SSM-based vision models through tailored token pruning techniques. Unlike ViTs, where existing token pruning methods have been extensively researched and optimized, SSM-based models present unique challenges. As demonstrated in Figure 2 and Table 1 of our main paper, traditional token pruning methods for ViTs do not perform well on SSMs.

We would like to highlight that, in comparison to existing designs, our method consistently delivers better accuracy across various SSM-based vision models and scales. For example, our approach surpasses the state-of-the-art (SOTA) token pruning technique for ViTs, EViT, by 3.8% on ViM-T, 4.0% on ViM-S, 2.4% on PlainMamba-L1, 2.7% on PlainMamba-L2, and 2.8% on PlainMamba-L3. This significant performance improvement (over 2.4% across different models) has been recognized and acknowledged by other reviewers (Reviewer M2Lk, TrFe, g8vG). Our method demonstrates superior performance compared to the ViT-based approaches, highlighting the need for SSM-specific pruning techniques.

Regarding the 0.5% accuracy drop compared to dense computations on SSM-based vision models, one possible explanation is the inherent difference in computational complexity. SSM-scans achieve linear complexity, which poses greater challenges in attaining high computation reductions compared to the quadratic complexity of the attention mechanism in ViTs, which tends to have more redundant computations. This makes direct comparisons of FLOPs reductions between SSM and ViT models challenging.

On the other hand, the ViM has relatively small and efficient architecture. Smaller models generally have less redundancy, making it more challenging to achieve significant FLOPs reductions without impacting accuracy. Our experiments on the larger PlainMamba-L3 model can show a 40% FLOPs reduction with only a 0.6% accuracy drop. This result is better than the performance reported for the transformer model EViT-DeiT-B, which achieved a 35% FLOPs reduction with a 0.5% accuracy degradation (as reported in Table 8 of the EViT paper).

---

W2. Speedup performance.

We agree that actual speed-up is an important metric to demonstrate the effectiveness of token pruning schemes. Thus, we would like to clarify that we also include the speedup performance of our methods in Table 4 in the main paper. For instance, on PlainMamba-L3, our method with 8.44G FLOPs achieves 1.43$\times$ throughput accelerations than the dense method. Here is the table in our paper,

---

Q1. How to determine the pruning protocol mentioned in line 262-265.

Thank the reviewer for raising this valuable question. We will provide more details for comprehensive experimental results. Our partitioning strategy was determined through zero-shot experiments. This allowed us to identify the most effective pruning locations for different model architectures.

The varying number of layers in different models necessitated a flexible approach to partitioning. We found that a uniform partitioning strategy was suboptimal due to these architectural differences. Existing works on ViT pruning also have different partitions. For example, Table 3 in the DynamicViT [1] paper and Section 5 of SPViT [2] paper demonstrate that different scales of ViT models use different pruning locations.

While the exact pruning locations vary between models, we maintained a consistent approach by ensuring fixed intervals and same pruning rate between pruning stages for each model. This also aligns with methodologies in existing works. In the paper, we constantly set the pruning ratio to 0.8 for each location. The FLOPs reduction difference is due to the difference in model size related to the pruning location of layer index as stated in Line 262-265.

We also include more detailed information about Pruning ratio and pruned layers in Global rebuttal A1.

[1] Dynamic Spatial Sparsification for Efficient Vision Transformers and Convolutional Neural Networks. TPAMI

[2] SPViT: Enabling Faster Vision Transformers via Soft Token Pruning, ECCV 2022

---

Q2. Results with different pruning ratios.

In the paper, we constantly set the pruning ratio to 0.8 for each location. The FLOPs reduction difference is due to the difference in model size related to the pruning location of layer index as stated in Line 262-265. We managed to finish evaluating more results with different pruning ratios on ViM-T and ViM-S, we will include more results in the revision.

Table A. Classification results of different ratios on ImageNet-1K.

---

Overall, we take the first step toward accelerating vision SSM models with token-based pruning. We conduct a comprehensive analysis of SSM-based blocks to identify the failure reason, as well as provide more insights for the SSM scan mechanism in vision tasks, shedding lights on future research on SSM-based vision models.

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

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:

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:

---

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.

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

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.