Разреженные модели пространства состояний представляют собой эффективные шоссейные сети.

Sparsified State-Space Models are Efficient Highway Networks

May 27, 2025

Авторы: Woomin Song, Jihoon Tack, Sangwoo Mo, Seunghyuk Oh, Jinwoo Shin

cs.AI

Аннотация

Модели пространства состояний (SSM) представляют собой перспективную архитектуру для моделирования последовательностей, предлагая альтернативу трансформерам за счет замены дорогостоящего самовнимания линейными рекуррентными операциями. В данной статье мы предлагаем простой, но эффективный прием для улучшения SSM в рамках заданных вычислительных бюджетов путем их разрежения. Наша интуиция заключается в том, что токены в SSM обладают высокой избыточностью из-за постепенных рекуррентных обновлений, а плотные рекуррентные операции блокируют передачу информации из прошлого. В частности, мы наблюдаем, что верхние слои SSM, как правило, более избыточны, так как кодируют глобальную информацию, в то время как нижние слои кодируют локальную информацию. Мотивируясь этим, мы представляем Simba — иерархический метод разрежения для SSM, основанный на обрезке токенов. Simba сильнее разрежает верхние слои по сравнению с нижними, побуждая верхние слои вести себя как "магистрали". Для достижения этого мы предлагаем новый критерий обрезки токенов для SSM, измеряющий глобальное влияние токенов на конечный выход путем накопления локальных рекуррентных операций. Мы демонстрируем, что Simba превосходит базовую модель Mamba при одинаковом количестве операций с плавающей запятой (FLOPS) в различных задачах обработки естественного языка. Более того, мы иллюстрируем эффект "магистралей", показывая, что Simba не только повышает эффективность, но и улучшает поток информации в длинных последовательностях. Код доступен по адресу 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.