Sparsified State-Space Models are Efficient Highway Networks

TL;DR

This paper introduces Simba, a hierarchical sparsification method for state-space models that prunes tokens to create highway-like layers, improving efficiency and information flow in sequence modeling tasks.

Contribution

The paper proposes a novel token pruning criterion and hierarchical sparsification technique for SSMs, enhancing their efficiency and long-range information flow.

Findings

Simba outperforms baseline models with the same FLOPS in NLP tasks.

Sparsification creates highway-like layers that improve long sequence processing.

The method enhances both efficiency and information flow in SSMs.

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.