Research map / Architectures and efficiency / Attention and transformers
Exploring Token Pruning in Vision State Space Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a token pruning method for SSM-based vision models (ViM, PlainMamba) that keeps original scan positions via pruning-aware hidden state alignment and ranks tokens by clipped SSM outputs, achieving large FLOPs cuts with little accuracy loss.
Problem
Token pruning designed for ViTs breaks the sequential neighborhood structure of SSM scans in vision state space models, causing large accuracy drops even after fine-tuning.
Approach
Keep original token positions during the SSM scan: pruned tokens still advance the hidden state (state = A_bar * state, no input), and token importance is the channel-wise sum of clipped (ReLU) SSM outputs rather than ViT-style attention scores.
Key results
- PlainMamba-L3-prune reaches 81.7% ImageNet top-1 at 8.44G FLOPs (-41.4%) versus 82.3% dense, outperforming EViT-pruned PlainMamba-L3 (78.9%).
- ViM-S-prune achieves 78.8% top-1 at 3.60G FLOPs versus 80.5% dense, while naive EViT pruning gives 74.8%.
- Without the alignment, ViM-S pruning falls to 75.4% top-1; with alignment it recovers to 78.8% at similar throughput (1.27x).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use proposed token pruning instead of EViT pruning | Vision SSMs · ImageNet-1K val | top-1 accuracy: 71.3 -> 75.1; 74.8 -> 78.8; 75.0 -> 77.4; 78.3 -> 81.0; 78.9 -> 81.7 | Helped |
| Apply proposed token pruning | PlainMamba-L3 · ImageNet-1K val | FLOPs reduction: (-41.4\%) | Helped |
| Use proposed token pruning instead of EViT pruning | PlainMamba-L3 · COCO 2017 | AP^b: 44.2 -> 46.5 | Helped |
| Use proposed token pruning instead of EViT pruning | PlainMamba-L3 · COCO 2017 | AP^b_50: 66.4 -> 67.7 | Helped |
| Use proposed token pruning instead of EViT pruning | PlainMamba-L3 · COCO 2017 | AP^b_75: 49.7 -> 50.8 | Helped |
| Use proposed token pruning instead of EViT pruning | PlainMamba-L3 · COCO 2017 | AP^m: 39.5 -> 40.6 | Helped |
| Use proposed token pruning instead of EViT pruning | PlainMamba-L3 · COCO 2017 | AP^m_50: 62.8 -> 64.1 | Helped |
| Use proposed token pruning instead of EViT pruning | PlainMamba-L3 · COCO 2017 | AP^m_75: 42.7 -> 43.4 | Helped |
| Add pruning-aware hidden state alignment (ablation) | ViM-S pruning · ImageNet-1K val | top-1 accuracy: 75.4 -> 78.8 | Helped |
| Add pruning-aware hidden state alignment (ablation) | PlainMamba-L3 pruning · ImageNet-1K val | top-1 accuracy: 79.3 -> 81.7 | Helped |
| Replace clipping with no clipping in token importance (ablation) | ViM-S pruning · ImageNet-1K val | accuracy: 78.8 -> 77.4 | Hurt |
| Replace clipping with no clipping in token importance (ablation) | PlainMamba-L3 pruning · ImageNet-1K val | accuracy: 81.7 -> 80.5 | Hurt |
Surprising findings
- Naive ViT-style token pruning drops ViM-S zero-shot accuracy by over 68%, and fine-tuning still leaves a 5.7% gap.
- The SSM scan accounts for only about 10-20% of total computation, so most speedup comes from reducing tokens in other layers.
Limitations
- Efficiency is limited by the baseline model architecture design, as the authors note.
- Pruning ratios are fixed per model with hand-chosen prune layers; the paper does not study adaptive pruning schedules.
Experimental setup
Datasets: ImageNet-1K COCO 2017 ADE20K
Benchmarks: ImageNet-1K COCO 2017 ADE20K
Metrics: top-1 accuracy FLOPs throughput AP^b AP^m mIoU
Baselines: EViT ViM-T ViM-S PlainMamba-L1 PlainMamba-L2 PlainMamba-L3 ViT-Base DeiT-Small PVT Swin
Builds on: EViT Vim PlainMamba Mamba VMamba
Closest papers by meaning
- Balanced Token Pruning: Accelerating Vision Language Models Beyond Local Optimization (NeurIPS 2025 · similarity 0.92)
- ATP-LLaVA: Adaptive Token Pruning for Large Vision Language Models (CVPR 2025 · similarity 0.90)
- LearnPruner: Rethinking Attention-based Token Pruning in Vision Language Models (ICLR 2026 · similarity 0.90)
- TopV: Compatible Token Pruning with Inference Time Optimization for Fast and Low-Memory Multimodal Vision Language Model (CVPR 2025 · similarity 0.90)
- QuietPrune: Query-Guided Early Token Pruning for Vision-Language Models (CVPR 2026 · similarity 0.90)
- When and What to Prune? Stage-Aware Visual Token Pruning for Efficient VLA (NeurIPS 2026 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.