Sparsified State-Space Models are Efficient Highway Networks

TL;DR

This work addresses the efficiency of sequence modeling with state-space models (SSMs) by introducing Simba, a hierarchical sparsification method that prunes tokens within pre-trained SSMs under a fixed compute budget. Simba computes a global token importance score $s(t)=\max(\Delta y_T(t))$, where $\Delta y_T(t)$ captures the influence of token $x_t$ on the final output via the SSM recurrence, and prunes tokens to create a trapezoidal network with sparse upper layers that function as highways. The approach is training-free and uses a linear pruning schedule preserving $10\%$ of tokens at the final layer; empirically it outperforms the baseline Mamba at the same FLOPs across 6 NLP benchmarks and improves language modeling perplexity on PG-19, with enhanced long-context information flow. Overall, Simba demonstrates how structured sparsification can yield both efficiency gains and improved information propagation in SSMs, approaching the performance of Transformers while maintaining linear-time computation.

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.

Figures (6)

• Figure 1: Simba: hierarchical sparsification of SSMs via token pruning. (a) We found that tokens in state-space models (SSMs) are highly redundant, especially in the upper layers. Motivated by this observation, we propose hierarchically 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.8b 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 setups in \ref{['sec:exp']}. Simba outperforms both models in accuracy and perplexity.
• Figure 2: Token redundancy of SSMs has a hierarchical structure. We measure the cosine similarity between adjacent tokens of the Mamba-2.8b model across layers, averaged over documents from the PG-19 test dataset. The tokens are highly redundant, especially in the upper layers.
• 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.
• Figure 4: Language modeling ability. We measure the FLOPs-perplexity curves on the PG-19 test dataset. Simba models are compared against Mamba models that use 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.
• Figure 5: Information flow across layers. We visualize the information flow using the normalized token influence score. We compare Mamba-2.8b 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 early tokens.