Sparsified State-Space Models are Efficient Highway Networks

2025-05-27

AI Summary

State-space models (SSMs) reduce computational cost by replacing self-attention with linear recurrences, yet their layer-wise recursive token updates induce severe redundancy—especially in deeper layers encoding global context—impeding long-range information propagation. This paper introduces Simba, the first hierarchical sparsification framework for SSMs. Its core innovation is a global influence-aware token pruning criterion, grounded in recursive cumulative sensitivity analysis, which automatically sparsifies higher layers into “information highways” while preserving fine-grained local details in lower layers. Under fixed FLOPS, Simba significantly improves efficiency for long-sequence modeling. Empirically, it consistently outperforms Mamba across multiple NLP benchmarks, achieving higher accuracy and superior long-range dependency modeling capability.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Summarization
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Enhancing SSMs efficiency via hierarchical sparsification
• Reducing token redundancy in upper SSM layers
• Improving long-sequence information flow with token pruning

Innovation

• Sparsifies SSMs via hierarchical token pruning
• Introduces Simba for efficient highway networks
• Measures token impact via local recurrences