Technical Summary: Spatial-Aware Reduction Framework (STORM)

1. Problem Statement

While State Space Models (SSMs) like Mamba have demonstrated superior efficiency in modeling long-range dependencies for vision tasks by reducing computational complexity from quadratic to linear, they face a critical bottleneck when integrated with token reduction techniques. Existing token reduction methods (e.g., pruning and merging), which are highly effective in Vision Transformers (ViTs), cause a severe performance collapse when applied to structurally enhanced Mamba variants like VMamba.

The authors attribute this degradation to the spatially agnostic nature of current reduction strategies. Unlike ViTs, which rely on global self-attention, Mamba architectures (specifically those utilizing the 2D Selective Scan, or SS2D, mechanism) depend on strict sequential scanning with recursive state dependencies. Standard reduction methods flatten the 2D token grid into a 1D sequence, perform reduction, and remap the tokens. This process disrupts the original spatial relationships and grid topology, violating the structural premises required for the SS2D mechanism to function correctly. Consequently, the causal state propagation is broken, leading to cascading information loss and catastrophic accuracy drops (e.g., a drop from ~83% to ~31% in VMamba-B when using ToMe).

2. Methodology: STORM

To address this, the authors propose STORM (Spatial-aware TOken Reduction), a plug-and-play framework designed to maintain structural integrity during compression without requiring additional training. STORM reframes token reduction as a structured operation on spatial primitives rather than a global operation on a flattened sequence.

Core Components

1. Structured Spatial Reduction (Dimension-Decomposed Workflow):
   Instead of processing a globally flattened sequence, STORM decomposes the reduction process into two decoupled stages aligned with the 2D grid:

   • Row-wise Reduction: The feature map is processed row by row. A base reduction operator (e.g., pruning or merging) is applied independently to each row.
   • Column-wise Reduction: The intermediate result is then processed column by column, applying the reduction operator independently to each column.
     This ensures the output remains a regular sub-grid, preserving the topological relationships required by the SS2D scanning paths.

2. Localized Window Mechanism:
   To prevent long-range interference that could disrupt local semantic coherence, STORM introduces a localized windowing constraint.

   • Within each row or column, the sequence is partitioned into non-overlapping windows.
   • Reduction decisions are confined strictly within these local neighborhoods.
   • This mechanism suppresses cross-region interference, ensuring that semantically unrelated tokens are not merged, thereby preserving fine-grained semantics essential for stable state evolution.

Theoretical Basis

The authors provide a theoretical analysis showing that global reduction methods cause unbounded spatial displacement between consecutive tokens in the scanning sequence, leading to error divergence. In contrast, STORM's dimension decomposition and window constraints strictly bound the spatial displacement of tokens. By limiting the maximum mapping displacement to the window size, the framework guarantees a tight upper bound on state propagation errors, ensuring stable performance even under aggressive compression.

3. Key Contributions

• Identification of Structural Integrity as Critical: The paper identifies that the failure of token reduction in vision Mamba is not merely due to token selection criteria but stems from the disruption of 2D spatial structure.
• Structured Reduction Paradigm: The authors propose a novel paradigm that operates on spatial primitives (rows and columns) rather than globally flattened tokens, effectively mitigating the performance collapse.
• Generalizable Framework: STORM is designed as a generic, training-free module that can be integrated with existing reduction pipelines (such as EViT and ToMe) to equip them with explicit spatial awareness.
• State-of-the-Art Performance: The framework establishes new performance benchmarks for vision Mamba reduction, delivering superior accuracy and throughput across diverse backbones.

4. Experimental Results

Experiments were conducted on ImageNet-1K (classification) and COCO 2017 (object detection and instance segmentation) using a training-free setting.

• ImageNet-1K Classification:

  • STORM significantly recovers accuracy on VMamba. For VMamba-B, applying ToMe alone drops accuracy to 35.7%, whereas STORM (ToMe) achieves 82.6% (only a 1.3% drop from the baseline).
  • On VMamba-S, STORM (ToMe) achieves 80.9% accuracy compared to 32.9% for standard ToMe, representing a 61.2% relative improvement in accuracy recovery.
  • STORM outperforms other Mamba-specific reduction methods (e.g., HSA, MTR, QuarterMap) by a large margin, with HSA dropping to 46.2% on VMamba-S while STORM reaches 82.0%.

• Object Detection (COCO):

  • STORM enables dramatic recovery on detection tasks. For VMamba-S, the baseline APb is 48.7. Standard ToMe collapses to 2.8, while STORM (ToMe) recovers to 44.2.
  • Similar improvements are observed on PlainMamba and LocalMamba backbones.

• Efficiency and Robustness:

  • Throughput: STORM achieves higher throughput than ToMe across all reduction ratios due to its parallel row/column operations, which avoid the scalability bottlenecks of global similarity computations.
  • Robustness: Even when using a random token selection strategy (without learned importance scores), STORM maintains robust performance (69.1% accuracy at 36% reduction), whereas ToMe drops to 11.7%. This highlights that the framework's structural preservation is the primary driver of success, not just the selection criterion.
  • Fine-tuning: Even with fine-tuning, STORM consistently outperforms baselines, recovering 78.4% accuracy on VMamba-T at a 0.32 reduction ratio compared to 32.7% for ToMe.

5. Significance and Claims

The paper claims that STORM represents a fundamental shift in how token reduction is approached for spatially sensitive architectures. By enforcing global layout regularity and local contextual fidelity, STORM resolves the incompatibility between token reduction and the chain-like state propagation of Mamba.

The authors emphasize that STORM is a training-free, plug-and-play solution that does not require retraining the backbone models. Its significance lies in enabling the practical deployment of vision Mamba models by allowing aggressive compression (up to 40% reduction) without the severe performance collapse previously observed. The work suggests that maintaining spatial consistency is more critical than merely retaining salient tokens for the efficacy of SS2D-based models.