Mamba-Shedder: Post-Transformer Compression for Efficient Selective Structured State Space Models

2025-01-28

AI Summary

Selective State Space Models (SSMs), such as Mamba, suffer from high computational overhead and memory footprint during deployment, hindering their practical adoption. Method: This paper introduces Mamba-Shedder—the first hardware-aware, multi-granularity compression framework tailored for SSMs. Leveraging systematic structural sensitivity analysis, it integrates module- and parameter-level collaborative pruning, structured sparsification, inference graph rewriting, and hardware-coordinated optimization to eliminate redundancy at fine granularity. Contribution/Results: Under <0.5% accuracy degradation, Mamba-Shedder achieves up to 1.4× inference speedup, significantly reducing latency and GPU memory consumption. The implementation is open-sourced, establishing a reproducible compression paradigm for efficient sequence modeling in the post-Transformer era.

Technology Category

• Machine Learning: Hardware-aware ML
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Data Mining & Knowledge Management: Data Compression

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Large pre-trained models have achieved outstanding results in sequence modeling. The Transformer block and its attention mechanism have been the main drivers of the success of these models. Recently, alternative architectures, such as Selective Structured State Space Models (SSMs), have been proposed to address the inefficiencies of Transformers. This paper explores the compression of SSM-based models, particularly Mamba and its hybrids. We study the sensitivity of these models to the removal of selected components at different granularities to reduce the model size and computational overhead, thus improving their efficiency while maintaining accuracy. The proposed solutions, collectively referred to as Mamba-Shedder, achieve a speedup of up to 1.4x during inference, demonstrating that model efficiency can be improved by eliminating several redundancies with minimal impact on the overall model performance. The code is available at https://github.com/IntelLabs/Hardware-Aware-Automated-Machine-Learning.

Problem

• Pretrained Models
• Sequence State Models (SSMs)
• Efficiency Optimization

Innovation

• Mamba-Shedder
• SSM Model Optimization
• Efficient Sequence Data Processing