Spatial-Aware Reduction Framework: Towards Efficient and Faithful Visual State Space Models

Vision Mamba models, which use state space models (SSMs) for efficient image processing, often suffer from severe performance degradation when standard token reduction techniques (like pruning or merging) are applied. The authors investigate why these models, particularly those using 2D Selective Scan (SS2D) mechanisms, fail under compression and whether this can be mitigated without retraining.

The authors identify that existing reduction methods are spatially agnostic, meaning they flatten 2D image tokens into a 1D sequence for processing. This disrupts the 2D grid topology required by the SS2D mechanism, causing a mismatch between the model's learned spatial dependencies and the compressed input. To solve this, the authors propose STORM, a plug-and-play framework that reformulates token reduction into a structured, two-stage operation: row-wise reduction followed by column-wise reduction. Additionally, STORM introduces a localized windowing mechanism that restricts reduction operations to small, coherent neighborhoods, preventing long-range interference that would otherwise distort the model's state propagation.

STORM effectively preserves the structural integrity of the feature map, allowing vision Mamba models to maintain high accuracy even under aggressive compression. In experiments on VMamba, STORM achieved a substantial accuracy recovery, outperforming prior methods by up to 63.3% in top-1 accuracy. On PlainMamba, the framework limited the accuracy drop to just 1.0%, achieving performance comparable to standard Vision Transformers. The authors provide a theoretical analysis demonstrating that by bounding the spatial displacement of tokens, STORM keeps the cumulative state error within a stable, manageable limit.