Spatial-Aware Reduction Framework: Towards Efficient and Faithful Visual State Space Models
Abstract
Mamba demonstrates strong efficiency in modeling long visual sequences. However, when token reduction is applied to structurally enhanced Mamba variants, these models exhibit a severe performance collapse. We attribute this degradation to the spatially agnostic nature of existing reduction methods, which violate the two-dimensional structural premise required by the selective scanning mechanism. In this work, we propose STORM, a spatial-aware token reduction framework designed to maintain structural integrity throughout the compression process. STORM reformulates reduction into a structured operation on spatial units, enforcing localized constraints to maintain both grid topology and neighborhood coherence. As a plug-and-play module, STORM equips existing reduction pipelines with explicit spatial awareness without any training. Empirical results demonstrate that STORM achieves state-of-the-art pruning accuracy across diverse vision Mamba backbones under training-free settings. Notably, STORM delivers a substantial accuracy recovery on VMamba, outperforming prior methods by up to 63.3\% in top-1 accuracy. Meanwhile, STORM incurs only a 1.0\% accuracy drop on PlainMamba, achieving performance comparable to ViT.
Lay Summary
To make image recognition models faster, researchers often remove redundant pieces of an image. However, when applied to a new class of efficient models called visual state space models, existing methods accidentally break the image’s spatial layout, causing accuracy to collapse. We propose STORM, a simple plug‑in solution that preserves the image grid structure by reducing tokens along rows and columns separately, within small local neighborhoods. STORM works without any training. Experiments show that STORM recovers most of the lost accuracy. On a standard benchmark, previous methods drop below 50% accuracy, while STORM achieves over 80%. By keeping spatial integrity intact, STORM makes efficient vision models both fast and reliable—paving the way for lightweight AI in real‑world applications like mobile devices or autonomous systems.