Sparsifizierte State-Space-Modelle sind effiziente Highway-Netzwerke.

Sparsified State-Space Models are Efficient Highway Networks

May 27, 2025

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

cs.AI

Zusammenfassung

Zustandsraummodelle (State-Space Models, SSMs) bieten eine vielversprechende Architektur für die Sequenzmodellierung und stellen eine Alternative zu Transformern dar, indem sie den rechenintensiven Selbstaufmerksamkeitsmechanismus durch lineare Rekurrenzen ersetzen. In diesem Artikel schlagen wir einen einfachen, aber effektiven Trick vor, um SSMs innerhalb vorgegebener Rechenbudgets durch Sparsifizierung zu verbessern. Unsere Intuition ist, dass Token in SSMs aufgrund schrittweiser rekursiver Aktualisierungen stark redundant sind und dichte Rekurrenzoperationen die Übermittlung vergangener Informationen blockieren. Insbesondere beobachten wir, dass die oberen Schichten von SSMs tendenziell redundanter sind, da sie globale Informationen kodieren, während die unteren Schichten lokale Informationen kodieren. Motiviert durch diese Beobachtung führen wir Simba vor, eine hierarchische Sparsifizierungsmethode für SSMs basierend auf Token-Pruning. Simba sparsifiziert die oberen Schichten stärker als die unteren Schichten und fördert so, dass die oberen Schichten wie „Highways“ agieren. Um dies zu erreichen, schlagen wir ein neuartiges Token-Pruning-Kriterium für SSMs vor, das den globalen Einfluss von Token auf das endgültige Ergebnis durch die Akkumulation lokaler Rekurrenzen misst. Wir zeigen, dass Simba das Basismodell Mamba bei gleicher Anzahl von FLOPS in verschiedenen Aufgaben der natürlichen Sprachverarbeitung übertrifft. Darüber hinaus veranschaulichen wir die Wirkung der „Highways“ und zeigen, dass Simba nicht nur die Effizienz steigert, sondern auch den Informationsfluss über lange Sequenzen verbessert. Der Code ist verfügbar unter 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.