AptMamba: enhancing vision Mamba efficiency by adaptively pruning tokens

• Main Author: Song, Shangye
• Other Authors: Jiang Xudong
• Format: Thesis-Master by Coursework
• Language: English
• Published: Nanyang Technological University 2025
• Subjects: Computer and Information Science; Deep learning
• Online Access: https://hdl.handle.net/10356/182039
• Institution: Nanyang Technological University

Mamba has recently emerged as a promising architecture for handling long se- quences in vision tasks. However, research on pruning methods specifically for vision Mamba remains unexplored. To address this gap, we aim to introduce to- ken pruning for vision Mamba to enhance efficiency, which faces several chal- lenges. First, without an attention mechanism, selecting tokens for pruning is less straightforward compared with vision transformers. Second, Mamba’s se- quential processing requires re-adjusting the token sequence after pruning to maintain calculation validity. To address these challenges, we propose to Adap- tively Prune Tokens for Mamba (AptMamba). Specifically, we design an Adap- tive Token Prediction (ATP) module to adaptively assess token importance, gen- erating binary decisions through Gumbel-Softmax. Pruned tokens are set to zero to minimize interference, and a multi-stage, progressive pruning strategy is ap- plied. To account for Mamba’s sequential scanning, we reorder tokens after pruning. Moreover, we present an auxiliary Pruning Reconstruction Decoder (PR-Decoder) during training to adaptively reconstruct pruned tokens, enhancing the representation capacity of the pruned Mamba without introducing extra in- ference costs. Additionally, we develop a Multi-Stage Pruning Loss (MSP-Loss) to adaptively confirm that pruned tokens contain minimal information, reducing the impact of token pruning. Extensive experiments show that our method can reduce the amount of computation by about 11.8% - 37.5% with 66% progres- sive token pruning, while maintaining accuracy impact within 1.1% - 3.5%.