Modelos de Espaço de Estados Esparsificados são Redes de Rodovias Eficientes

Sparsified State-Space Models are Efficient Highway Networks

May 27, 2025

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

cs.AI

Resumo

Modelos de espaço de estados (SSMs) oferecem uma arquitetura promissora para modelagem de sequências, fornecendo uma alternativa aos Transformers ao substituir a custosa auto-atenção por recorrências lineares. Neste artigo, propomos um truque simples, porém eficaz, para aprimorar SSMs dentro de orçamentos computacionais dados, esparsificando-os. Nossa intuição é que os tokens em SSMs são altamente redundantes devido a atualizações recorrentes graduais, e operações de recorrência densas bloqueiam a entrega de informações passadas. Em particular, observamos que as camadas superiores dos SSMs tendem a ser mais redundantes, pois codificam informações globais, enquanto as camadas inferiores codificam informações locais. Motivados por isso, introduzimos Simba, um método hierárquico de esparsificação para SSMs baseado em poda de tokens. O Simba esparsifica mais as camadas superiores do que as inferiores, incentivando que as camadas superiores se comportem como rodovias. Para alcançar isso, propomos um novo critério de poda de tokens para SSMs, medindo o impacto global dos tokens na saída final ao acumular recorrências locais. Demonstramos que o Simba supera o modelo de referência, Mamba, com o mesmo número de FLOPS em várias tarefas de linguagem natural. Além disso, ilustramos o efeito das rodovias, mostrando que o Simba não apenas aumenta a eficiência, mas também melhora o fluxo de informações em sequências longas. O código está disponível em 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.