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

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.

## introduction we have seen an outstanding increase in the number of transformer-based models @xcite developed to tackle tasks from natural language processing (nlp) and other domains @xcite @xcite due to their effectiveness at modeling sequences. however, these models also present critical efficiency challenges. for example, the cost of training these models scales quadratically in the sequence length. in the generation stage, transformers, in their original form, require large caches to store the previously seen tokens. several variants of transformers have been proposed to address these efficiency challenges, but researchers have also explored alternative post-transformer architectures to address these limitations. structured state space models (ssms), e.g., * co-first authors. s4 @xcite , followed by selective state space models, e.g., mamba @xcite have been proposed as efficient alternatives that achieve training time with linear scaling in sequence length, and during generation, maintain constant state size. model compression methods, e.g., pruning and quantization, have been broadly explored and applied to transformer-based models. however, more must be done to explore compression in their structured state space counterparts. this paper explores the pruning of these alternative architectures, providing insights into potential opportunities to increase their efficiency without sacrificing accuracy. we discuss the following contributions: the rest of the paper is organized as follows: §2 provides details of the alternative architectures utilized in our study and popular strategies for element removal in large models. §3 describes methods to study network pruning in mamba and hybrid architectures. §4 presents the results of our experiments and ablation studies, and we offer concluding remarks in §5. a related work section is included in the appendix. mamba: selective state space models s4 and other structured ssms are linear time-invariant (lti), i.e., their parameters are fixed, limiting their effectiveness for sequence modeling. for instance, structured state space models fail in many contentand context-based reasoning tasks. these limitations have motivated the development of timevarying alternatives, e.g., mamba @xcite , which incorporate selection mechanisms and are suitable for solving tasks previously ssm generations failed. specifically, mamba's ssm module, s6, allows its parameters to depend on the input, thereby modifying the formulation from time-invariant to time-varying. a second improvement proposed in mamba compared to previous ssms is a hardware-aware algorithm that speeds up execution while reducing memory ios. furthermore, mamba-2 (dao and gu, 2024) improves the original mamba architecture by proposing state space duality (ssd), which improves its efficiency on hardware accelerators compared to s6. this improvement is achieved by changing the state matrix, a, which directly controls the latent state, h. a is modified from being structured as a diagonal matrix to a formulation that utilizes a scalar-times-identity structure. additionally, mamba-2 introduces the concept of heads in ssms inspired by how multi-head attention (mha) works and implementing a grouped-value attention (gva) head structure. overall, the mamba-2 architecture, with its ssd core component, allows for improved parallelism of the block's projections. mamba block mamba models comprise several blocks stacked after each other. figure 1 on the left illustrates a single mamba block. each block has the selective ssm mechanism (s6 for mamba-1 and ssd for mamba-2) at its core, placed within a larger structure that combines a gated multilayer perceptron (mlp), a convolution, and silu activation functions @xcite . for more details about selective structured state space models, we refer the reader to @xcite and @xcite . ## methodology due to the large sizes of current state-of-the-art sequence models, mamba-shedder requires an efficient strategy to identify structures that can be removed without significantly affecting the model's accuracy. we approach this problem using a training-free approach, in which the least essential elements are considered for removal. similar strategies have been explored in transformer-based large language models @xcite @xcite . however, to our knowledge, no study explores the removal of structures in selective structured state space models. mamba-shedder conducts structure removal of mamba models and their hybrid variants at different granularities. as illustrated in the left of figure 1 , in the case of models with only mamba blocks, we explore the iterative removal of entire mamba blocks ( §2.1), or their ssm subcomponents, either s6 or ssd modules depending on the version of mamba (figure 1 ). the proponents of the mamba architecture do not provide a rationale for the number of mamba blocks required to build robust models, opening an opportunity for mamba-shedder to investigate whether some components might be redundant and hence removed from the model with a minor impact in accuracy. in addition to these components, in the case of hybrid models that also contain transformer blocks (middle of figure 1 ), we also explore the removal of entire transformer blocks or their subblocks: multilayer perceptrons (mlp) modules and multihead attention (mha) modules. in hybrid models, mamba-shedder also explores the removal of structures at a finer granularity by targeting groups of channels in the mlp's linear layers, i.e., based on a channel group size, g, mamba-shedder explores the removal of ng channels, where n is the number of groups that could be removed based on their impact of the overall model performance. ## experiments we evaluate mamba-shedder and study the removal of structures from ssm-based models utilizing several open-source models and datasets. we analyze their absolute and relative drop in accuracy and quantify the inference speedup obtained by the pruned models. next, we discuss the resources utilized for our experiments and details of our setup and results. ## conclusion selective structure state space models have become an efficient alternative to transformer-based models. in this paper, we propose mamba-shedder and investigate structured pruning strategies to remove elements from mamba and hybrid models and reduce model size, accelerating inference. the results demonstrate that selective structured state space architectures have several redundancies that can be removed without significantly affecting the model's performance.