I Modelli Spaziali di Stato Sparsificati sono Reti Autostradali Efficienti

Sparsified State-Space Models are Efficient Highway Networks

May 27, 2025

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

cs.AI

Abstract

I modelli spazio-stato (SSM) offrono un'architettura promettente per la modellazione di sequenze, fornendo un'alternativa ai Transformer sostituendo il costoso self-attention con ricorrenze lineari. In questo articolo, proponiamo un semplice ma efficace stratagemma per migliorare gli SSM entro determinati budget computazionali attraverso la loro sparsificazione. La nostra intuizione è che i token negli SSM siano altamente ridondanti a causa degli aggiornamenti ricorrenti graduali, e che le operazioni di ricorrenza dense ostacolino il passaggio delle informazioni passate. In particolare, osserviamo che gli strati superiori degli SSM tendono a essere più ridondanti poiché codificano informazioni globali, mentre gli strati inferiori codificano informazioni locali. Motivati da ciò, introduciamo Simba, un metodo di sparsificazione gerarchica per gli SSM basato sulla potatura dei token. Simba sparsifica maggiormente gli strati superiori rispetto a quelli inferiori, incoraggiando gli strati superiori a comportarsi come autostrade. Per raggiungere questo obiettivo, proponiamo un nuovo criterio di potatura dei token per gli SSM, che misura l'impatto globale dei token sull'output finale accumulando le ricorrenze locali. Dimostriamo che Simba supera il modello di riferimento, Mamba, con lo stesso numero di FLOPS in vari compiti di elaborazione del linguaggio naturale. Inoltre, illustriamo l'effetto delle autostrade, mostrando che Simba non solo migliora l'efficienza, ma anche il flusso di informazioni attraverso sequenze lunghe. Il codice è disponibile all'indirizzo 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.