AIRE-Prune: Asymptotic Impulse-Response Energy for State Pruning in State Space Models

Resumen Técnico: AIRE-Prune

Planteamiento del Problema

Los Modelos de Espacio de Estados (SSM) profundos han demostrado eficacia en el modelado de datos secuenciales mediante la compresión de la historia de entrada en estados internos. Sin embargo, la dimensión del estado $n$ sigue siendo un factor principal de los costes computacionales y de memoria. Aunque las arquitecturas modernas suelen utilizar parametrizaciones diagonales o diagonalizables (por ejemplo, pilas multi-SISO o MIMO como S5) para garantizar la estabilidad y la eficiencia, estos modelos sufren frecuentemente de sobreparametrización. Los mecanismos existentes de post-entrenamiento para optimizar $n$ son limitados. Otros enfoques previos, como la poda de pesos no estructurada, no reducen directamente la dimensión del estado. Otros métodos estructurados, como el reciente LAST (Gwak et al., 2025), dependen de la ganancia en el dominio de la frecuencia en el peor de los casos ($H_\infty$) para calificar los estados. Aunque es efectivo, las medidas del peor caso pueden ser excesivamente conservadoras para cargas de trabajo típicas, ya que enfatizan la amplificación máxima que las tareas rara vez excitan, lo que puede conducir a ratios de compresión subóptimos.

Metodología: AIRE-Prune

Los autores presentan AIRE-Prune (Asymptotic Impulse-Response Energy for State PRUN(E)ing), un método de poda estructurado, de post-entrenamiento y adaptativo por capa, diseñado para SSMs diagonales (o diagonalizables). A diferencia de otros métodos que se centran en la ganancia del peor caso, AIRE-Prune alinea la poda con la energía de la respuesta al impulso asintótica, un criterio de caso típico que representa la energía de salida total que un estado aporta a lo largo de un horizonte de tiempo infinito.

Mecanismo Central

1. Clasificación Basada en Energía:
   Para una capa de SSM diagonal de tiempo discreto estable definida por $(\Lambda, B, C)$, donde $\Lambda = \text{diag}(\lambda_1, \dots, \lambda_n)$ y $|\lambda_i| < 1$, el método calcula la energía total de salida al cuadrado generada por una entrada de impulso unitario para cada estado (modo) $i$.
   El score de forma cerrada para el $i$-ésimo modo se deriva como:
   $$\text{EnergyScore}_{\text{local}}(x_i) = E_i = \frac{\|C_{:,i}\|_2^2 \|B_{i,:}\|_2^2}{1 - |\lambda_i|^2}$$
   Este score integra tres factores:

   • Controlabilidad: Representada por $\|B_{i,:}\|_2^2$ (facilidad de excitación).
   • Observabilidad: Representada por $\|C_{:,i}\|_2^2$ (facilidad de medición).
   • Amortiguamiento/Memoria: Representada por el polo $\lambda_i$; los estados con polos más cercanos al círculo unitario (decaimiento más lento) contribuyen más a la energía a largo plazo.

2. Normalización entre Capas:
   Dado que las escalas de energía pueden variar significativamente entre diferentes capas (por ejemplo, debido a las ganancias del codificador/decodificador), los autores proponen una estrategia de puntuación con prefijo normalizado para permitir la comparación global:

   • Dentro de cada capa $\ell$, los modos se ordenan por su energía local $E^{(\ell)}_i$ en orden descendente.
   • Se calcula una suma de prefijos $S^{(\ell)}_{(i)}$ para las energías ordenadas.
   • El score global se define como la relación de "tasa de riesgo" (hazard-rate):
     $$\text{AIREPrune}(x^{(\ell)}_{(i)}) = \frac{E^{(\ell)}_{(i)}}{S^{(\ell)}_{(i)} + \epsilon}$$
     Esta normalización asegura que los scores sean monótonamente no crecientes dentro de una capa, permitiendo que se aplique un único umbral global $\tau$ a todo el modelo.

3. Ejecución de la Poda:
   Dado un ratio de poda o umbral objetivo, el algoritmo retiene el prefijo más largo de estados en cada capa donde el score normalizado supera a $\tau$. Esto permite una poda adaptativa por capa donde algunas capas pueden ser fuertmente podadas o incluso eliminadas por completo si sus estados caen por debajo del umbral global, mientras que otras permanecen mayoritariamente intactas.

Contribuciones Clave

• Score de Energía Asintótica de Forma Cerrada: El artículo proporciona un score de forma cerrada derivado matemáticamente para cada modo aprendido basado en la energía de salida de horizonte infinito. Esto evita la necesidad de simulaciones numéricas costosas o re-entrenamiento.
• Extensión de la Truncación Modal: El trabajo extiende la clásica truncación modal (típicamente aplicada a sistemas lineales individuales) a pilas profundas de SSMs no lineales, preservando la estructura diagonal esencial para la eficiencia.
• Caso Típico vs. Peor Caso: Al cambiar la métrica de importancia de la ganancia $H_\infty$ del peor caso a la energía total de la respuesta al impulso, el método apunta al comportamiento agregado del sistema en lugar de a las amplificaciones pico raras.
• Agnóstico a la Arquitectura: El método es aplicable tanto a construcciones multi-SISO como a pilas MIMO (por ejemplo, estilo S5), así como a otras familias de SSM como S4D y Mamba.

Resultados Experimentales

Los autores evaluaron AIRE-Prune en el benchmark Long Range Arena (LRA) y el conjunto de datos Speech Commands utilizando backbones S5, S4D y Mamba.

• Eficiencia de Compresión: A través de diversos benchmarks, AIRE-Prune logró una poda de estado promedio del 60.8% con una caída de precisión promedio de solo 0.29% sin ningún re-entrenamiento.
• Comparación con Baselines:
  • En la suite LRA, AIRE-Prune superó significativamente a los baselines incluyendo Random, Uniform Magnitude, Global Magnitude, LAMP, Uniform $H_\infty$, Global $H_\infty$ y LAST.
  • Bajo una restricción estricta de $\le 1$ punto porcentual (pp) de pérdida de precisión, AIRE-Prune permitió ratios de poda sustancialmente mayores que el baseline más fuerte (LAST). Por ejemplo, en la tarea Path-X, AIRE-Prune podó de forma segura el 70% de los estados comparado con el 30% de LAST.
  • En la tarea Speech Commands, se podó el 45% de los estados con una pérdida de precisión insignificante (96.40% vs. 96.43% para el modelo completo).
• Eficiencia de Inferencia: La poda se tradujo directamente en ganancias a nivel de sistema. En una GPU NVIDIA H100, el método produjo aceleraciones de inferencia que oscilaron entre 1.23$\times$ y 2.86$\times$ y reducciones de parámetros del 19% al 65%, dependiendo de la tarea y el ratio de poda.
• Generalización: El método demostró un rendimiento robusto en arquitecturas S4D y Mamba, logiendo hasta un 90% de poda en la tarea Text para S4D y un 80% en la tarea Image para Mamba con una degradación de precisión mínima.

Significado y Reivindicaciones

El artículo afirma que AIRE-Prune revela una redundancia sustancial en los espacios de estado de los SSMs entrenados, la cual puede eliminarse sin re-entrenamiento. Al alinear la poda con la energía de respuesta asintótica en lugar de la ganancia del peor caso, el método logra una compresión agresiva con un coste insignificante para el rendimiento de la tarea.

Los autores destacan que el perfil de caída de precisión "en forma de escalón" observado con AIRE-Prune (donde la precisión permanece estable hasta que se alcanza un alto umbral de poda, y luego cae bruscamente) indica una separación fuerte entre los estados importantes e irrelevantes. Esto contrasta con los baselines que muestran una degradación suave, sugiriendo que sus métodos de puntuación mezclan estados importantes e irrelevantes.

Prácticamente, el método permite:

1. Presupuesto Global: Establecer un único presupuesto objetivo sin necesidad de ajuste por capa.
2. Reducción a Nivel de Bloque: La capacidad de podar capas enteras cuando sus estados caen por debajo del umbral, convirtiendo la escasez fina en reducciones estructuradas por bloques que son más amigables para el hardware.
3. Optimización Sin Entrenamiento: Un mecanismo de post-entrenamiento que no requiere entrenamiento adicional, lo que lo hace inmediatamente aplicable a modelos desplegados.

El trabajo concluye que AIRE-Prune ofrece un enfoque basado en la energía y con principios sólidos para la reducción del orden de modelos en SSMs profundos, superando las estrategias de poda actuales en términos de ratio de compresión y preservación de la precisión.