SparseSSM: Efficient Selective Structured State Space Models Can Be Pruned in One-Shot
Abstract
State-space language models such as Mamba match Transformer quality while permitting linear complexity inference, yet still comprise billions of parameters that hinder deployment. While existing one-shot pruning methods are effective for generic linear and attention blocks, they are not designed with the overall Mamba architecture in mind and fail to account for the time-shared and discretized state-transition matrix at the heart of the selective state-space module (SSM). In this paper, we introduce SparseSSM, the first training-free pruning framework that extends the classic optimal brain surgeon (OBS) framework to state space architectures. Our layer-wise algorithm (i) derives an approximate second-order saliency score that aggregates Hessian-trace information across time steps, (ii) incorporates a component sensitivity analysis to guide feed-forward network (FFN) pruning, which also sheds light on where redundancy resides in mamba architecture, (iii) can be easily extended to semi-structured and structured sparsity, and generalized to other SSM-based architectures. Empirically, we prune 50% of SSM weights without fine-tuning and observe only limited performance degradation, achieving the current state-of-the-art one-shot pruning algorithm for Mamba-based LLMs.
Lay Summary
Large language models are powerful, but they are often expensive to run because they contain many parameters. Mamba, a recent alternative to Transformer-based language models, is designed to process long sequences efficiently, yet Mamba models can still be too large for practical deployment. Existing compression methods were mostly built for Transformers and do not account for how Mamba stores and updates information over time. We introduce SparseSSM, a method that removes less important parts of Mamba models in one-shot, without extra training. The key idea is to assess the importance of parameters in Mamba’s core modules and identify where redundancy lies within the Mamba architecture, then remove the parts that matter least while protecting more sensitive components. We test SparseSSM on Mamba models of multiple sizes and find that it can remove half of the state-space weights with little or no loss in performance. In structured settings, the compressed models also run faster and use less memory. This could make efficient long-sequence language models easier and cheaper to deploy on limited hardware.