arXiv:2505.20698cs.LG2025-05中稿 · TMLR 2025

通过分层稀疏化提升状态空间模型效率,让高层像高速路一样快速传递信息。

Sparsified State-Space Models are Efficient Highway Networks

  • 按层级稀疏化:高层比低层删更多冗余令牌
  • 同计算量下超越Mamba,在多语言任务上表现更优
  • 新剪枝准则衡量令牌全局影响,改善长序列信息流动

状态空间模型(SSMs)为序列建模提供了一种替代Transformer的方案,用线性递推代替昂贵的自注意力机制。本文提出一种简单有效的稀疏化方法——Simba,旨在不增加计算成本的前提下提升SSM性能。我们观察到:由于渐进式递推更新,SSM中的令牌高度冗余,密集递推会阻碍历史信息传递。特别是高层编码全局信息,冗余度更高;而低层主要处理局部信息。因此,我们设计了基于令牌剪枝的层次化稀疏策略,对高层进行更激进的稀疏化,使高层行为类似高速公路。为此,提出一种新型令牌剪枝标准,通过累积局部递推来评估令牌对最终输出的全局影响。实验表明,相同浮点运算量(FLOPS)下,Simba在多项自然语言任务中优于基线模型Mamba。此外,验证了高速公路效应:不仅提升效率,还显著改善长序列的信息流动。代码已公开于https://github.com/woominsong/Simba。

原文摘要 · Abstract (English)

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.

状态空间模型稀疏化高效序列建模信息流动

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。