Sparsified State-Space Models are Efficient Highway Networks
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.
6847486d3ec10bdd8ab4dd07
Woomin Song
6847486d3ec10bdd8ab4dd08
Jihoon Tack
6847486d3ec10bdd8ab4dd09
Sangwoo Mo
6847486d3ec10bdd8ab4dd0a
Seunghyuk Oh
6847486d3ec10bdd8ab4dd0b
Jinwoo Shin
Sparsified State-Space Models are Efficient Highway Networks
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.
Simba, a hierarchical sparsification method for state-space models, enhances efficiency and information flow in natural language tasks by pruning tokens more aggressively in upper layers.
state-space models sequence modeling self-attention linear recurrences sparsification token pruning hierarchical sparsification Simba global information local information highways FLOPS natural language tasks