# Exploring Token Pruning in Vision State Space Models

arXiv: https://arxiv.org/abs/2409.18962 · HTML: https://mlsanity.com/paper/2409.18962/

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Placed in Token pruning and merging for ViTs, one of 12 directions in Efficient LLM inference and serving.
- Builds on 36 of its 50 references that are in the map.
- First posted on arXiv 27 Sep 2024.

## Placed in

- Efficient LLM inference and serving / [Token pruning and merging for ViTs](https://mlsanity.com/inference/token-pruning-and-merging-for-vits/)

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Mask R-CNN](https://mlsanity.com/paper/1703.06870/) (2017)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Unified Perceptual Parsing for Scene Understanding](https://mlsanity.com/paper/1807.10221/) (2018)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model](https://mlsanity.com/paper/2401.09417/) (2024)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [VMamba: Visual State Space Model](https://mlsanity.com/paper/2401.10166/) (2024)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Vision Transformer Adapter for Dense Predictions](https://mlsanity.com/paper/2205.08534/) (2022)
- [Learning Structured Sparsity in Deep Neural Networks](https://mlsanity.com/paper/1608.03665/) (2016)
- [Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration](https://mlsanity.com/paper/1811.00250/) (2018)
- [Soft Filter Pruning for Accelerating Deep Convolutional Neural Networks](https://mlsanity.com/paper/1808.06866/) (2018)
- [Long Range Language Modeling via Gated State Spaces](https://mlsanity.com/paper/2206.13947/) (2022)
- [NISP: Pruning Networks using Neuron Importance Score Propagation](https://mlsanity.com/paper/1711.05908/) (2017)
- [PlainMamba: Improving Non-Hierarchical Mamba in Visual Recognition](https://mlsanity.com/paper/2403.17695/) (2024)
- [IA-RED2: Interpretability-Aware Redundancy Reduction for Vision Transformers](https://mlsanity.com/paper/2106.12620/) (2021)
- [Selective Structured State-Spaces for Long-Form Video Understanding](https://mlsanity.com/paper/2303.14526/) (2023)
- [The Hidden Attention of Mamba Models](https://mlsanity.com/paper/2403.01590/) (2024)
- [Learning to Merge Tokens in Vision Transformers](https://mlsanity.com/paper/2202.12015/) (2022)
- [HeatViT: Hardware-Efficient Adaptive Token Pruning for Vision Transformers](https://mlsanity.com/paper/2211.08110/) (2022)

## Built on by (most cited first, 28 of 28)

- [Token Reduction Should Go Beyond Efficiency in Generative Models -- From Vision, Language to Multimodality](https://mlsanity.com/paper/2505.18227/) (2025)
- [VOTE: Vision-Language-Action Optimization with Trajectory Ensemble Voting](https://mlsanity.com/paper/2507.05116/) (2025)
- [TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training](https://mlsanity.com/paper/2501.04765/) (2025)
- [Rethinking Token Reduction for State Space Models](https://mlsanity.com/paper/2410.14725/) (2024)
- [Fast and Memory-Efficient Video Diffusion Using Streamlined Inference](https://mlsanity.com/paper/2411.01171/) (2024)
- [Structured Agent Distillation for Large Language Model](https://mlsanity.com/paper/2505.13820/) (2025)
- [RoRA: Efficient Fine-Tuning of LLM with Reliability Optimization for Rank Adaptation](https://mlsanity.com/paper/2501.04315/) (2025)
- [QuartDepth: Post-Training Quantization for Real-Time Depth Estimation on the Edge](https://mlsanity.com/paper/2503.16709/) (2025)
- [Vision Mamba in Remote Sensing: A Comprehensive Survey of Techniques, Applications and Outlook](https://mlsanity.com/paper/2505.00630/) (2025)
- [UniQL: Unified Quantization and Low-rank Compression for Adaptive Edge LLMs](https://mlsanity.com/paper/2512.03383/) (2025)
- [Faster Vision Mamba is Rebuilt in Minutes via Merged Token Re-training](https://mlsanity.com/paper/2412.12496/) (2024)
- [Representation Shift: Unifying Token Compression with FlashAttention](https://mlsanity.com/paper/2508.00367/) (2025)
- [Open-Source Multimodal Moxin Models with Moxin-VLM and Moxin-VLA](https://mlsanity.com/paper/2512.22208/) (2025)
- [Collaborative Compression for Large-Scale MoE Deployment on Edge](https://mlsanity.com/paper/2509.25689/) (2025)
- [TSLA: A Task-Specific Learning Adaptation for Semantic Segmentation on Autonomous Vehicles Platform](https://mlsanity.com/paper/2508.12279/) (2025)
- [Vision SmolMamba: Spike-Guided Token Pruning for Energy-Efficient Spiking State-Space Vision Models](https://mlsanity.com/paper/2604.25570/) (2026)
- [Fast Enough to Act: Spatio-Temporal Visual Token Merging for Low-Latency Robotic VLMs and VLAs](https://mlsanity.com/paper/2606.29350/) (2026)
- [Fast Vision Mamba: Pooling Spatial Dimensions for Accelerated Processing](https://mlsanity.com/paper/2502.00594/) (2025)
- [QuarterMap: Efficient Post-Training Token Pruning for Visual State Space Models](https://mlsanity.com/paper/2507.09514/) (2025)
- [Keyword Mamba: Spoken Keyword Spotting with State Space Models](https://mlsanity.com/paper/2508.07363/) (2025)
- [Towards Efficient Vision State Space Models via Token Merging](https://mlsanity.com/paper/2508.13599/) (2025)
- [PSTTS: A Plug-and-Play Token Selector for Efficient Event-based Spatio-temporal Representation Learning](https://mlsanity.com/paper/2509.22481/) (2025)
- [Sequential Token Merging: Revisiting Hidden States](https://mlsanity.com/paper/2509.22691/) (2025)
- [MambaScope: Coarse-to-Fine Scoping for Efficient Vision Mamba](https://mlsanity.com/paper/2512.00647/) (2025)
- [CLASP: Class-Adaptive Layer Fusion and Dual-Stage Pruning for Multimodal Large Language Models](https://mlsanity.com/paper/2604.12767/) (2026)
- [InterPruner: Interactive Structured Pruning via Taylor-Implicit Criterion and Language-Prior Modulator for Multimodal Object Detection](https://mlsanity.com/paper/2608.10724/) (2026)
- [LEAP: Learned Block-wise Evidence Retrieval for Long Audio-Video Perception](https://mlsanity.com/paper/2609.39938/) (2026)
- [From Patching to Pruning Visual Computation in Vision Language Models](https://mlsanity.com/paper/2610.03389/) (2026)

## Abstract (arXiv)

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.