Gesparseerde State-Space Modellen zijn efficiënte Highway Netwerken

Sparsified State-Space Models are Efficient Highway Networks

May 27, 2025

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

cs.AI

Samenvatting

State-space modellen (SSMs) bieden een veelbelovende architectuur voor sequentiemodellering en vormen een alternatief voor Transformers door dure zelf-attentie te vervangen met lineaire recurrenties. In dit artikel introduceren we een eenvoudige maar effectieve truc om SSMs te verbeteren binnen gegeven rekenbudgetten door ze te versparsen. Onze intuïtie is dat tokens in SSMs sterk redundant zijn vanwege geleidelijke recurrentie-updates, en dat dichte recurrentie-operaties de overdracht van informatie uit het verleden blokkeren. In het bijzonder merken we op dat de bovenste lagen van SSMs vaak meer redundant zijn omdat ze globale informatie coderen, terwijl de onderste lagen lokale informatie coderen. Gemotiveerd door dit inzicht introduceren we Simba, een hiërarchische versparsingsmethode voor SSMs gebaseerd op token pruning. Simba versparst de bovenste lagen meer dan de onderste lagen, waardoor de bovenste lagen zich meer als snelwegen gaan gedragen. Om dit te bereiken, stellen we een nieuw token pruning-criterium voor SSMs voor, dat de globale impact van tokens op de uiteindelijke uitvoer meet door lokale recurrenties op te tellen. We demonstreren dat Simba het baseline-model, Mamba, overtreft met hetzelfde aantal FLOPS in verschillende natuurlijke-taaltaken. Bovendien illustreren we het effect van snelwegen, waarbij we aantonen dat Simba niet alleen de efficiëntie verbetert, maar ook de informatiestroom over lange sequenties optimaliseert. Code is beschikbaar op 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.