Sparsified State-Space Models are Efficient Highway Networks
Authors: Woomin Song
Jihoon Tack
Sangwoo Mo
Seunghyuk Oh
Jinwoo Shin
Venue: Accepted by TMLR ()

Abstract: State-space models (SSMs) offer a promising architecture for sequence modeling, providing an alternative to Transformers by replacing expensive self-attention with linear recurrences. In this paper, we propose a simple yet effective trick to enhance SSMs within given computational budgets by sparsifying them. Our intuition is that tokens in SSMs are highly redundant due to gradual recurrent updates, and dense recurrence operations block the delivery of past information. In particular, we observe that upper layers of SSMs tend to be more redundant as they encode global information, while lower layers encode local information. Motivated by this, we introduce Simba, a hierarchical sparsification method for SSMs based on token pruning. Simba sparsifies upper layers more than lower layers, encouraging the upper layers to behave like highways. To achieve this, we propose a novel token pruning criterion for SSMs, measuring the global impact of tokens on the final output by accumulating local recurrences. We demonstrate that Simba outperforms the baseline model, Mamba, with the same FLOPS in various natural language tasks. Moreover, we illustrate the effect of highways, showing that Simba not only enhances efficiency but also improves the information flow across long sequences. Code is available at https://github.com/woominsong/Simba.

Published in Transactions on Machine Learning Research (03/2025)
Sparsified State-Space Models are Efficient Highway Networks
Woomin Song
woomin.song@haris..c.kkk
Korea Advanced Institute of Science 8 Technology (KAIST)
Jihoon Tack
jihoontackkkaist..aco
Korea Advanced Institute of Science & Technology (KAIST)
Sangwoo Mo
sumo@umich.coo edu
University of Michigan, Ann Arbor
Seunghyuk Oh
seunghyukoo.Bkaist..co oc.kr
Korea Advanced Institute of Science & Technology (KAIST)
Jinwoo Shin
jinwoos@kaist.ooc oc.kr
Korea Advanced Institute of Science & Technology (KAIST)
Reviewed on OpenReview: https: openrevvee... net/ forum? id= G1d0Yuuuu
Abstract
State-space models (SSMs) offer a promising architecture for sequence modeling, providing an alternative to Transformers by replacing expensive self-attention with linear recurrences. In this paper, we propose simple yet effective trick enhance SSMs within given computational budgets by sparsifying them. Our intuition is that tokens in SSMs highly redundant due to gradual recurrent updates, and dense recurrence operations block the delivery of past information. In particular, observe that upper layers of SSMs tend to be more redundant as they encode global information, while lower layers encode local information. Motivated by this, we introduce Simba, a hierarchical sparsífication method for SSMs based on token pruning. Simba sparsifies upper layers more than lower layers, encouraging the upper layers to behave like highways. To achieve this, we propose a novel loken pruning criterion for SSMs, measuring the global impact of tokens on the final output by accumulating local recurrences. We demonstrate that Simba outperforms the bascline model, Mamba, with the same FLOPS in various natural language tasks. Moreover, illustrate the effect of highways, showing that Simba not only enhances efficiency but also improves the information flow across long sequences. Code is available at
wtwp:/thuu..co/oooooo/ooo con//woominsoon//iinta.
1 Introduction
State-space models (SSMs) Gu et al. 2022b Gu & Dao 2023 offer a promising architecture for sequence modeling, efficiently handling sequences using linear recurrence structures. Thanks to this efficiency, SSMs have shown potential mechanism, incurring high computational costs for long sequences. In particular, Mamba (Gu & Dao 2023 has recently demonstrated that SSMs can scale up to billions of parameters and show comparable performance
alternative to Transformers (Vaswani al. 2017), which use the self-attention
with Transformers in various domains (Zhu et. al. 2024 Li et. al. 2024a bb_
After their success, numerous works have aimod to enhance SSMs and Mamba. further. One popular approach
involves hybrid models combining Transformers and SSMs (Lieber et al. [024 Poli al.) 2024). This approach assists the models in retaining past information through the global memory of Transformers (Vardasbi et. al.

---

Published in Transactions on Machine Learning Research (03/2025)
Network Layers
SSM black
Pruned block
Layer
Highwey connection
Pythia
Mambe
Simba (ours)
62.51
Layer
54.94
Accuracy (1 is better)
Layer 2
Pythia
Mamba
Simba (ours)
2.30
Layer
2.35
Inputs
x1
x2
x3
x4 x5 X6
Perplexity (1 is better)
(a) Conceptual illustration of Simba
(b) Performance highlights
Figure 1: Simba: hierarchical sparsification of SSMs via token pruning. (a) We found that tokens in state-space models (SSMs) are highly redundant, especially the upper layers. Motivated by this observation, we propose hicrarchically sparsifying pre-trained SSMs by progressively pruning tokens across layers. This results in models with a trapezoidal shape, featuring sparse upper layers that act like highways, enhancing efficiency and information flow of the original SSM. (b) We highlight results comparing Simba-2.88 with Mamba and Pythia, all with the same number of FLOPS. We report the mean accuracy over 6 NLP benchmarks and perplexity on the PG-19 dataset with 2k context, following the sctups in Section Simba outperforms both models in accuracy and perplexity.
2023:
Jelassi et al. 2024). However, this compromiscs the efficiency of SSMs by reintroducing expensive self- attention. Instead of sacrificing efficiency, we explore an alternative direction to improve SSMs by comparing models on a fixed computational budget. Specifically, we investigate sparsification of large pre-trained SSMs, known yield better models than training small ones from scratch (Frankle & Carbin 20189).
To this end, we first analyze the behavior of tokens in pre-trained SSMs. We observe that tokens in SSMs are highly redundant, they gradually updated over sequences. This redundancy is particularly noticeable in the upper layers. Furthermore, dense recurrence operations over the redundant tokens block the delivery of past information, potentially harming the contextual understanding of SSMs.
Inspired by this, we propose Simba, a simple yet effective method to sparsify SSMs through token pruning (i.e., it is training-free). Our core idea is sparsify SSMs into a hierarchical form, enforcing more sparsity in upper layers than in lower layers. As result, the upper layers behave as highways to transmit past. information, enabling efficient inference and facilitating the information flow across long sequences. Figure illustrates the visual concept of our approach, obtaining trapezoidal-shappr sparsified network.
To implement this, we propose a novel token pruning criterion for SSMs. Specifically, our score measures the global influence of each token on the final output by reformulating SSM equations to accumulate the effect from local recurrences. This approach can also be viewed as an SSM extension of attention based token pruning criteria used for Transformers (Goyal et al. 2020). We found that our criterion outperforms intuitive basclines, such as uniform pruning of tokens with even intervals.
Our experiments show that Simba, obtained by sparsifying Mamba without any fine-tuning, significantly outperforms Mamba using the same number of FLOPS in various tasks. For instance, Simba consistently achieves better FLOPS-accuracy curves on 6 NLP benchmarks, including Lambada (Paperno et al. 2016).
HellaSwag [Zellers et al. 2019), PIQA (Bisk et. al.) 2020). ARC-Challenge (Clark et al.) 2018). ARC- Easy (Clark et al. 2018), and WinoGrande (Sakaguchi et al. 2021). Here, Simba obtained from Mamba-2.8h
performs on par with the original Mamba-2.88 despite using similar FLOPS Mamba- 1.4b, as highlightod in Figure lb For example, it achieves an average accuracy of 62.5% for 6 downstream NLP tasks, improving 58.8% of Mamba-1.4h.
2

---

Published in Transactions on Machine Learning Research (03/2025)
We also demonstrate the language modeling ability of Simba by measuring perplexity on the PG-19 dataset (Rac et al. 2019) across different context lengths. Like the NL.P benchmarks, Simba achieves better perplexity than Mamba using the same number of FLOPS. More importantly, Simba performs robustly over long sequences exceeding the pre-trained context length, such as twice longer than the trained length, unlike Mamba, which significamlly deteriorates with length extrapolation. This supports the idea that the highway structures in Simba facilitate long sequence modeling.
We further investigate the effect of highways in Simba. Somewhat unexpectedly, we found that Simba performed better than its original unpruned Mamba in some of our experiments, potentially benefiting from the highways created at the upper layers. To further investigate the positive effect of highways, we examine the information flow across layers by assessing the influence of the sequence tokens on the final output. We observe that Mamba relies on tokens near the end across all layers, while Simba also focuses on earlier tokens at the upper layers, showcasing the role of highways.
2 Related work
Statc-space models (SSMs) are a powerful architecture for sequence modcling, integrating concepts from classic control theory (Kalman 1960) with recurrent. neural networks (Elman 1990]. The key idea of SSMs is
to employ linear recurrence (Katharopoulos et al. 2020 Gu et al. 2022b a Mehta et al. 2023 Smith et al.,
2023 Fu et al. 2023 Orvieto et al., 2023] Poli et al. 2023 Peng et al. 2023 Sun et al.] 2023 De et al. 2024).
enabling efficient parallel inference and effective training, unlike Transformers using sel self-attention (Vaswani
et all (017),, which requires quadratic computation over the sequence length. As result, SSMs have shown success in handling long sequences (Tay et al., 2021). Recently, Mamba (Gu & Dao, 2023) further scaled up SSMs through a selection mechanism and hardware-aware algorithm, showing the potential of SSMs in
challenging tasks such as language, audio, and video (Zhu et al.) 2024: Li et al., 2024a b). We aim to further
improve Mamba through network sparsification, efficiently utilizing fixed computational budget.
Sparsifying networks have been widely studied, primarily for training efficient models (Han et al. 2016 Most prior work focused on weight, pruning, which removes unnecessary edges in weight. matrices (Zhu & Guptal 2017 Gale et al. 2019 Park et al. 2020 Lee et al., 2021). However, while weight pruning reduces model size, it does not enhance inference speed due to the batch computation nature of GPUs. To address this, structured pruming removes entire blocks at once, such as channels in CNNs (Li et al., 2017) or attention heads in Transformers (Michel et al. 2019). On the other hand, some works focused benefits of sparsified networks, suggesting that sparsifying large networks yields better models than training
the performance
small models from scratch, known as lottery tickets (Frankle & Carbin 20189 Li et al. 2020 Our work relates to structured pruning, as removing tokens speeds up computation and reduces memory usage, and to lottery tickets, as it also improves performance.
Token pruning (or merging) is widely applied in Transformers to reduce heavy computation over long sequences (Goyal al.) 2020 Kim et al., 2022 Liu al. 2022 Bolya et al., 2023 Ke et al. 2024 Shang et al.) 2024). In particular, HOMER (Song et al. 2024 has shown that hicrarchical token pruning not only reduces computation but also enhances long context understanding by condensing global information into sparse tokens in the upper layers. Our work shares a similar spirit with HOMER but has notable differences. First, we explore token pruning for SSMs, unlike prior works focused on Transformers. To this end, we propose a novel token pruming criterion based on the global importance of tokens to the final output, derived from reformulating SSM equations to accumulate local recurrences. Second, by targeting SSMs, our hierarchical pruning scheme offers a novel interpretation of highway networks, connecting SSMs with classical recurrent networks with long-term memory. As result, our token pruning approach enhances both the inference efficiency and information flow of SSMs.
Highway networks have been proposed to bypass information loss from dense computation through local residuals (He el al.] 2016) or long skip connections (Srivastava et. al. 2015 Zilly et al.] 2017 Huang et al.] 2017 Ronneberger et al. 2015). Highways, also termed long-term memory, were used as a standard approach for sequence modeling before global computation methods like solf-attention gained popularity (Hochreiter & Schmidhuber, 1997 Chung et al. 2014 Feichtenhofer et al. 2019). However, integrating highways with SSMs proved challenging due to the linear recurrence structure of SSMs, not, easily combined with skip connections.
3

---

Published in Transactions on Machine Learning Research (03/2025)
Instead of explicitly using such modules in SSM architectures, we introduce simple and effective way to integrate highways by pruning dense token connections from the upper layers.
Tree RNNs have been explored for processing hierarchical data structures, such as word, sentence, and paragraph hierarchies in language (Hihi Bengio, 1995 Wang et al. 2019 Despite aligning with human intuition, most methods were unsuccessful due complex architectures, while simple linear sequence modeling demonstrates its power (Achiam et al., 2023). Our token pruming naturally incorporates this hierarchical structure into SSMs, while favoring the success and scalability of linear sequence modeling.
3 Simba: Hierarchical sparsification for state-space models
In Section 3.1 we review the mathematical formula of state-space models (SSMs) and discuss observations the token redundancy of Mamba. In Section 3.2 we describe our proposed hierarchical sparsification approach for SSMs and explain the token pruming criteria.
3.1 Motivation: Hierarchy in SSM token redundancy
Structured state-space model (54) Gu et al. 2022b) is family of recently proposed SSMs. In its continuous form, the SSM updates the state h(t)) using the input sr(() and produces the output v(() according to Eq. where A, B, and C' refer to the parameters of SSMs. It discretizes the update rules for discrete sequences, as shown Eq. (2)
(t()) AM()) Bx((),
u(t) Ch(()
(1)
Ah + Brg Yt Ch.
(2)
Mamba (Gu & Dao 2023) further improves this formulation using an input selectivity mechanism, creating the matrices A, B, and C dependent on the imput zt. Thus, the state update equation can be rewritten as follows, where A,, Be, and C, denote the input-dependent matrices created using 2t-
h = Ache-1 BYT,
Yt C.h.
(3)
Hierarchy in token redundancy. As the state up- date depends solely the input and the immediate previous states, the model must compress all previ- ous information into the state of each token. Thus, the states of SSM tokens are likely highly redundant, as states at similar positions would compress a sim- ilar set of information. To verify this, we visualize the token redundancy of Mamba in Figure 2 by mea- suring the cosine similarity between adjacent tokens across the network layers.
60
50
0.6
Layers
30
20
0.2
10
Our analysis indicates that tokens in the upper lay- ers tend to exhibit more redundancy than those in the lower layers. One possible explanation is that SSMs process information a hierarchical manner: lower layers focus on local information, while upper layers focus global information. This aligns with findings from attention visualizations of Mamba mod- cls Ali ct al. (2024). where lower layers emphasize diagonal elements in the allention map, while up- per layers highlight nover-triangulat elements. Since upper layer states encapsulate global information, states similar positions contain similar informa- tion, resulting in higher redundancy.
1000
2000
Token Position
Figure 2: Token redundancy of SSMs has a hicr- archical structure. We measure the cosine similarity between adjacent tokens of the Mamba-2.88 model across layers, averaged over documents from the PG-19 test dataset. The tokens are highly redundant, espe- cially in the upper layers.
4

---

Published in Transactions on Machine Learning Research (03/2025)
3.2 Hierarchical sparsification for SSMs
We introduce Simba, a simple yet effective sparsification approach that can be directly applied to any pre-trained state-space models (SSMs) in a plug-and-play manner. This introduces an efficient highway for effective sequential modeling. Based on prior motivation, Simba sparsifies the full SSM in a hierarchical manner: it sparsifies redundant tokens upper layers while preserving the local information captured lower layers, guided by our novel token importance criteria.
Hierarchical sparsification through token pruning. To achieve sparsity in SSMs in hierarchical manner, we propose loken pruning al each layer with specified rate. Implementing this pruning technique at a particular layer automatically decreases the computational burden on subsequent layers, as pruned tokens are no longer propagated to the next layer. Consoquently, by soquentially applying token pruning from lower to upper layers, the pruning ratio consistently increases during propagation, thereby introducing hierarchical sparsification to SSMs.
Token pruning criterion. We propose a novel token importance score for SSMs to prune the least important token from each layer. To achieve this, we leverage the influence that each token has on the final token's output. This calculation is straightforward due to the linear recurrence property of SSMs. Formally, for a given token sequence (Iı,--- ,TT) of length T and the final token output yr of the layer, we estimate the influence of token 2, position by considering the updated output yا (obtained by removing I from yr), as follows:
(yr = V/-
V1
(t)
T
T
T
T
Σ Cr
II
Ap
BAT
Σ
Cr
II kart1
As
BArr
(4)
k=r+1
r=1.rt
T
Cr
Av
BeTr.
k=1+1
Here, we found that the inffuence measure Ayr(t) aggregated with the max pooling (or the & norm) serves effective pruning criterion, denoted s(t) max(Ayr (t)); we max pooling throughout the paper as it was slightly belter than & norm. Also, empirically observed that excluding the bias term when computing Aş in fully connected layers was beneficial, as it potentially helps normalize token importance scores across different channels before applying max pooling, resulting in more consistent pruning decisions. Based on the proposed score s(l), we prune each layer by removing the tokens with the lowest scores to obtain sparsified SSM. The pruning process performed per input sequence during inference and is based solely on token importance scores computed for that specific input. As a result, no external pruning data is required to determine the pruning pattern.
Pruning schedule. The pruning schedule aims to balance the trade-off between efficiency and performance. Upper layers, with their ability to model global context, are better equipped to identify important tokens. Conversely, pruning tokens at carlier layers offers more significant computational savings. Therefore, we propose a linear pruning schedule, inspired by prior work on Transformers (Bolya et al. 2023), where the number of active tokens lincarly reduced across all layers, resulting in a trapezoidal-shaped network after sparsification.
Highways in sparsified SSMs. Long-term dependency is well-known challenge for recurrent mod- cls (Hochreiter & Schmidhuber 1997). Dense recurrence operations tend to attenuate previous information, restricting the information flow across distant tokens. Our sparsification scheme addresses this issue by reducing the number of recurrence operations. As unimportant tokens are pruned during sparsification, upper layers can selectively focus on processing more important information without being burdened by dense recurrence operations on redundant tokens. This highway effect enables our pruning scheme not only enhance inference efficiency but also facilitate the information flow across distant tokens (Section 4.3
5

---

Published in Transactions on Machine Learning Research (03/2025)
4 Experiments
In this section, we demonstrate the performance of Simba on diverse tasks. In Section we evaluate its performance on 6 NLP benchmarks, consistently showing superior performance compared to dense models with equivalent computational resources. In Section 4.2 we assess the language modeling ability of Simba by measuring perplexity conditioned on various context lengths. Tn Section 4.3 we further investigate the highway effects of Simba. Finally in Scction 4.4 we perform ablation studies on token pruning criteria and different pruning ratios, along with a simple fine-tuning experiment.
Common setups and baselines. We primarily apply sparsification method to pre-trained Mamba models of various scales. For our method, we implement linear pruning schedule that preserves 10% of the tokens at the final layer unless specified otherwise. The Simba models are generated by sparsifying Mamba models without any fine-tuning, following a alug-and-play approach. We conduct performance comparisons of Simba against Mamba (Gu Dao 2023) and Pythia (Biderman et al., 2023 models that utilize similar amount of computation. In NLP tasks consisting of prompt and a label, we exclusively apply sparsification to the prompts to prompt remains unpruned, as its output is utilized for computing the label logits. Moreover, to accommodate the task structure, the token importance score is computed with respect to the final token of the prompt. For additional details, refer to Appendix A
accurate evaluation of the label logits. The final token of the
60
Acourapy (%)
Accuracy (%)
Accuracy (%)
50
70
30
Pythia
Pythia
65
Pythia
Mamba
Mamba
Mamba
Simba (ours)
30
Simba (ours)
Simba (ours)
0
4
0.5 FLOPs (x1e12)
1.0
FLOPs (x1e12)
FLOPs (x1e12)
(a) Lambada
(b) HellaSwag
(c) PTQA
70
40
60
Accuracy (%)
Accuracy (%)
35
Accuracy (%)
60
30
55
Pythia
Pythia Mamba
Pythia
Mamba
25
Mamba
50
Simba (ours)
Simba (ours)
Simba (ours)
4
FLOPs (x1e12)
FLOPs (xte12)
FLOPs (x1e12)
(d) Arc-Easy
(e) Arc-Challenge
(f) WinoGrande
Figure 3: Performance on NLP Benchmarks. We visualize the FLOPs-accuracy curve of Mamba, Pythia, and Simba models of various scales on 6 NLP benchmarks. Across all benchmarks, Simba consistently outperforms the baselines using the same number of FLOPs.

---

Published in Transactions on Machine Learning Research (03/2025)
4.1 NLP benchmarks
In this section, we assess the language understanding capability of Simba by evaluating six downstream NLP tasks. Specifically, we prcsent the performance and computational efficiency of Simba on the Lambada Paperno et al. 2016). HellaSwag (Zellers et al., 2019), PIQA (Bisk et al. 2020), ARC- CRCC-hallenge (Clark et al, 2018), ARC-Easy (Clark al. 2018)018), and WinoGrande (Sakaguchi al. 2021) benchmarks. Consistent with Mamba (Gu & Dao 2023), we report accuracy normalized by sequence length for HellaSwag and ARC-Challenge, and accuracy for the other datasets. All evaluations use the LM evaluation harness from AleutherAT Gao et al. 2021
We report evaluation accuracy and computational cost for each benchmark in Figure 3 results, Simba provides the best the same number FLOPs. This demonstrates that Simba can successfully make dense Mamba models sparse, advancing the frontier of the accuracy-cfficiine trado-off. We provide the full results for all benchmarks Table of the Appendix.
As evident from the trade-off, consistently outperforming other models with
4.2 Language modeling
3.1
Mamba Simba (ours)
2.7
Mamba
Simba (ours)
Perplexity
Perplexity
2.9
2.6
2.7
2.5
2.6
1024
2048
3072
4096
1024
2048
3072
4096
Context Length
Length
(a) Mamba-1300 130m and Simba-370m
(b) Mamba-370m and Simba-790m
Mamba Simba (ours)
Mamba Simba (ours)
2.50
2.40
Perplesity
Perplexily
2.35
2.30
2.40
1024
2048
3072
4096
1024
2048
3072
4096
Context Length
Context Length
(c) Mamba-790m and Simba-1.4b
(d) Mamba-1.4h and Simba-2.8h
Figure Language modeling ability. We measure the FLOPs-perplexitg curves on the PG-19 test dataset. Simba models are compared against Mamba models that similar computation. Simba not only outperforms Mamba with the same computation but also shows decreasing perplexity after its pre-trained context limit of 2k tokens.
In this section, evaluate the language modeling ability of Simba by measuring perplexity long documents. Specifically, we measure the perplexity of short document snippets sampled from PG-19 dataset (Rac et al. 2019), conditioned on varying amounts of context. We keep the 100-token snippet fixed for all experiments to ensure that all perplexity measurements are done on the same set of Lokens.
We report perplexity values conditioned on varying contexts Figure The results show that Simba consistently shows improved perplexity with similar computation, outperforming the dense Mamba models. We provide the full results, including Pythia, in Table of the Appendix.
7

---

Published in Transactions on Machine Learning Research (03/2025)
Long context capability. An intriguing observation is that Simba, unikee Mamba models, exhibits decreasing perplexity even after surpassing its pre-trained context limit of 2k tokens. Context limits, defining the maximum number of tokens model can handle, often result in catastrophic performance drop if exceeded (Song et al. 2024h. While SSMs are generally more resilient to this issue, the results indicate increasing perplexity for Mamba models with longer contexts, suggesting ing ineflective utilization of additional context. In contrast, Simba consistently demonstrates decreasing perplexity even with extended contexts, highlighting its adeptness at leveraging extra information.
This benefit is likely attributed to the highways formed in the upper layers due to extensive sparsification. The performance drop with longer inputs typically results from distribution shift in the training data, as the model never saw the long imput length during training. For dense models like Mamba, this shift affects all network layers. Conversely, sparse models like Simba process significantly fewer tokens in the upper layers, thanks to extensive sparsification. Consequently, the upper layers remain unaffocted by distribution shifts input lengths, processing inputs more effectively.
In summary, our experiments demonstrate that Simba not only outperforms Mamba with the same computation but also exhibits superior long-context handling abilities, suggesting the benefits of highways.
4.3 Sparsified SSMs as highway networks
In this section, we further investigate the highway effects of Simba. First, we identify some scenarios where Simba models perform better than the original dense models despite using fewer FLOPs, potentially benefiting from the highways. Second, we examine the information flow in the model, showing that highways assist in obtaining information from earlier tokens, unlike dense SSMs, which over-rely on later tokens.
Table 1: Comparsion between same model scales. We compare Simba and Mamba on NLP benchmarks: Lambada (Lbd.), HellaSwag (HS), PIQA, Arc-Easy (Arc-E), Arc-Challenge (Arc-C), and WinoGrande (WG). We use moderale pruning ratio for Simba, leaving 70% of the tokens at the final layer. Bold denotes the best results, showing that Simba often improves Mamba while using fewer FLOPS.
Model
FLOPs (x1e12)
Lbd.
HS
PIQA
Arc-F.
Arc-C
WG
Avg.
Model
Scale
Dim.
acc. (t)
acc. (t)
acc. (t)
acc. (t)
acc. (t)
acc. (t)
acc. (t)
Mamba
130m
768
0.48
31.84
34.88
64.36
50.76
23.72
49.80
42.56
Simba (ours)
130m
768
0.39
32.43
35.05
64.42
50.38
24.32
49.88
42.75
Mamba
370m
1024
1.38
46.61
46.46
69.59
59.05
27.17
55.80
50.83
Simba (ours)
370m
1024
1.15
47.09
46.69
69.64
59.51
27.65
55.49
51.01
Mamba
790m
1536
2.94
51.33
55.12
72.47
65.49
32.76
57.54
55.78
Simba (ours)
790m
1536
2.47
51.82
54.93
72.20
65.87
32.76
57.22
55.80
Comparison under same model sizes. In the previous experiments, mainly compare models with the same number of FLOPS. Here, we provide an additional comparison between the dense and sparse models that have the same scale. Specifically, we evaluate Simba models with a more moderate pruning ratio, using a linear pruming schedule with 70% of tokens remaining at the final layer. We benchmark Mamba and Simba models using the 6 NLP benchmarks, following the setup in Section 4.1 We provide the evaluation results in Table
Somewhat unexpectedly, we found that Simba models sometimes perform even better than the original model, indicating the possible benefits of highways created at the upper layers. The gain is more evident for smaller models, possibly because dense recurrence operations are more harmful to smaller state sizes, so highways provide more benefits.
Highways facilitate information flow from early tokens. We investigate how information flows through the dense and sparse SSM layers. Specifically, we mcasure the influence of tokens at each position on the final token, using our token importance score Eq. (4) but normalized to equalize the contribution of each
input document, Le., (tt//llrrl| See Appendix A.2 for more details.
8

---

Published in Transactions on Machine Learning Research (03/2025)
Mamba
10'3
Mamba
10
Mamba
10"2
Simba (ours)
+
Simba (ours)
Simba (ours)
Influence
Influence
Influence
10
10-3
10
10
10
10-6
0.0
0.2
0.4
0.0
0.2
0.8
0.0
0.4
0.6
Relative Position
Relative Position
Relative Position
Layer
(b) Layer 13
Layer 26
10
10
Mamba Simba (ours)
Mamba Simba (ours)
Mamba Simba (ours)
10
10
Influence
Influence
Influence
10-3
10-3
10
10-10
10
10
10-15
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.8
Relative Position
Relative Position
Relative Position
(d) Layer 38
(e) Layer 51
(f) Layer 64
Figure 5: Information flow across layers. We visualize the information flow using the normalized token influence score. We compare Mamba-2.8h and Simba-2.8b, averaging scores over the PG-19 test dataset samples. The information flow of Simba flattens at the upper layers, indicating better information flow from carly tokens.
We illustrate the results in Figure Mamba demonstrates consistent information flow pattern across all layers, with tokens near the end exerting more significant influence on the final token. Simba displays similar trend in the lower layers. However, the slope flattens in the upper layers, indicating that they function as highways, facilitating the flow of information from early tokens.
40
Accuracy (%)
36
34
4.4 Ablation study and analysis
Dense
Uniform
Random
This scction presents ablation studics token pruning criteria and different pruning ratios. We also perform simple fine-tuning experiment to further improve the performance of Simba.
32
Influence
2.0
2.5
3.0
4.0
4.5
FLOPs (x1e12)
Pruning criteria. We compare our proposed pruning criteria Eq. (4) with two baselines: "Random," which chooses tokens from random positions, and "Uniform," which chooses tokens from evenly distributed intervals. We visualize the efficiency- perfoormaceeeeeeeeeeeeee trade-off of Mamba- 2.8b and Simba-2.8b models on the Are-Challenge dataset in Figure 6 Random pruning significantly hurts perfor- mance, while Uniform pruning forms a strong baseline, highlighting the necessity of proper token selection. Simba further improves upon Uniform pruning by considering token influence.
Figure 6: Ablation study on pruning critcria. Performance of sparsified Mamba-2.8h models the Are- Are-Challenge dataset (Clark et al., 2018 evaluated using different pruning criteria across various pruning ratios. Our proposed token in- fluence score Eq. (4) performs the best, even remaining robust under severe sparsification. We also report the performance of the dense Mamba- 1.4b and Mamba-2. -28b models for comparison.

---

Published in Transactions on Machine Learning Research (03/2025)
Table Fine-tuning results. We compare the language modeling perplexity on the test split the PG-19
dataset, measured with different context lengths. The fine-tuned Simba model consistently outperforms its training- frre counterpart.
FLOPs (xle12)
Within Context
Extrapolation
Model
Scale
0.5k
1k
1.5k
2k
2.5k
3k
3.5k
4k
Mamba
130m
0.48
2.943
2.917
2.907
2.899
2.899
2.900
2.901
2.905
Simba (training-free)
370m
0.68
2.727
2.686
2.662
2.650
2.643
2.633
2.630
2.625
Simba (fine-tuned)
370m
0.68
2.723
2.678
2.658
2.645
2.637
2.630
2.626
2.621
Mamba
370m
1.38
2.645
2.616
2.607
2.600
2.599
2.601
2.604
2.617
Pruning ratio. We compare the performance of Simba using different sparsity levels, with linear pruning schodules leaving 90%, 80%, 70%, 60%, 50,, 30,, and 10% of tokens at the final layer. Figure presents the performance curves for different. pruning criteria. Simba is robust extreme sparsity, retaining performance even when pruning 90% of tokens the final layer.
Fine-tuning. Although method can be applied to pre-trained SSMs without training, investigate if the performance be improved with further fine-tuning. To this end, perform simple fine-tuning experiment, further training the Mamba-370m model with MiniPile dataset (Kaddour) 2023), which is a subset of the pre-training dataset (the Pile /ao et al. 2020)) used for training Mamba. We provide the detailed training configurations in Appendix A.3
Following the sctup Soction 4.2 we evaluate the language modeling perplexity of 100 tokens using the PG-19 dataset, conditioned on varying amounts of context. We report the results in Table Ø For all context. lengths, the fine-tuned Simba model consistently outperforms the training-free free Simba model, suggesting that fine-tuning can further improve the performance of sparsified SSMs.
5 Conclusion
We propose Simba, which sparsifies pre-trained SSMs into a hicrarchical form through token pruning. Simba outperforms Mamba with the same number of FLOPS in both accuracy on downstream NL.P benchmarks and language modeling perplexity. Additionally, our pruning scheme creates highways in the upper layers, enhancing length extrapolation for long sequences and facilitating the information flow across distant tokens. We hope Simba inspires a broad community, including state-space models, sparse and efficient networks, and classic recurrent networks with highways.
Limitations. Our paper mainly locused on applying Simba to the pre-trained Mamba without adjustment. However, token pruning incurs distribution shifts from the original models, and further fine-tuning could reduce this misalignment. We demonstrate that simple fine-tuning can improve the performance of small models within fixed computational budget in Section 4.4 However, more sophisticated fine-tuning scheme tailored for sparse SSMs could be investigated. Also, while sparsifying large model provides high performance at reduced inference-ting time computation, preparing large model would require a higher pre-training cost compared to using a smaller, dense model.
Broader Impact Statement Our paper studies sequence models, with broad applications such as language, audio, and video generation. As our method enhances the efficiency and efficacy of these models, it holds the potential to impact broader audience in generative AI. Hence, users of our method and sequence models should carefully read and follow the guidance from the commmityy (Bai et al. 2022
6 Acknowledgements
This work was conducted by Center for Applied Rescarch Artificial Intellgenno(CARAA]] grant funded by Defense Acquisition Program Administration(AAPA)) and Agency for Defense Development(ADD (UD230017TD)) and partially supported by Institute for Information & communications Technology Pro-
10

---

Published in Transactions on Machine Learning Research (03/2025)
motion(IITP)) grant funded by the Korca government(MSII) W.o.RS-2019-111900 Artificial Intelligence Graduate School Program (KATST); RS-2022-T1220955, Few-shot Learning of Causal Inference in Vision and Language for Decision Making).
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Lconi Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint W:Xii:2303088774 2023.
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. or preprint arXiv:2403.0153 2024.
Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, et al. Constitutional ai: Harmlessness from ai feedback. Xix preprint or arXin:2212.08077 2022.
Edward Beeching, Clémentine Fourrier, Nathan Habib, Sheon Han, Nathan Lambert, Nazneen Rajani, Omar Sanseviero, Lewis Tunstall, and Thomas Wolf. Open Ilm leaderboard. https://huggingfacee http:///gggngaaacc.cc////aaaaaea
fugginngPacei4//ppon, 11m Imm_eeaderboard 2023.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle OBrien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, 2023.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, el al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Daniel Bolya, Cheng-Yaue Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. International Conference on Learning Regresenlations, 2023.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXie preprint 15h3:1803.0543 2018.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Bolev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. orXiv preprint ar 4xX4::2402.19227, 2024.
Jeffrey L Elman. Finding structure in time. Cognitive science, 14(2):179 211, 1990.
Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He. Slowfast networks for vidco recognition. IEEE International Conference on Computer Vision, 2019.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. International Conference on Learning Representations, 20189.
Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. In International Conference on Learning Representations, 2023.
Trevor Gale, Erich Elsen, and Sara Hooker, The state of sparsity in deep neural networks, orXiv preprint 1rXi::1902.09570 2019.
11

---

Published in Transactions on Machine Learning Research (03/2025)
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa. Nabeshima, et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXic:2101.0000 2020.
Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Authony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Mucnnighoff, et al. A framework for fow-shot language model evaluation. Version v0. 0.1. Sepl, pp. 8, 2021.
Saurabh Goyal, Anamitra Roy Choudhury, Saurabh Raje, Venkatesan Chakaravarthy, Yogish Sabharwal, and Ashish Verma. Power-bert: Accelerating bert inference via progressive word-vector elimination. Tn International Conference on Machine Learning, 2020.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXi::2312.000 00752,, 2023.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Advances in Newrul Information Processing Syslems, 2022a.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. Tn International Conference on Learning Represenlations, 2022b.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. International Conference on Learning Representations, 2016.
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference Computer Vision and Pallern Recognilion, 2016.
Salah Hihi and Yoshua Bengio. Hierarchical recurrent neural networks for long-term dependencies. Advances in Neural Information Processing Systems, 1995.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural compulation, 0(8):1735-1788) 1997.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In IEEE Conference Computer Vision and Pattern Recognition, 2017.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying. ar preprint urXiv:2402.01034 2024.
Jean Kaddour. The minipile challenge for data-efficient language models. arXiv preprint arXie:2304.08442, 2023.
Rudolph Emil Kalman. A new approach lincar filtering and prediction problems. 1960.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, 2020.
Tsung- Wei Ke, Sangwoo Mo, and Yu Stella. Learning hicrarchical image segmentation for recognition and by recognition. In International Conference on Learning Representations, 2024.
Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers. In ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2022.
Jacho Loc, Sejun Park, Sangwoo Mo, Sungsoo Ahn, and Jinwoo Shin. Layer-adaptive sparsity for the magnitude-based pruming. International Conference on Learning Representations, 2021.
Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient. convnets. International Conference Learning Representations, 2017.
12

---

Published in Transactions on Machine Learning Research (03/2025)
Kunchang Li, Xinhao Li, Yi Wang, Yinan He, Yali Wang, Limin Wang, and Yu Qiao. Videomamba: State space model for efficient video understanding. arXiv preprint arx:::2433.069 2024a.
Shufan Li, Harkanwar Singh, and Aditya Grover. Mamba-nd: Selective state space modeling for multi- dimensional data. arXiv preprint orXie:200..0 05892, 2024b.
Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gonzalez. Train big, then compress: Rethinking model size for efficient training and inference of transformers. In International Conference on Machine Learning, 2020.
Youwei Liang, Chongjian Ge, Zhan Tong, Yibing Song, Jue Wang, and Pengtao Xie. Not all patches are what you need: Expediting vision transformers via. token reorganizations. arXie preprint Xiu:2202.07800 2022.
Opher Licher, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmodigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: hybrid transformer-mambl language model. arXiv preprint arXiv:2408.1100 19887, 2024.
Xiangcheng Liu, Tianyi Wu, and Guodong Guo. Adaptive sparse vit: Towards learnable adaptive token pruning by fully exploiting self-attention. arXiv preprint arXiv:2209.1385 2022.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In International Conference on Learning Representations, 2023.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than onc? Advances in Neural Information Processing Systems, 2019.
Antonio Orvieto, Samuel Smith, Albert Gu, Amushan Fernando, Caglar Gulcchre, Razvan Pascamu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, 2023.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prodiction requiring broad discourse context. arXiv preprinl Xim: 1606. 00000, 2016.
Sejum Park, Jaeho Lee, Sangwoo Mo, and Jinwoo Shin. Lookahead: A fir-tighted alternative of magnitude- based pruning. International Conference on Learning Representations, 2020.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Mattco Grella, Kranthi Kiran GV, al. Rwkv: Reinventing rns for the transformer cra. arXiv preprinl ar Xiv:22300 13048, 2023.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. Tn International Conference on Machine Learning, 2023.
Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiscroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, et al. Mechanistic design and scaling of hybrid archilectures. Xiv preprind or Xr3:22403. 17844, 2024.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL httpss: //apssv///xvv.. org/abs/ abs/ 1911.05507Z
Olaf Ronneberger, Philipp Fischer, and Thomas Brox U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Compuler Assisted Intervention, 2015.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Commanications of the ACM, 64(9):99- 2021.
13

---

Published in Transactions on Machine Learning Research (03/2025)
Yuzhang Shang, Mu Cai, Bingxin Xu, Yong Jae Lee, and Yan Yan. Llava-prunerge: Adaptive token reduction for efficient large multimodal models. arXiv preprint arXi:24003 15388, 2024.
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. In International Conference Learning Representations, 2023.
Woomin Song, Seunghyuk Oh, Sangwoo Mo, Jaehyung Kim, Sukmin Yun, Jung- Woo Ha, and Jinwoo Shin. Hierarchical context merging: Better long context understanding for pre-traines llms, International Conference on Learning Representations, 2024.
Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. ar preprint 1XXiv:505.0000 2015.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive norwork: A successor to transformer for large language models. arXie preprint ar arxv:22077.066 2023.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Scbastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representalions, 2021.
Alexandre Terres-Legue. mamba.py: simple, hackable and efficient mamba implementation in pure pytorch and mlx., 2024. URL https://githhb. httpa///..m/mlammm////aaaaaaaaaa py
Ali Vardasbi, Telmo Pessoa Pires, Robin M Schmidt, and Stephan Peitz. State spaces aren't enough: Machine translation needs attention. arXiv preprint arXiv:2304.12776 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Yau-Shiar Wang, Hung-Yi Yi Lee, and Yun-Nung Chen. Tree transformer: Integrating tree structures into self-attention. Conference Empirical Methods in Natural Language Processing, 2019.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint Xrx:::1905. 07830, 2019.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint 24ii:2401.0911 2024.
Michael Zhu and Suyog Gupta. To prune, not to prune: exploring the efficacy of pruning for model compression. Xix preprint 1710::1710.01879 2017.
Julian Georg Zilly, Rupesh Kumar Srivastava, Jan Koutmi, and Jürgen Schmidhuber. Recurrent highway networks, In International Conference on Machine Learning, 2017.
14

---

Published in Transactions on Machine Learning Research (03/2025)
A Detailed setups
A.1 NLP benchmarks
The six benchmarks were chosen according to the experiment setup of (Gu & Dao 2023). The metrics are also selected accordingly, where we measure answer perplexity and accuracy for Lambada, accuracy for PIQA, ARc-Easy, and Winogrande, and normalized accuracy for HellaSwag and Aro-Challenge. All measurements use the LM evaluation harness from Eleuther (ao al. p.21p.
Here, we provide the details for downstream experiments in Section 4.1
Following the widely adopted practice (Beeching et. al. 2023 evaluate the model's downst performance using few-shot prompts. Following the setup in the open LLM leaderboard (Beeching et al. 2023), we evaluate the model with 10-shot prompts for HellaSwag, 25-shot prompts for ARC-Easy and Arc-Challenge, and 5-shot prompts for WinoGrande. For benchmarks not considered the leaderboard, we evaluate the models with 5-shot prompts.
For the visualization in Figure we measure FLOPs separately for each benchmark and each model. First, measure each benchmark's ' mean prompt and answer length, shown in Table 3 Then, measure FL.PPs by forwarding an input that matches the mean lengths.
Table 3: Downstream task details. We provide the details of our downstream evaluations in Section 4.1 including the number of few-shot prompts, average prompt lengths, average answer lengths, and metrics used for evaluation.
Few-shot Prompts
Prompt Length (avg.)
Answer Length (avg.)
Dataset
Metrics
Lambada Paperno et al. 2016
5
507.44
1.47
ppl./ acc.
HellaSwag Zellers et al. PIQA Bisk et al. 2020
2019
10
877.25
29.86
acc norm
5
229.53
22.82
acc.
ARC-Easy [Clark et al.
2018
25
913.81
5.00
acc.
ARC-Challenge Clark al.
2018
25
913.81
5.00
WinoGrande (Sakaguchi et al.
202
5
143.17
5.67
acc.
A.2 Information flow visualization
Here, we provide the details for our information flow visualization experiments in Figure 5 We measure the normalized token influence score s(t) /tt/////lll 12 across all documents from the PG-19 test set, truncated at 1000 tokens. For all samples, we gather the influence score into five bins according to the position of the tokens. We report the average influence scores for each bin.
A.3 Fine-tuning with Simba
We provide the details for our fine-tuning experiments in Section 4.4
Loss design. The language modeling loss cannot be dircetly applied to sparsified models bocause the output logit shape does not match the label shape due to pruning. We apply the language modeling loss only to the available output logits, which are trained to predict the next token in the input sequence. We also add the standard language modeling loss computed using dense forwarding for stable training.
Training. We train the model for 400 steps on the MiniPile dataset (Kaddour 2023), a subset of the Pile dataset (Gao et al. 2020) with similar data distribution. We use AdamW optimizer with a learning rate of 5c-5. We schodule the lcarning rate with a lincar warmup for 10% of the total training steps and cosine learning rate decay for the remaining steps. We randomly select the pruning ratio between 0% and 90% for each sample.
15

---

Published in Transactions on Machine Learning Research (03/2025)
B Additional results
The tables below present the full values reported in our experiments.
B.1
Detailed results for downstream evaluations
Table 4: The full results table corresponding to Figure
Bold denotes the best results.
FLOPs
Lbd.
Lbd.
HS
PIQA
Are-E
Arc-C
WG
Avg.
Model
Scale
(x1e12)
perp. (4)
acc. (t)
acc. (1)
acc. (t)
acc. (1)
anc. (t)
acc. (f)
8cc. (1)
Simba (ours)
130m
0.23
31.63
36.28
35.29
64.64
49.12
24.66
49.96
43.32
Pythia
160m
0.66
10a
22.21
30.42
62.46
45.29
22.95
50.36
38.95
Mamba
130m
0.53
35.81
31.84
34.88
64.36
50.76
23.72
49.80
42,56
Simba (ours)
370m
0.75
16.79
45.10
46.72
69.59
58.84
28.07
55.41
50.62
Pythia
410m
1.86
25.65
22.22
40.95
68.06
56.73
26.45
53.51
47,32
Mamba
370m
1.51
12.34
46.61
46.46
69.59
59.05
27.47
55.80
50.83
Simba (ours)
790m
1.67
10.10
51.84
54.82
72.25
65.40
32.42
58.17
55.82
Pythia
lb
4.27
16.01
42.27
47.92
71.44
60.65
29.61
52.88
50.80
Mamba
790m
3.24
9.44
51.33
55.12
72.47
65.49
32.76
57.54
55.78
Simba (ours)
1.4b
2.95
8.41
55.61
58.61
73.34
65.32
34.47
58.72
57.68
Pythia
1.4b
6.19
10.35
49.46
53.54
70.89
65.32
33.02
57.38
54.94
Mamba
1.4b
5,60
7.31
56.43
59.60
74.05
66.50
36.18
60.14
58.82
Simba (ours)
2.8b
6.07
6.29
60.06
65.70
75.46
71.21
38.99
63.61
62.51
Pythia
2.8b
12.21
7.69
54.93
60.85
74.37
68.10
36.26
60.93
59.24
Mamba
2.8b
11.31
5.80
60.85
66.73
76.12
71.97
39.93
63.69
63.21
B.2
Detailed results for perplexity evaluations
Table
The full results table corresponding to Figure
Bold denotes the best results.
FLOPs
Within Context 1k
Extrapolation
Model
Scale
(xle12)
0.5k
1.5k
2.5k
3k
3.5k
4k
Simba (ours)
130m
0.21
3.060
3.006
2.986
2.974
2.961
2.956
2.948
2.942
Pythia
160m
0.60
3.166
3.134
3.128
3.120
3.195
6.333
7.883
7.986
Mamba
130m
0.48
2.943
2.917
2.907
2.899
2.899
2.900
2.901
2.905
Simba (ours)
370m
0.68
2.727
2.686
2.662
2.650
2.643
2.633
2.630
2.625
Pythia
410m
1.69
2.750
2.713
2.702
2.691
2.923
7.611
8.918
9.314
Mamba
370m
1.38
2.645
2.616
2.607
2,600
2.599
2.601
2,604
2.617
Simba (ours)
790m
1.52
2.543
2.509
2.488
2.477
2.474
2.470
2.468
2.471
Pythia
1b
3.88
2.607
2.573
2.560
2.559
5.665
6.582
6.817
7.003
Mamba
790m
2.94
2.486
2.457
2.445
2.437
2,436
2.443
2.458
2.497
Simba (ours)
1.4b
2.68
2.495
2.466
2.443
2.422
2.413
2.404
2.396
2.394
Pythia
1.4b
5.63
2.513
2.474
2.465
2,456
3.003
6.696
7.321
7.775
Mamba
1.4b
5.09
2,389
2.363
2.354
2.348
2.347
2.352
2.358
2.372
Simba (ours)
2.8b
5.52
2.380
2.330
2.317
2.304
2.294
2.289
2.279
2.284
Pythia
2,8b
11.11
2,382
2.345
2.332
2.326
6.689
7.275
7.266
7.244
Manba
2.8b
10.28
2.281
2.254
2.242
2.235
2.236
2.254
2.314
2.547
16

---

Published in Transactions on Machine Learning Research (03/2025)
C Additional information
C.1 Compute resources
All experiments are done on RTX-3090 or RTX-208S0 GPUs. All FLOPs reported in the paper indicate the computation required for running & single forward pass, so the number of data samples must be multiplied calculate the amount of computation required for the full experiments. We further acknowledge that the full research project required more computing than the experiments reported in the paper, including the preliminary experiments and failed experiments.
C.2 License for existing assets
Here, we provide the license for all datasets and models used in our experiments. Apache 2.0 license is applied for the pretrained Mamba models, Pythia models, PG-19 dataset, and WinoGrande dataset. CC BY 4.0 license is applied for the Lambada dataset. CC BY-SA 4.0 license is applied for ARC- RRC-Challenge and ARC-Easy datascts. MIT license applied for the PIQA dataset, HellaSwag datasct, and MiniPile datasct.
D Further Discussion and Analysis
D.1 Inference speed analysis
Throughout the paper, we have used FLOPs as the primary metric for evaluating efficiency. In this section, we present actual inference time measurements demonstrate that Simba achieves tangible real-world speedups. Specifically, we compare the inference time of Simba and Mamba models at two scales: 790M and 2.8B parameters. To ensure fair comparison, base evaluation the Py implementation of Mamba's parallel scan operation (Torres- Leguei, 2024), with optimizations to eliminate redundant operalions involved in score computation. These modifications allow us to measure the relative gains provided by Simba, and we expect comparable improvements with bardware-optimized implementations, given appropriate optimizations. The results are summarized Table which reports the average inference time for both Mamba and Simba across various context lengths.
Table 6: Inference time for various input lengths. All measurements are (aken on RTX 2080 Ti GPUs. We also report the percentage of the speedup.
Context Length
Model
Scale
lk
2k
3k
4k
Mamba
790m
0.683
1.356
2.438
2.709
Simba
790m
0.598
(14.2% speedup)
1.115 (21.6% speedup) 1.724 (41.4% speedup) 2.164 (25.2% speedup)
Mamba
2.8b
1.608
3.294
5.819
6.469
Simba.
2.8b
1.070 (50.3% speedup) 2.041 (61.4% speedup) 3.217 (80.9% speedup) 4.023 (60.8% speedup)
The results indicate that Simba delivers substantial real-world efficiency improvements, achieving up to 80% speedup for Mamba-2.8B models. We believe that these gains could be further enhanced by leveraging existing techniques, such as performing token pruning at specific layers pt all 2022 This represents a promising direction for future exploration.
rather than at every layer (Liang
D.2 Application to chunked forwarding
While originally designed for standard prefill scenarios where all tokens are forwarded at once, Simba can be scamlessly adapted to chunked prefill scenarios. In this context, long input sequences are divided into fixed-size chunks for sequential processing. To maintain computational efficiency, token pruning can be performed locally within each chunk.
17

---

Published in Transactions on Machine Learning Research (03/2025)
This can be done by caleulating token eviction scores based the influence of tokens on the last token of each chunk, instead of the last token of the entire input sequence. This localized scoring approach ensures that pruming decisions remain contextually relevant within each chunk, while still preserving the overall efficiency benefits of sparsification. By adopting this method, Simba can achieve similar efficiency gains in chunked prefill scenarios as standard prefill settings.
D.3 Analysis for performance improvement under same scale
Table 7: Performance on downstream tasks. We evaluate Simba and Mamba models on 6 NLP benchmarks (Lambada, HellaSwag, PIQA, Arc-Easy, Are- Are-Challenge, and WinoGrande) and report the averaged accuracy. We also report the performance difference between the two approaches with the same model scale. State size corresponds the size of the state tensor corresponding single token, at a single SSM layer. All experiment setup follows Table
Model Scale
130m
370m
790m
1.4b
2.8b
State Size
24k
32k
48k
64k
80k
Mamba
42.56
50.83
55.78
58.82
63.21
Simba
42.75
51.01
55.80
58.21
62.86
Difference
+0.19
+0.18
+0.02
-0.60
-0.35
While token pruning is often associated with a decrease in performance compared to dense models, we observed that moderate pruning can enhance performance under certain conditions in Section 4.3 In this section, we conduct further analysis on this phenomenon.
In Table 7 we compare the averaged accuracy of Mamba and Simba models across 6 NLP benchmarks considered throughout the paper. Noticably, the performance gains are more evident for smaller models with smaller state sizes. We hypothesize that this occurs because smaller state sizes have a limited capacity to retain information and are more prone to forgetting in dense recurrence operations. By introducing highway-like connections through our sparsification method in the upper layers, Simba could mitigate these issues by improving the information flow across distant tokens.
While these results are promising, believe that further exploration in this direction will be essential to fully understand the underlying mechanisms and to optimize the use of token pruning across a broader range of settings.
D.4 Comparison with Transformers
Transformers excel at modeling long-range dependencies, as their atlention mechanism gives each token direct access to all previous tokens. However, this capability comes at a high computational cost, with complexity growing quadratically with input length.
In contrast, SSMs process sequences with linear complexity by compressing all prior information into single state and updating it sequentially. While efficient, this design restricts information flow, tokens only access the immediate previous state. Over time, dense recurrence operations degrade earlier information, leading to a loss of long-range dependencies.
Simba mitigates the information flow issue in SSMs by sparsifying tokens in upper layers, reducing the number of recurrence steps. This approach alleviates information degradation while preserving the computational efficiency of SSMs. By balancing efficiency and improved long-range information flow, Simba narrows the gap between SSMs and Transformers.
18

---

Published in Transactions on Machine Learning Research (03/2025)
D.5 Potential failure cases
One potential failure case for Simba arises in scenarios where tokens containing essential information are prumed, leading to loss of critical context. While this issue would be less severe in comparison to Transformer- based models duc to the inherently compressod nature of SSMs, where information from pruned tokens is retained in the states of subsequent tokens, it may still impact the models ability to effectively utilize that information. This limitation could result performance degradation compared the dense model, particularly in tasks requiring precise retention of long-range dependencies complex reasoning.
19


Reviews:
[{"review_id": "kDL9DFeYOa", "reviewer": "Reviewer_QBdD", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nThe paper clearly introduces a hypothesis, verifies and improves it. The hypothesis is that fewer tokens are necessary in the upper layers because the information there is less influenced by the last new token, which is clearly show in Section 3.1. A reasonable solution is proposed, and it is verified that this solution actually does what is intended (results in Figure 4 and 5), apart from only on NLP benchmarks. This thorough approach is highly appreciated.\n\n**Weaknesses**\n\nThere are a few smaller things that I think can improve the paper further.\n\n* The paper is concerned about a fixed computational budget, but if I understand everything correctly, this is only considering the computational cost of inference. The training cost of a Simba model with the same inference cost as a Mamba model still required a larger training cost. I think this could be made more explicit when talking about the computation cost throughout the paper. \n* The scores in Table 1 are not that different. As far as I could tell, the input data of the pruning can be the only source of variation in the pruning. It is unclear which data and how much is used for the pruning. Both may be a source of variability in the pruning and eventual results. It implies some important questions: is the pruning asymptotic and are the same filters pruned if enough data is used, or is there always some variability? Does the data source that is used matter? Is it important that the data that is used to prune is close to the eventual test data? If possible, it wouldn't be bad to provide standard deviations in Table 1. I understand that it is not the goal of this method to have higher accuracy than Mamba, but right now this is claimed in the paper and given that there likely is some variability, this is not sufficiently backed up.\n\n\n\n* In the introduction there is the reasoning step that more redundancy is necessarily a result of more global information. I don't think this is the only reason that redundancy can happen, and it is not a necessary argument for the method. I think simply observing that there is more redundancy is enough to motivate the method, with speculating on the 'kind' of information. \n* SSMs and especially Mamba are proposed as an less computational complex alternative for transformer models. It is not absolutely necessary, but a comparison between the downsides and benefits of SSMs and transformers and how Samba possibly resolves these would be interesting.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose a method to sparsify state space models (SSM), in particular a recent improvement 'Mamba'. This has two benefits. One the one side it reduces the computational complexity of SSMs further. On the other side, it allows better performance on sequences with long contexts, as they are more efficiently propagated through the network. The sparsification that is proposed is hierarchical, it prunes more weights 'upper' layers of the network, as they are shown to be more similar to each other than the lower layers. The filters that are pruned are selected based on how much the output of a model changes when they are removed. Experiments on various NLP benchmarks are performed, and under the same computational budget, the proposed Simba outperforms Mamba. Additional experiments show that the fewer tokens in the upper layers indeed improve the information flow between distant input tokens. A final ablation study verifies that the same results are not reached when using random or uniform pruning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Necessary** \n\n* A clearer indication that the computational cost that is concerned only relates to inference, and that training costs are higher, should be included.\n* A clearer discussion of the data that is used to prune the network is necessary (but not all the ablations I proposed, although they would improve the paper). \n\n** Optional **\n\nEverything else I mentioned in the above weakness section.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "LAnHBFnzJF", "reviewer": "Reviewer_hcUQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The paper is well-written.\n2. The motivation of Simba is very clear.\n3. It shows better performance when using same flops in comparison to Mamba. In addition, by dropping redundant information during inference, it even increase performance despite using \n\nWeaknesses:\n\n1. Efficiency: It seems this method can only reduce prefilling time, while during generation we cannot drop tokens as we need to compute the current state to predict the next token. Moreover, when it comes to prefilling, we will compute it in a chunkwise manner. Based on the current implementation, we will only prune tokens in each chunk, which should be computed in parallel, therefore it is not clear about the practical improvement in training speed by using less FLOPs. It is difficult for us to completely prune a whole block based on the current implementation explained in the paper. Based on these, I am concerning about the practical speedup and suitable settings.\n2. Performance: It is very impressive that Simba can improve performance from Mamba when using less FLOPs, however, the authors do not provide a clear explanation of this part. Additionally, I would expect the authors to mainly focus this improvement in the motivation.\n3. Experiments: In Table 2, why Mamba and Simba with different scales are compared. This is a little strange.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this work, the authors propose Simba, which is  a simple yet effective method to sparsify SSMs through token pruning. Here, the main motivation is that since SSMs generate the next states based on the immediate previous states and current input, compressing all information into one state will result into redundancy. Therefore, the authors will prune more tokens for upper layers, which  behave like highways. To determine which token to be pruned, the influence function is used, which measures the importance of token $x_t$ by subtracting the final token output with and without this token. Experimental results show the effectiveness of proposed method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Add results about improvement in practical efficiency to show how much Simba can accelerate prefilling\n2. Discuss about the chunkwise prefilling, and clarify how to prune tokens within and among chunks.\n3. Add analyses for the better performance of Simba compared to Mamba even using less FLOPs.\n4. Update Table 2 with comparison among models with the same scale.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "VQqMcqDorG", "reviewer": "Reviewer_9kSj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\nThe key strength of this work lies in its direct and intuitive approach. While Transformers rely on KV Cache mechanisms that make skip connections challenging to implement, SSMs can naturally leverage their hidden states to create such pathways. This makes the proposed sparsification approach particularly well-suited for the SSM architecture. The idea itself is straightforward yet effective, demonstrating clear improvements in computational efficiency as measured by FLOPs.\n### Weaknesses\nThe primary weakness centers on the reliance on FLOPs as the main efficiency metric. As someone with experience in pre-training large models, I find this concerning because FLOPs improvements often fail to translate into actual training speed gains. Many factors beyond raw FLOPs affect real-world training performance, including memory access patterns, hardware utilization, and implementation efficiency. The paper would be strengthened by demonstrating whether these theoretical FLOPs advantages can actually deliver practical speedups in training or inference scenarios.\n\nFurthermore, while the authors propose a token pruning criterion based on measuring global influence through local recurrences, they don't adequately address potential failure cases. For instance, in scenarios where critical information is distributed across multiple tokens rather than concentrated in individual \"important\" tokens, their pruning approach might discard essential contextual information. This could be particularly problematic for tasks requiring complex reasoning or long-range dependencies.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Simba, a hierarchical sparsification method for SSM. The main contributions include a token pruning approach that increasingly sparsifies upper layers while preserving lower layers, creating a trapezoidal network structure. The authors develop a new token importance scoring criterion specifically designed for SSMs that measures global token influence by accumulating local recurrences. The work demonstrates that the resulting sparse upper layers act as \"highways\" that improve information flow across long sequences. Through extensive experiments, the authors show empirical improvements over baseline Mamba models while using equivalent computational resources across multiple NLP tasks and language modeling benchmarks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The authors should demonstrate whether their FLOPs improvements translate to actual training or inference time speedups. This is particularly important for pre-training scenarios where theoretical efficiency gains don't always materialize in practice.", "broader_impact_concerns": "N/A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces Simba, a hierarchical token pruning method for state-space models (SSMs) that sparsifies upper layers more aggressively, enhancing both computational efficiency and long-range information flow. While initially focused on FLOPs as an efficiency metric, the authors convincingly addressed reviewer concerns by demonstrating substantial real-world inference speedups and clarifying increased training costs and adaptability to chunked pre-filling. They also added a relevant discussion regarding failure scenarios where information might be more evenly distributed. The paper is recommended for acceptance as is; it offers a novel, well-supported, and empirically validated approach, making a strong contribution to sequence modeling.
Decision: Accept (poster)