Venue: NAACL 2025

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

J. Pablo Munoz \({}^{1}\), Jinjie Yuan\({}^{2*}\), Nilesh Jain\({}^{1}\)

\({}^{1}\)Intel Labs, \({}^{2}\)Intel Corporation

{pablo.munoz, jinjie.yuan, nilesh.jain}@intel.com

Co-first authors.

###### 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.

## 1 Introduction

We have seen an outstanding increase in the number of Transformer-based models Vaswani et al. (2017) developed to tackle tasks from Natural Language Processing (NLP) and other domains Parmar et al. (2018); Dosovitskiy et al. (2021); Arnab et al. (2021); Gong et al. (2021) 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., S4 Gu et al. (2022), followed by _Selective state space models_, e.g., Mamba Gu and Dao (2023); Dao and Gu (2024) 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:

* A pruning solution, **Mamba-Shedder**, which targets structures in selective structured state space models, improving their computational and memory efficiency.
* Comprehensive experiments to determine the tolerance of SSM-based models to the removal of their structures.
* Insights on how the differences in the SSM building blocks and their interaction with Transformer blocks in hybrid models affect the trade-off between efficiency and accuracy.

The rest of the paper is organized as follows: SS2 provides details of the alternative architectures utilized in our study and popular strategies for element removal in large models. SS3 describes methods to study network pruning in Mamba and hybrid architectures. SS4 presents the results of our experiments and ablation studies, and we offer concluding remarks in SS5. A Related Work section is included in the Appendix.

## 2 Preliminaries

### State Space Models

State space models (SSMs) have a long history of modeling sequences and dynamic systems. Re

cently, _structured_ SSMs, e.g., S4 Gu et al. (2022), have been proposed as an alternative to Transformers because of their efficient capabilities for mapping input to output signals. When dealing with discrete sequences as in Natural Language Processing (NLP), the parameters \(A\), \(B\) and \(C\) of these models are discretized to transform an input sequence, \(x_{t}\), and hidden state, \(h_{t}\), to obtain the output sequence, \(y_{t}\). It can be formalized as:

\[h_{t}=\mbox{$A$}h_{t-1}+\mbox{$B$}x_{t},y_{t}=\mbox{$C$}^{\top}h_{t}. \tag{1}\]

Mamba: Selective State Space ModelsS4 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 content- and context-based reasoning tasks. These limitations have motivated the development of time-varying alternatives, e.g., Mamba Gu and Dao (2023), 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-2Dao 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 blockMemba 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 Elfwing et al. (2018).

For more details about selective structured state space models, we refer the reader to Gu and Dao (2023) and Dao and Gu (2024).

### Hybrid Models

Lately, new models achieve the best of both worlds (Transformers and Selective SSMs) by proposing architectures with both classes of blocks. Zamba Glorioso et al. (2024) is one example of such a hybrid model. It combines the strengths of Mamba's backbone and the efficiency of selective SSMs with a shared Transformer block that incorporates Transformers' powerful in-context learning capabilities. The _shared attention_ mechanism, in which two attention blocks are reused and interleaved in an ABAB pattern throughout the network, is a characteristic innovation of Zomba. This model also applies LoRA adapters Hu et al. (2022) to the shared MLP blocks, achieving specialization when interacting with the affected layers, memory efficiency, and faster inference with reduced computational overhead.

Another example of a hybrid model is Hymba Dong et al. (2024). This model takes a different approach than Zamba, proposing an entirely new hybrid-head module, illustrated in Figure 1 on the right, in which the SSM and Attention mechanisms contribute in parallel to the sequence modeling. Additionally, Hymba benefits from group query attention, cross-layer KV cache sharing, and learnable meta-tokens, resulting in higher throughput, reduced memory requirements, and competitive performance compared to models of similar size.

### Model Pruning

A popular model compression technique, _pruning_LeCun et al. (1989), has been effectively used to reduce the size of deep learning models and improve their efficiency. Network pruning operates at two levels: (1) _Unstructured pruning_ identifies the importance of individual weights that can be masked to minimize their impact on overall model behavior. At a different level, (2) _structured pruning_ focuses on removing more significant structural components of the model, such as whole Transformer blocks Men et al. (2024), or reducing the granularity to target subcomponents of these layers Zhong et al. (2024). Other dimensions for pruning include groups of channels in the Transformer's MLPs or heads from the MHA layer Munoz et al. (2025). In this paper, the focus is solely on structured pruning applied to Mamba-based models. Next, we 

discuss Mamba-Shedder's methodology to study redundancies in Mamba and hybrid models.

## 3 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 Ashkboos et al. (2024); Men et al. (2024); Zhong et al. (2024). 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 (SS2.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.

```
1:Set of blocks/modules \({\cal M}\) from a model \(m\), Calibration dataset \({\cal C}\), Metric \(\phi\), Target pruning steps \(t\).
2:Pruned model \(m^{*}\)
3:for\(k\gets 1\)to\(t\)do
4:for all\(M_{i}\in{\cal M}\)do
5:\(S_{i}\leftarrow\text{Importance}(M_{i},m,{\cal C},\phi)\)
6:endfor
7:\(M_{\min}\leftarrow\arg\min_{M_{i}\in{\cal M}}S_{i}\)
8:\({\cal M}\leftarrow{\cal M}\setminus\left\{M_{min}\right\}\)\(\triangleright\) Block/Module Pruning
9:endfor
10:return\(m^{*}\)with the remaining blocks/modules in \({\cal M}\)
```

**Algorithm 1** Block / Module Pruning

Algorithm 1 details the procedure to remove entire structures, e.g., Mamba or Transformer blocks, MLPs, MMA, or SSM modules. Given a set \({\cal M}\) of structures selected for potential removal, a proxy data set \(C\) and a metric \(\phi\) are used to measure the importance of an individual structure and the impact of removing it from the model Zhong et al. (2024). In addition to entire structures, Mamba-Shedder follows the same logic to remove channel

Figure 1: Overview of Mamba-Shedder. This figure illustrates the pruning strategy for three types of Mamba-based models. The first type includes Mamba models such as Mamba-1 Gu and Dao (2023), Mamba-2 Dao and Gu (2024), and Falcon-Mamba Zuo et al. (2024). The second type comprises Mamba + Transformers architectures, including Zamba Glorioso et al. (2024). The third type is Hypha Dong et al. (2024), a novel architecture with hybrid heads. Red dashed lines indicate potential removal. In Transformers, channel pruning can also be applied to MLP block (width pruning).

 

groups as detailed in Algorithm 2.

```
0: Set of MLP blocks \({\cal M}_{\text{MLP}}\) from a model \(m\), Calibration dataset \({\cal C}\), Metric \(\phi\), Target pruning steps \(t\), MLP channel group size \(g\).
0: Pruned model \(m^{*}\)
1:for\(k\gets 1\) to \(t\)do
2:forall\(M_{i}\in{\cal M}_{\text{MLP}}\)do
3:\(S_{i}\leftarrow\text{Importance}(M_{i}[:,:-g],m,{\cal C},\phi)\)
4:endfor
5:\(M_{\text{mn}}=\arg\min_{M_{i}\in{\cal M}_{\text{MLP}}}S_{i}\)
6:\(M_{\text{mn}}=M_{\text{min}}[:,:-g]\)\(\triangleright\) Channel Pruning
7:endfor
8:return\(m^{*}\) with the altered MLP blocks in \({\cal M}\)
```

**Algorithm 2** MLP Channel Pruning

Depending on the pruning objective, Mamba-Shedder might treat these pruning targets in isolation, but Section 4 also presents the results of configurations in which Mamba-Shedder sequentially prunes larger structures (e.g., Mamba blocks) and, at a later stage, smaller components, e.g., SSM modules in the remaining Mamba blocks. Future work will explore larger search spaces with more complex configurations of candidate structures for removal. For example, the importance of Mamba blocks and their SSM modules can be assessed in the same pruning iteration.

## 4 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.

### Models

Our experiments employed the following pretrained Mamba and hybrid models: **Mamba-2.8b**Gu and Dao (2023), consists of 64 S6 blocks1. **Mamba2-2.7b**Dao and Gu (2024), consists of 64 SSD blocks 2. Both Mamba models were trained on 300B tokens on the Pile dataset Gao et al. (2020). For our choice of a hybrid model, we explored **Zamboa2-2.7B**Glorioso et al. (2024)3. It has 54

Figure 2: Pruning Mamba blocks. _Avg. Accuracy_ indicates the average accuracy for seven tasks. The model composed of Mamba 1 blocks (left) can tolerate the removal of entire blocks without significantly increasing its perplexity or decreasing accuracy compared to Mamba-2 and Zamba-2. In all three models, removing each Mamba block reduces 0.04B parameters from the model. These are _training-free_ results, and drops in accuracy can be reduced by a subsequent fine-tuning stage (§4.5).

 

layers, including 45 single Mamba-2 Blocks and 9 hybrid layers composed of both Mamba-2 Blocks and Transformer Blocks. Zamba-2 was trained on 3T tokens from open web datasets, including Zyda (Tokpanov et al., 2024), and subsequently annealed with 100B additional tokens. The aforementioned models are all of the same size and can be compared directly. For Mamba models of different sizes, we also explored **Falcon-Mamba-7B**Zuo et al. (2024)4, which is based on the Mamba-1 architecture and is the best-performing Mamba model at this scale in the literature, as well as **Hymba-1.5B-Base**Dong et al. (2024)5, which features a hybrid architecture incorporating both Mamba and Attention heads.

Footnote 4: [https://huggingface.co/tiuae/falcon-mamba-7b](https://huggingface.co/tiuae/falcon-mamba-7b)

Footnote 5: [https://huggingface.co/nvidia/Hymba-1.5B-Base](https://huggingface.co/nvidia/Hymba-1.5B-Base)

### Datasets

Following the language modeling evaluation of Mamba Gu and Dao (2023); Dao and Gu (2024), we utilize _lm-eval-harness_Gao et al. (2023) to assess the zero-shot performance, which includes measuring perplexity on Lambada Paperno et al. (2016), and accuracy on the following downstream tasks: HellaSwagzellers et al. (2019), Physical Interaction Question Answering (PIQA) Bisk et al. (2020), AI2 Reasoning Challenges Arc-e, Arc-c) Clark et al. (2018), Large-scale Winograd Schema Challenge WinoGrande) Sakaguchi et al. (2021), and the Open Book Question Answering Mihaylov et al. (2018) dataset.

Regarding the calibration dataset, we follow BlockPruner Zhong et al. (2024) in using the Alpaca dataset 6 as the calibration dataset and employ perplexivity as the metric for calculating importance scores. All the hyperparameters used in our experiments are detailed in the Appendix.

Footnote 6: [https://github.com/tatsu-lab/stanford_alpacas](https://github.com/tatsu-lab/stanford_alpacas)

### Results

#### 4.3.1 Pruning Target: Mamba Block

This section explores the impact of pruning Mamba blocks on model performance. Figure 2 and Table 1 present the results of applying Mamba-Shedder to Mamba-2.8B, Mamba-2.7B, and Zamba2-2.7B

\begin{table}
\begin{tabular}{l l c c c c c c c c c} \hline \hline \multirow{2}{*}{**Model**} & \multirow{2}{*}{**Method**} & \multicolumn{1}{c}{\multirow{2}{*}{**Num. of**}} & \multirow{2}{*}{**Lambda**} & \multirow{2}{*}{**HelaS**} & \multirow{2}{*}{**PIQA**} & \multirow{2}{*}{**ARC+**} & \multirow{2}{*}{**ARC-c**} & \multirow{2}{*}{**WinoG**} & \multirow{2}{*}{**OBQA**} & \multirow{2}{*}{**Average**} \\  & & & & PPL (↓) & & & & & & & \\ \hline \multirow{4}{*}{**Mamba-2.8B**} & Dense & 0 / 64 & 4.23 & 69.2 & 66.1 & 75.2 & 69.7 & 36.3 & 63.5 & 39.6 & 59.9 \\  & 16 / 64 & **9.23,540** & 55.2 & 52.1 & 68.1 & 57.8 & 28.4 & 55.6 & 31.6 & **49.8,61** \\  & SSM Pruning & 20 / 64 & **10.10,857** & 57.1 & 48.2 & 65.5 & 50.9 & 25.9 & 56.0 & 29.4 & **47.6,123** \\  & 24 / 64 & **22.55,183** & 44.4 & 43.2 & 64.4 & 47.4 & 25.8 & 53.6 & 29.8 & **44.1,158** \\ \hline \multirow{4}{*}{**Mamba-2.7B**} & Dense & 0 / 64 & 4.10 & 69.7 & 66.6 & 76.4 & 69.6 & 36.4 & 64.0 & 38.8 & 60.2 \\  & 16 / 64 & **4.26,146** & 66.9 & 66.1 & 76.4 & 68.6 & 37.2 & 64.0 & 39.2 & **59.8,4** \\  & SSM Pruning & 20 / 64 & **5.89,179** & 59.8 & 66.0 & 76.1 & 68.9 & 36.7 & 63.6 & 39.2 & **58.6,148** \\  & 24 / 64 & **14.95,185** & 43.4 & 65.8 & 74.8 & 67.1 & 36.6 & 62.9 & 38.0 & **55.5,47** \\ \hline \multirow{4}{*}{**Zambia-2.7B**} & Dense & 0 / 54 & 4.01 & 69.7 & 77.0 & 79.8 & 77.5 & 48.5 & 72.1 & 45.8 & 67.2 \\  & 16 / 54 & **4.14,0.13** & 69.2 & 75.8 & 79.2 & 75.8 & 46.5 & 72.2 & 45.8 & **66.4,0.8** \\ \cline{1-1}  & SSM Pruning & 20 / 54 & **5.07,166** & 64.2 & 75.8 & 79.3 & 75.5 & 46.2 & 73.2 & 46.0 & **65.7,15** \\ \cline{1-1}  & 24 / 54 & **5.46,145** & 62.3 & 74.7 & 79.0 & 75.4 & 44.3 & 70.9 & 46.4 & **64.7,2.8** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Detailed results of Mamba-Shedder with _training-free_ SSM pruning. The remaining tasks represent their respective accuracy. Here, we do not consider the pruning ratio, as the number of SSM’s parameter weights is small. Its benefit is the reduction of computational overhead. Underlined numbers indicate the smallest gap with Dense under the same level of pruning.

Figure 3: Pruning SSM (S6 and SSD modules). Mamba-2.8B and Mamba2-2.7B have 64 SSM modules, while Zamba2-2.7B has 54 SSM modules. _Avg. Accuracy_ is for the seven tasks evaluated.

 

models with a focus on removing redundant entire Mamba blocks. The model that utilizes the first version of Mamba blocks (S6) appears to tolerate a higher number of removed blocks without significantly affecting its performance. Specifically, the Mamba-2.8B model demonstrates robustness, with its perplexity (PPL) increasing from 4.23 to 7.51 and average accuracy dropping from 59.9 to 53.8 when the pruning ratio reaches 20.86%. In contrast, the Mamba2-2.7B and Zamba2-2.7B models exhibit more significant performance degradation, although they performed better before pruning (Dense). The poorer pruning performance of Zamba2-2.7B may be attributed to the pruning of Mamba blocks disrupting a certain balance within the hybrid layers. Overall, the effects of Mamba block pruning vary across different models, depending on the model architecture and the characteristics of the pre-training stage.

#### 4.3.2 Pruning Target: SSM Module

In this section, we delve into assessing the impact of pruning only the SSM modules within Mamba blocks on the performance of various models, as illustrated in Table 2 and Figure 3. When using the same target in Mamba-2.8B, we observe that further pruning SSMs results in a noticeable increase in perplexity, soaring to 22.55 and decreasing average accuracy to 44.1. This result indicates a significant sensitivity to SSM pruning for Mamba-1, where performance degradation is pronounced even at moderate pruning levels. Conversely, Mamba2-2.7B and Zambia2-2.7B exhibit remarkable resilience to SSM pruning. Even with 24 SSMs pruned, the model maintains a relatively stable performance. This robustness suggests that Mamba-2 blocks can tolerate higher SSM module pruning, potentially due to Mamba-2's optimizations or different training strategies with Mamba-1. The Zamba2-2.7B model, with the hybrid architecture, outperforms both Mamba-1 and Mamba-2. Pruning 12 out of its 54 SSMs results in a negligible PPL increase from 4.01 to 4.02, while the average accuracy slightly decreases from 67.2% to 67.0%. The hybrid nature of Zamba2-2.7B may contribute to its ability to maintain performance despite SSM pruning. Overall, these findings underscore the importance of model architecture and training strategies in determining the impact of SSM pruning. They offer valuable insights for optimizing model efficiency without compromising performance.

4.3.3 Pruning Target: Finer-grained removal of Mamba and Transformer blocks, and their subcomponents

Table 3 presents the results of pruning various components of the Zamba2-2.7B model, including combinations of Mamba-2 blocks, entire Transformer blocks, and their subcomponents, i.e., MHA blocks, MLP blocks, MLP channels, and SSM modules. We design four search spaces to study the effectiveness of different granularities and their combinations. "&" indicates that the pruning targets are considered together in the same pruning step, while "+" signifies the distinction between pruning stages, with pruning occurring sequentially:

Mamba Block & Transformer Block PruningThis experiment involves pruning the entire Mamba-2 blocks and Transformer blocks.

\begin{table}
\begin{tabular}{l c c c c c c c c c c} \hline \hline
**Pruning Target** & \multicolumn{2}{c}{**Ratio**} & \multicolumn{2}{c}{**Additional**} & \multicolumn{2}{c}{**Lambda**} & \multicolumn{2}{c}{**Hella’s PIQ/A ARC-c WC-wino OBOQA**} & \multicolumn{2}{c}{**Avg.**} \\  & \multicolumn{2}{c}{**(Block, Width)**} & \multicolumn{2}{c}{**Puned SSMs**} & \multicolumn{2}{c}{**PPL (L)**} & & & & & & & \\ \hline Mamba Block & 0\% & 0 / 54 & 4.01 & 69.7 & 77.0 & 79.8 & 77.5 & 48.5 & 72.1 & 45.8 & 67.2 \\ \hline Mamba Block \& Transformer Block & 10.40\% & 0 / 54 & 9.18,5\_17 & 53.5 & 67.3 & 76.3 & 63.5 & 37.8 & 64.3 & 40.6 & 57.6\_abs_ \\ Mamba Block \& MLP \& MHz & 10.33\% & 0 / 54 & **5.01,10** & **65.6** & 73.6 & 78.5 & 75.3 & 43.8 & 69.3 & 45.2 & 64.5\_2,7 \\ Mamba Block \& MLP \& MHz + MLP Channel & 10.27\% & 0 / 54 & 5.45,i4.4 & 63.4 & **74.9** & **80.1** & **79.0** & **49.7** & **70.9** & 46.0 & **66.34,9** \\ Mamba Block \& MLP \& MHz + MLP Channel + SSM & 10.27\% & **18 / 54** & 5.18,11.7 & 63.4 & 73.9 & 80.0 & **79.0** & 48.7 & 69.5 & **46.6** & 65.9\_abs_ \\ \hline Mamba Block \& Transformer Block & 15.89\% & 0 / 54 & 10.38,4\_97 & 51.4 & 65.6 & 74.0 & 61.7 & 37.7 & 63.5 & 39.6 & 56.2\_1,0 \\ Mamba Block \& MLP \& MHz & 15.54\% & 0 / 54 & 10.64,6.6j & 49.3 & 69.2 & 76.9 & 66.1 & 38.1 & 66.0 & 41.8 & 58.2\_abs_ \\ Mamba Block \& MLP \& MHz + MLP Channel & 15.48\% & 0 / 54 & **7.39,38** & **57.6** & **70.0** & **78.5** & **74.5** & **43.9** & 67.5 & **43.8** & **62.34,9** \\ Mamba Block \& MLP \& MHz + MLP Channel + SSM & 15.48\% & **18 / 54** & 7.43,3¡ 

Mamba Block & MLP & MHA PruningThis experiment decomposes the transformer block into sub-blocks, pruning Mamba-2 blocks as well as MHA and MLP.

**Mamba Block & MLP & MHA + MLP Channel Pruning** This experiment prunes the Mamba-2 blocks, MHA, and MLP at the first stage and further prunes the MLP channels at the next stage.

**Mamba Block & MLP & MHA + MLP Channel Pruning + SSM** Add additional SSM pruning following the previous solution.

The results indicate that pruning Mamba blocks and Transformer blocks alone leads to significant performance degradation. However, **more granular pruning strategies show a more favorable trade-off between pruning ratio and performance**. Specifically, pruning Mamba blocks, MLP, MHA (single stage), and MLP channels subsequently performs the best. Inspired by the SSM pruning of Mamba-2 in Section 4.3.2, we further add SSM pruning to the third strategy, and the results show that removing around 18 SSMs can maintain accuracy performance while reducing computational overhead. An interesting finding is that pruning SSMs can even lower PPL; for instance, at a 10% pruning ratio, PPL decreases from 5.45 to 5.18, suggesting that some SSM modules are redundant after the second pruning stage. Overall, these findings indicate that multi-granularity pruning methods, particularly those including MLP channels and SSM modules, can effectively reduce the complexity of hybrid Mamba models while maintaining a higher level of performance.

3.4 Pruning Mamba Models of Other Sizes HymbaTable 4 shows the results of Mamba-Shedder with training-free Hypha Block pruning for Hypha-1.5B-Base. The dense configuration achieves an average accuracy of 63.8, which decreases as more blocks are pruned, dropping to 60.5 when 8 blocks are pruned, indicating a general decline in performance across benchmarks. Further analysis of inference acceleration and recovery tuning experiments for Hypha-1.5B-Base will be discussed in the subsequent sections.

Falcon-LambdaWhile the previous sections focused on exploring the pruning of Mamba models with sizes around 2.7B or 2.8B, we also investigated the impact of Mamba-Shededer on a larger-scale Mamba model, specifically Falcon-Mamba-7B (Table 5). Pruning SSM modules in the Falcon-Mamboa-7B model shows better tolerance in terms of perplexity, suggesting that SSM pruning is more effective in maintaining lower perplexity. Regarding average accuracy, pruning entire Mamba blocks is more beneficial.

Additionally, it is important to note that pruning entire Mamba blocks yields more significant computational benefits than SSM pruning, suggesting that while SSM pruning is advantageous for maintaining perplexity, pruning Mamba blocks offers a better trade-off between computational efficiency and accuracy. The choice of pruning strategy should be guided by the specific performance metric of interest and the desired balance between computational efficiency and model accuracy.

None of the above results have undergone fine-tuning to improve the performance of the pruned models. As in other works, the drop in the accuracy performance of pruned models can be recovered by fine-tuning, which will be incorporated in SS4.5.

### Inference Acceleration

Through the above analysis, we have gained a good understanding and insights into the impact of Mamba-Shedder's structured pruning on model accuracy and perplexity performance. In this section, we discuss the impact of pruning on inference acceleration. All the following tests were conducted on a single Tesla V100 32GB GPU.

Mamba-1When removing entire Mamba blocks, as shown in Table 6, Mamba-Shedder speeds up the decoding stage up to 1.29x when removing 14 blocks, and 1.13x when removing only 7 blocks, which highlights the potential of Mamba-Shedder to optimize computational efficiency in Mamba models. The user's decision on how aggressively to prune will impact the average accuracy or the perplexity as observed in Table 1.

Mamba-2As detailed in Table 7, removing 24 SSM modules (44% of the total number of modules) results in up to a 1.20x speedup in the prefill stage and a 1.18x speedup in the decoding stage during of inference. A more conservative pruning ratio achieves 1.11x speedup when removing 16 SSM modules. Based on previous observations, the impact on performance metrics is minimal (0.4% for accuracy and 0.16 for PPL). These results underscore the effectiveness of SSM pruning in enhancing computational efficiency while barely affecting model performance, making it a viable strategy for optimizing Mamba models 

Zamba-2.7B, with _training-free_ MAMba block and SSM pruning for Falcon-Lambda-7B.

Zamba-2We observe significant acceleration on inference after multiple granularities pruning of Zamba-2 (Table 8). Specifically, pruning Mamba blocks, MLP, and MHA blocks along with MLP channels results in a 1.34x speedup in the decoding stage. When SSM prun

\begin{table}
\begin{tabular}{l l 

speedup increases to 1.39x, indicating that a comprehensive pruning strategy that includes multiple components can significantly enhance inference speed while maximizing the preservation of model performance.

HymbaBlock pruning of Hymba-1.5B-Base demonstrates notable improvements in inference speed (Table 9). By removing 7 out of 64 Hymba blocks, Mamba-Shedder achieves a 1.15x speedup in the prefill stage and a 1.24x speedup in the decoding stage, suggesting that significant computational efficiency gains can be realized even with a relatively modest pruning ratio. The results highlight the potential of Mamba-Shedder to optimize the performance of Hymba models, making them more efficient for real-time applications without substantial sacrifices in model accuracy.

### Recovery Tuning of the Pruned Model

Following related work Ma et al. (2023); Zhong et al. (2024), we performed post-training on the Mamba-Shedder compressed model using the cleaned version of Alpaca. The results summarized in Tables 10, 11, and 12 demonstrate substantial performance gains after just two epochs of recovery tuning. For instance, the Mamba-Shedder model obtained by removing Mamba Blocks & MLPs & MHAs + MLP Channels + SSM in Zamba-2 (Table 3), initially exhibits a perplexity of 5.18 and an average accuracy of 65.9 when 18 out of 54 SSMs are pruned. However, after recovery tuning, it achieves a significantly reduced PPL of 4.58 and an improved average accuracy of 67.0, which is almost on par with the Dense model. Similarly, the recovery tuning of the Hymba-1.5B-Base model also yields significant improvements (Table 12). Initially, the pruned model with 7 out of 32 Hymba blocks removed shows an average accuracy of 61.7. After recovery tuning, the average accuracy increases to 63.7, which is nearly equivalent to the dense model's accuracy of 63.8. This phase effectively enhances the performance of the pruned model, bringing it closer to the original dense model's performance while maintaining computational efficiency. In summary, recovery tuning is crucial to optimize pruned models, making them more viable for practical applications.

### Insights on the Compression Sensitivity of the Variants of Mamba

The proponents of Mamba modified the original architecture to restrict the expressivity in Mamba-2 and increase the training efficiency. As illustrated on the left side of Figure 4, our experiments suggest that these changes make Mamba-2 models less robust to removing entire blocks than the previous version of the Mamba block. As soon as we remove blocks with the least importance, Mamba-1 exhibits a more robust behavior. However, Mamba-2 demonstrates a significantly higher tolerance to removing SSMs, maintaining a stable perplexity even as more SSMs are pruned, suggesting that while Mamba-2's architectural improvements have made it more sensitive to the removal of Mamba blocks, they have also enhanced its robustness to SSM pruning.

## 5 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.

Figure 4: Close examination of the impact of removing Mamba blocks or SSMs from the two versions of Mamba models reveals distinct differences in their tolerance levels. Mamba-1 exhibits a higher tolerance for removing its blocks, while Mamba-2 exhibits greater tolerance for removing the SSM subcomponent.

 



### Limitations

Despite their outstanding results, large sequence models are still under investigation to better understand their capabilities and limitations. Mamba-Shedder is, to the best of our knowledge, the first work to investigate the removal of structures in Mamba-based models, including hybrids with Transformer blocks. Our goal is to motivate the research community to better understand this class of models to identify opportunities for future improvements in the model architecture and applicable compression techniques. The results indicate that these models contain redundant elements that might be removed to improve their efficiency. However, future work must explore and attempt to better understand the trade-offs between efficiency and accuracy when removing these models' components. Even more research questions can be entertained when considering Transformer blocks and hybrid models, as in the case of Zamba. For instance, there is much to understand about the right mix of the SSM- and Transformer-based elements.

## Ethics Statement

Due to the well-known flaws in modern sequence models, e.g., hallucinations, many guard rails must be in place when considering deploying them in production. Our research focuses on improving the efficiency of these models in existing downstream tasks and datasets. However, further experimentation and analysis are needed when considering deploying these compressed models in environments where their output might affect people's well-being.

## References

* A. Arnab, M. Dehghani, G. Heigold, C. Sun, M. Lucic, and C. Schmid (2021)Vivit: a video vision transformer. In 2021 IEEE/CVF International Conference on Computer Vision (ICCV), Vol. , Los Alamitos, CA, USA, pp. 6816-6826. External Links: Document Cited by: SS3.4.
* S. Ashkboos, M. L. Croci, M. Gennari do Nascimento, T. Hoefler, and J. Hensman (2024)SliceGPT: compress large language models by deleting rows and columns. In The Twelfth International Conference on Learning Representations, Cited by: SS5.
* I. Beltagy, M. E. Peters, and A. Cohan (2020)Longformer: the long-document transformer. arXiv:2004.05150. Cited by: SS1.
* Y. Bisk, R. Zellers, R. Le Bras, J. Gao, and Y. Choi (2020)Piqa: Reasoning about physical commonsense in natural language.. In Thirty-Fourth AAAI Conference on Artificial Intelligence, Cited by: SS1.
* T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, K. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei (2020)Language models are few-shot learners. In Advances in Neural Information Processing Systems, Vol. 33, pp. 1877-1901. External Links: Link, Document Cited by: SS1.
* K. Choromanski, V. Likhosherstov, D. Dolhan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Quincy Davis, A. Mohiuddin, L. Kaiser, D. Benjamin Belanger, L. J. Colwell, and A. Weller (2021)Rethinking attention with performers. In International Conference on Learning Representations, Cited by: SS1.
* P. Clark, Isaac Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tarjford (2018)Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXivabs/1803.05457. Cited by: SS1.
* G. M. Correia, V. Niculae, and A. F. T. Martins (2019)Adaptively sparse transformers. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), Hong Kong, China, pp. 2174-2184. External Links: ISBN 978-3-319-31

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. 2021. An image is worth 16x16 words: Transformers for image recognition at scale. _Preprint_, arXiv:2010.11929.
* Elfwing et al. (2018) Stefan Elfwing, Eiji Uchibe, and Kenji Doya. 2018. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. _Neural Networks_, 107:3-11. Special issue on deep reinforcement learning.
* Frantar et al. (2022) Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. 2022. GPTQ: Accurate post-training compression for generative pretrained transformers. _arXiv preprint arXiv:2210.17323_.
* Fu et al. (2023) Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Re. 2023. Hungry Hippos: Towards language modeling with state space models. In _International Conference on Learning Representations_.
* Gao et al. (2020) Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2020. The pile: An 800gb dataset of diverse text for language modeling. _Preprint_, arXiv:2101.00027.
* Gao et al. (2023) Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Mueminghoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. 2023. A framework for few-shot language model evaluation.
* Glorioso et al. (2024) Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. 2024. Zamba: A compact 7b ssm hybrid model. _Preprint_, arXiv:2405.16712.
* Gong et al. (2021) Yuan Gong, Yu-An Chung, and James Glass. 2021. AST: Audio Spectrogram Transformer. In _Proc. Interspeech 2021_, pages 571-575.
* Gu and Dao (2023) Albert Gu and Tri Dao. 2023. Mamba: Linear-time sequence modeling with selective state spaces. _arXiv preprint arXiv:2312.00752_.
* Gu et al. (2022) Albert Gu, Karan Goel, and Christopher Re. 2022. Efficiently modeling long sequences with structured state spaces. In _The International Conference on Learning Representations (ICLR)_.
* Gu et al. (2024) Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Re. 2024. Combining recurrent, convolutional, and continuous-time models with linear state-space layers. In _Proceedings of the 35th International Conference on Neural Information Processing Systems_, NIPS '21, Red Hook, NY, USA. Curran Associates Inc.
* Hoefler et al. (2021) Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikolid Dryden, and Alexandra Peste. 2021. Sparsity in deep learning: pruning and growth for efficient inference and training in neural networks. _J. Mach. Learn. Res._, 22(1).
* Hu et al. (2022) Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022. LoRA: Low-rank adaptation of large language models. In _International Conference on Learning Representations (ICLR)_.
* Katharopoulos et al. (2020) Angelos Katharopoulos, Apoorty Vyas, Nikolaos Papapas, and Francois Fleuret. 2020. Transformers are rnns: fast autoregressive transformers with linear attention. In _Proceedings of the 37th International Conference on Machine Learning_, ICML'20. JMLR.org.
* Lagunas et al. (2021) Francois Lagunas, Ella Charlaix, Victor Sanh, and Alexander Rush. 2021. Block pruning for faster transformers. In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_, pages 10619-10629, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
* LeCun et al. (1989) Yann LeCun, John Denker, and Sara Solla. 1989. Optimal brain damage. In _Advances in Neural Information Processing Systems_, volume 2. Morgan-Kauffmann.
* Ma et al. (2023) Xinyin Ma, Gongfan Fang, and Xinchao Wang. 2023. LLM-pruner: On the structural pruning of large language models. In _Thirty-seventh Conference on Neural Information Processing Systems_.
* Men et al. (2024) Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianguei Han, and Weipeng Chen. 2024. ShortgtpT: Layers in large language models are more redundant than you expect. _Preprint_, arXiv:2403.03853.
* Mihalvo et al. (2018) Todor Mihalvo, Peter Clark, Tschar Khot, and Ashish Sabharwal. 2018. Can a suit of armor conduct electricity? a new dataset for open book question answering. In _Conference on Empirical Methods in Natural Language Processing_.
* Munoz et al. (2024) J. Pablo Munoz, Jinjie Yuan, and Nilesh Jain. 2024. Shears: Unstructured sparsity with neural low-rank adapter search. In _Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 6: Industry Track)_, pages 395-405, Mexico City, Mexico. Association for Computational Linguistics.
* Munoz et al. (2025) J. Pablo Munoz, Jinjie Yuan, and Nilesh Jain. 2025. Multipruner: Balanced structure removal in foundation models. _Preprint_, arXiv:2501.09949.
* Paperno et al. (2016) Denis Paperno, German Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Geimma Boleda, and Raquel Fernandez. 2016. The LAMBADA dataset: Word 

prediction requiring a broad discourse context. In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_, pages 1525-1534, Berlin, Germany. Association for Computational Linguistics.
* Parmar et al. (2018) Niki J. Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. 2018. Image transformer. In _International Conference on Machine Learning (ICML)_.
* Radford et al. (2021) Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. 2021. Learning transferable visual models from natural language supervision. In _International Conference on Machine Learning (ICML)_.
* Sakaguchi et al. (2021) Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2021. Winogrande: An adversarial winograd schema challenge at scale. _Commun. ACM_, 64(9):99-106.
* Sun et al. (2023) Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. 2023. A simple and effective pruning approach for large language models. _arXiv preprint arXiv:2306.11695_.
* Tokpanov et al. (2024) Yury Tokpanov, Beren Millideg, Paolo Glorioso, Jonathan Pilault, Adam Ibrahim, James Whittington, and Quentin Anthony. 2024. Zyda: A 1.3 t dataset for open language modeling. _arXiv preprint arXiv:2406.01981_.
* Touvron et al. (2023) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023. LLama 2: Open foundation and fine-tuned chat models. _arXiv preprint arXiv:2307.09288_.
* Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasaz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In _Advances in Neural Information Processing Systems_, volume 30. Curran Associates, Inc.
* Xu et al. (2024) Peng Xu, Wenqi Shao, Mengzhao Chen, Shitao Tang, Kaipeng Zhang, Peng Gao, Fengwei An, Yu Qiao, and Ping Luo. 2024. Besa: Pruning large language models with blockwise parameter-efficient sparsity allocation. _Preprint_, arXiv:2402.16880.
* Zellers et al. (2019) Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. Hellaswag: Can a machine really finish your sentence? In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_.
* Zhang et al. (2023) Li Zhang, Jiachen Lu, Sixia Zheng, Xinxuan Zhao, Xian Zhu, Yanwei Fu, Xiang Tao, and Jianfeng Feng. 2023. Vision transformers: From semantic segmentation to dense prediction. _arXiv_.
* Zheng et al. (2022) Lin Zheng, Chong Wang, and Lingpeng Kong. 2022. Linear complexity randomized self-attention mechanism. In _International Conference on Machine Learning_, pages 27011-27041. PMLR.
* Zhong et al. (2024) Longguang Zhong, Fanqi Wan, Ruijun Chen, Xiaojun Quan, and Liangzhi Li. 2024. Blockprunner: Fine-grained pruning for large language models. _Preprint_, arXiv:2406.10594.
* Zuo et al. (2024) Jingwei Zuo, Maksim Velikanov, Dhia Eddine Rhaitem, Ilyas Chahed, Younes Belkada, Guillaume Kunsch, and Hakim Hacid. 2024. Falcon mamba: The first competitive attention-free 7b language model. _arXiv preprint arXiv:2410.05355_.

 



## Supplementary Material

### Related Work

Transformers Vaswani et al. (2017) and its variants are the primary building block of successful deep learning architectures, e.g., Llama Touvron et al. (2023) and GPT Brown et al. (2020), that have revolutionized Natural Language Processing (NLP) Devlin et al. (2019); Gao et al. (2023), Computer Vision (CV) Parmar et al. (2018); Radford et al. (2021); Zhang et al. (2023), and many other domains. Due to the Transformer's popularity, researchers have proposed variants to improve their computational and memory efficiency further and tackle issues like their quadratic complexity in sequence length during training Correia et al. (2019); Beltagy et al. (2020); Dai et al. (2020); Choromanski et al. (2021); Katharopoulos et al. (2020); Zheng et al. (2022).

A parallel research effort investigates alternatives to Transformers in the form of _structured state space models_ (SSMs) that can power the next generation of sequence models. The initial proposals of structured SSMs were linear time-invariant, e.g., LSSL Gu et al. (2024), S4 Gu et al. (2022), H3 Fu et al. (2023). Recent improvements to the state space model formulation have resulted in the proposal of time-varying selective SSMs, e.g., Mamba Gu and Dao (2023); Dao and Gu (2024).

To our knowledge, Mamboa-Shedder is the first study on pruning selective structured state space models (Mamba) and their hybrids. On the other hand, many works have proposed pruning techniques for Transformer-based models Hoefler et al. (2021). Several of these works focus on _unstructured_ pruning Sun et al. (2023); Xu et al. (2024); Frantar et al. (2022), which can achieve higher sparsity levels. However, it requires highly optimized runtimes to realize the benefits of sparsity. Sophisticated solutions have been proposed to fine-tune sparse models and recover any accuracy drop from the pruning stage Munoz et al. (2024). Recently, _training-free_ approaches have been proposed for _structured_ pruning of Transformers. These approaches cannot achieve high sparsity levels as the _unstructured_ pruning approaches. However, they are very convenient because their compressed models do not require specialized runtimes and exhibit beneficial inference acceleration. In this line of research, LLMPruner Ma et al. (2023), ShortGPT Men et al. (2024), BlockPruner Lagunas et al. (2021), SliceGPT Ashkboos et al. (2024), and MultiPruner Munoz et al. (2025) have demonstrated efficient methods for Transformer pruning. BlockPruner improved over many previous approaches by proposing a global metric that can be used to determine the importance of a selected network structure. MultiPruner extended this approach to pruning the width dimension, as well. Mamboa-Shedder builds on these works and the rest of the extensive literature on _structured_ block pruning to explore opportunities for removing redundancies in models with Mamboa blocks.

### Hyperparameters

Table 13 offers a detailed summary of the hyperparameters employed in our experiments, promoting both reproducibility and clarity.

\begin{table}
\begin{tabular}{l c} \hline \hline
**Hyper-parameter** & **Value** \\ \hline _Prunning Stage:_ & \\ Calibration Dataset & tatsu-lab/alpaca \\ Importance Metric & Perplexity (PPL) \\ Number of Calibration Samples & 256 \\ MLP Channel Group Size (Zamba2) & 1024 \\ Steps of MLP Channel Pruning (Zamba2) & 20 \\ \hline \hline \end{tabular}
\end{table}
Table 13: Hyper-parameters used in the experiments.