Los Modelos de Espacio de Estados Esparsificados son Redes de Carreteras 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

Resumen

Los modelos de espacio de estados (SSMs, por sus siglas en inglés) ofrecen una arquitectura prometedora para el modelado de secuencias, proporcionando una alternativa a los Transformers al reemplazar la costosa autoatención con recurrencias lineales. En este artículo, proponemos un truco simple pero efectivo para mejorar los SSMs dentro de presupuestos computacionales dados mediante su esparsificación. Nuestra intuición es que los tokens en los SSMs son altamente redundantes debido a las actualizaciones recurrentes graduales, y las operaciones de recurrencia densa bloquean la entrega de información pasada. En particular, observamos que las capas superiores de los SSMs tienden a ser más redundantes, ya que codifican información global, mientras que las capas inferiores codifican información local. Motivados por esto, introducimos Simba, un método de esparsificación jerárquica para SSMs basado en la poda de tokens. Simba esparsifica más las capas superiores que las inferiores, incentivando a que las capas superiores se comporten como autopistas. Para lograr esto, proponemos un nuevo criterio de poda de tokens para SSMs, que mide el impacto global de los tokens en la salida final acumulando recurrencias locales. Demostramos que Simba supera al modelo de referencia, Mamba, con el mismo número de FLOPS en diversas tareas de lenguaje natural. Además, ilustramos el efecto de las autopistas, mostrando que Simba no solo mejora la eficiencia, sino que también optimiza el flujo de información a lo largo de secuencias largas. El código está disponible en 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.