Resumen Técnico: Un Método de Compresión de Modelos de Espacio de Estados Profundos usando Cota Superior del Error de Salida

Enunciado del Problema

Los Modelos de Espacio de Estados Profundos (Deep SSMs), como S5 y Mamba, han demostrado un rendimiento de vanguardia en el manejo de dependencias de largo alcance sin mecanismos de atención. Sin embargo, el despliegue de estos modelos a menudo requiere reducir su número de parámetros manteniendo el rendimiento. Las técnicas existentes de Reducción de Orden de Modelo (MOR) para Deep SSMs típicamente aplican compresión de forma independiente a cada bloque de espacio de estados lineal (por capa) basándose en criterios locales (por ejemplo, Truncamiento Balanceado o normas $H_2$).

Los autores identifican una limitación crítica en estos enfoques existentes: minimizar el error de aproximación de bloques lineales individuales no garantiza la minimización del error de salida global de todo el Deep SSM. Debido a que los Deep SSMs involucran capas en cascada con interacciones no lineales (específicamente salidas cuadráticas y Normalización de Capa), los errores locales pueden acumularse o propagarse de manera impredecible. En consecuencia, existe una falta de un marco teórico que vincule directamente las estrategias de reducción por capa con el rendimiento de salida global del Deep SSM.

Metodología

1. Marco Teórico: Sistemas LQO y Límites de Error

El artículo modela los Deep SSMs como una cascada de sistemas Lineales con Salida Cuadrática (LQO). Un sistema LQO se define por una actualización de espacio de estados lineal seguida de un mapa de salida cuadrático, lo que generaliza arquitecturas como S5 donde una activación cuadrática sigue a un bloque lineal.

Los autores derivan una cota superior demostrable sobre el error de salida $\ell_\infty$ entre un Deep SSM de orden completo y un Deep SSM de orden reducido.

• Insight Clave: El error de salida global está acotado por una suma ponderada de los errores de aproximación de la norma $h_2$ de los sistemas LQO individuales en cada capa.
• Sensibilidad por Capa: La cota derivada revela que la contribución del error de aproximación de una capa al error de salida total no es uniforme. Específicamente, los errores en las capas superficiales (al inicio de la red) tienen un impacto significativamente mayor en el error de salida final que los errores en capas más profundas. Esto se debe al efecto multiplicativo de las constantes de Lipschitz asociadas con la Normalización de Capa y la propagación de errores a través de las capas subsiguientes.

2. Formulación de Optimización

Basándose en la cota superior derivada, los autores formulan un problema de Reducción de Orden de Modelo (MOR). En lugar de minimizar el error de salida directo (que es no suave y carece de una caracterización de forma cerrada), la optimización apunta a la minimización de la cota superior derivada.

• Objetivo: Minimizar la suma ponderada de los errores de la norma $h_2$ entre los sistemas LQO originales y reducidos en cada capa.
• Principio de Diseño: Para minimizar esta cota, el método prescribe asignar dimensiones de estado reducidas mayores a las capas superficiales y dimensiones menores a las capas más profundas.

3. Algoritmo: MOR Basado en Gradientes

El artículo desarrolla un algoritmo de optimización basado en gradientes (Algoritmo 1) para resolver el problema de MOR:

• Derivación del Gradiente: Los autores derivan gradientes de forma cerrada para el error de la norma $h_2$ con respecto a los parámetros de los sistemas LQO reducidos (matrices $\hat{A}, \hat{B}, \hat{C}, \hat{M}$). Estos gradientes se adaptan para sistemas de tiempo discreto y de valor complejo con matrices de estado diagonales.
• Garantía de Estabilidad: El algoritmo emplea una búsqueda de línea con retroceso con una condición de Armijo. Crucialmente, hace cumplir explícitamente la estabilidad de Schur (todos los eigenvalores dentro del disco unitario) para la matriz de estado reducida $\hat{A}$ en cada iteración, evitando la generación de modelos inestables que pueden ocurrir en métodos MOR estándar.
• Inicialización: El método puede inicializarse con técnicas existentes de MOR de tiempo limitado (por ejemplo, Truncamiento Balanceado de Tiempo Limitado o reducción $H_2$ de Tiempo Limitado), pero los refina utilizando el descenso de gradiente propuesto para satisfacer las restricciones específicas de Deep SSM.

Contribuciones Clave

1. Cota Teórica: El artículo proporciona la primera derivación teórica de una cota superior sobre el error de salida de un Deep SSM en términos de los errores de la norma $h_2$ de sus bloques LQO constituyentes. Esto establece un vínculo directo entre la reducción por capa y el rendimiento global.
2. Análisis de Sensibilidad por Capa: El análisis demuestra que reducir los errores de aproximación en capas superficiales es más efectivo para minimizar la cota de error de salida global que reducir errores en capas más profundas.
3. Algoritmo de Optimización: Se propone un novedoso algoritmo MOR basado en gradientes que minimiza la cota superior derivada mientras garantiza la estabilidad del modelo reducido y preserva la estructura diagonal de los Deep SSMs.
4. Compresión sin Reentrenamiento: El método está diseñado para construir modelos comprimidos precisos a partir de pesos preentrenados sin necesidad de reentrenamiento (compresión de un solo disparo).

Resultados Experimentales

Los autores evaluaron el método en la tarea de análisis de sentimientos de IMDb del benchmark Long Range Arena (LRA) y en la tarea de modelado de lenguaje WikiText-103.

• Tarea IMDb:
  • Un Deep SSM preentrenado con 207,490 parámetros alcanzó una precisión de prueba del 86.66%.
  • Utilizando el método propuesto con una asignación específica de dimensiones por capa ($r_{list} = [96, 48, 36, 12]$), el modelo se comprimió a 83,650 parámetros (una reducción de ~60%) manteniendo una precisión de prueba del 86.69%.
  • Los resultados confirmaron que asignar dimensiones mayores a capas superficiales (por ejemplo, $[96, 48, 36, 12]$) produjo una precisión significativamente mejor que una asignación uniforme (por ejemplo, $[16] \times 4$) para la misma cantidad total de parámetros.
  • El método propuesto superó las líneas base existentes de compresión sin reentrenamiento (como las de [12]) en dimensiones de estado comparables.
• Tarea WikiText-103:
  • Se observaron tendencias similares donde minimizar la cota superior derivada se correlacionó con una mejora en la Perplejidad (PPL).
  • Los modelos con asignación de dimensiones no uniforme (mayores en capas superficiales) alcanzaron un PPL cercano al del modelo preentrenado original, mientras que la asignación uniforme resultó en un PPL más alto.

Significado y Afirmaciones

El artículo afirma que su principal significado radica en proporcionar una estrategia de compresión fundamentada teóricamente y sin reentrenamiento para Deep SSMs. Al cambiar el enfoque de la aproximación de bloques locales a una cota de error global, el método ofrece una manera sistemática de determinar dónde comprimir un modelo.

Los autores enfatizan que su enfoque permite un despliegue de bajo costo y alto rendimiento en escenarios donde el reentrenamiento es impráctico debido a restricciones computacionales. El método valida que la estrategia de reducción "primero las capas superficiales" no es solo heurística, sino que está respaldada por la cota de error de salida derivada, ofreciendo una alternativa práctica a las técnicas MOR existentes que a menudo fallan en preservar el rendimiento general del sistema cuando se aplican por capa sin consideración global.