Les modèles d'état espace éparsifiés sont des réseaux autoroutiers efficaces.

Sparsified State-Space Models are Efficient Highway Networks

May 27, 2025

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

cs.AI

Résumé

Les modèles à espace d'états (SSMs) offrent une architecture prometteuse pour la modélisation de séquences, constituant une alternative aux Transformers en remplaçant l'auto-attention coûteuse par des récurrences linéaires. Dans cet article, nous proposons une astuce simple mais efficace pour améliorer les SSMs dans des budgets de calcul donnés en les rendant parcimonieux. Notre intuition est que les tokens dans les SSMs sont hautement redondants en raison des mises à jour récurrentes graduelles, et que les opérations de récurrence dense bloquent la transmission des informations passées. En particulier, nous observons que les couches supérieures des SSMs ont tendance à être plus redondantes car elles encodent des informations globales, tandis que les couches inférieures encodent des informations locales. Motivés par cela, nous introduisons Simba, une méthode de parcimonie hiérarchique pour les SSMs basée sur l'élagage de tokens. Simba rend les couches supérieures plus parcimonieuses que les couches inférieures, encourageant les couches supérieures à se comporter comme des autoroutes. Pour y parvenir, nous proposons un nouveau critère d'élagage de tokens pour les SSMs, mesurant l'impact global des tokens sur la sortie finale en accumulant les récurrences locales. Nous démontrons que Simba surpasse le modèle de référence, Mamba, avec le même nombre de FLOPS dans diverses tâches de traitement du langage naturel. De plus, nous illustrons l'effet des autoroutes, montrant que Simba améliore non seulement l'efficacité mais aussi le flux d'information à travers les longues séquences. Le code est disponible à l'adresse 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.