スパース化状態空間モデルは効率的なハイウェイネットワークである

Sparsified State-Space Models are Efficient Highway Networks

May 27, 2025

著者: Woomin Song, Jihoon Tack, Sangwoo Mo, Seunghyuk Oh, Jinwoo Shin

cs.AI

要旨

状態空間モデル（SSMs）は、シーケンスモデリングにおいて有望なアーキテクチャを提供し、高コストな自己注意機構を線形再帰に置き換えることで、Transformerに代わる選択肢を提示します。本論文では、与えられた計算予算内でSSMsを強化するためのシンプルかつ効果的な手法として、疎化を提案します。我々の直感は、SSMsにおけるトークンは漸進的な再帰的更新により高度に冗長であり、密な再帰操作が過去の情報の伝達を妨げるというものです。特に、SSMsの上位層はグローバルな情報を符号化するためより冗長であるのに対し、下位層はローカルな情報を符号化する傾向があることを観察しました。これに基づき、トークンプルーニングに基づくSSMsの階層的疎化手法であるSimbaを導入します。Simbaは上位層を下位層よりも多く疎化し、上位層がハイウェイのように振る舞うことを促します。これを実現するために、ローカルな再帰を累積することでトークンの最終出力に対するグローバルな影響を測定する、SSMs向けの新しいトークンプルーニング基準を提案します。Simbaが、同じFLOPSでベースラインモデルであるMambaを様々な自然言語タスクにおいて上回ることを実証します。さらに、ハイウェイの効果を説明し、Simbaが効率を向上させるだけでなく、長いシーケンス間の情報の流れを改善することを示します。コードはhttps://github.com/woominsong/Simbaで公開されています。

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.