StableMamba: Distillation-free Scaling of Large SSMs for Images and Videos
This addresses a major limitation for vision applications by enabling more scalable and robust SSMs, though it appears incremental as it builds on existing Mamba and attention methods.
The paper tackled the scalability issue of large state-space models (SSMs) for image and video tasks by proposing a Mamba-Attention interleaved architecture, which improved accuracy by up to +1.7 on benchmarks like ImageNet-1K without requiring knowledge distillation.
State-space models (SSMs), exemplified by S4, have introduced a novel context modeling method by integrating state-space techniques into deep learning. However, they struggle with global context modeling due to their data-independent matrices. The Mamba model addressed this with data-dependent variants via the S6 selective-scan algorithm, enhancing context modeling, especially for long sequences. However, Mamba-based architectures are difficult to scale with respect to the number of parameters, which is a major limitation for vision applications. This paper addresses the scalability issue of large SSMs for image classification and action recognition without requiring additional techniques like knowledge distillation. We analyze the distinct characteristics of Mamba-based and Attention-based models, proposing a Mamba-Attention interleaved architecture that enhances scalability, robustness, and performance. We demonstrate that the stable and efficient interleaved architecture resolves the scalability issue of Mamba-based architectures for images and videos and increases robustness to common artifacts like JPEG compression. Our thorough evaluation on the ImageNet-1K, Kinetics-400 and Something-Something-v2 benchmarks demonstrates that our approach improves the accuracy of state-of-the-art Mamba-based architectures by up to $+1.7$.