Zubnet AIAprenderWiki › Modelo de espacio de estados
Modelos

Modelo de espacio de estados

También conocido como: SSM, Mamba
Una alternativa a los Transformers que procesa secuencias manteniendo un "estado" comprimido en lugar de usar atención sobre todos los tokens. Mamba es la arquitectura SSM más conocida. Los SSMs escalan linealmente con la longitud de la secuencia (vs. cuadraticamente para la atención), haciendolos potencialmente mucho más eficientes para contextos muy largos.

Por qué importa

Los SSMs son el principal retador al dominio de los Transformers. Son más rapidos para secuencias largas y usan menos memoria, pero la investigación aun esta madurando. Las arquitecturas hibridas (mezclando capas SSM con atención) podrían terminar siendo lo mejor de ambos mundos.

En profundidad

Los modelos de espacio de estados toman prestado su marco matemático de la teoría de control, donde los SSMs se han usado durante decadas para modelar sistemas dinamicos. La idea central es una recurrencia lineal: el modelo mantiene un estado oculto que se actualiza en cada paso temporal mediante una transformación lineal aprendida, y luego se mezcla con la entrada actual. En tiempo continuo, esto es una ecuación diferencial (dx/dt = Ax + Bu, y = Cx + Du). Discretizarla te da una recurrencia que puede procesar secuencias token por token, actualizando un estado de tamaño fijo en cada paso. La parte elegante es que durante el entrenamiento, esta recurrencia puede desplegarse en una convolución, haciendola paralelizable en GPUs igual que la atención. Durante la inferencia, vuelves a la forma de recurrencia y procesas tokens uno a la vez con memoria constante — sin un KV cache que crece.

El avance de Mamba

Mamba (Albert Gu y Tri Dao, 2023) fue el avance que hizo a los SSMs competitivos con los Transformers en lenguaje. Los SSMs anteriores como S4 y H3 usaban matrices de transición de estado fijas, lo que limitaba su capacidad para realizar razonamiento basado en contenido — el modelo no podia cambiar como procesaba un token basandose en lo que ese token era. Mamba introdujo espacios de estados selectivos, donde las matrices A, B y C son funciones de la entrada. Esto permite que el modelo decida, en cada token, cuanto recordar y cuanto olvidar. Piensa en ello como un mecanismo de compuerta aprendido y diferenciable, pero operando a través del lente de la recurrencia lineal en lugar de la atención. Mamba-2 luego reformulo esto como dualidad de espacio de estados estructurado (SSD), revelando que los SSMs selectivos y la atención lineal están matemáticamente relacionados, y permitiendo implementaciones GPU aun más rapidas mediante algoritmos basados en multiplicación de matrices.

Las ventajas practicas son reales y medibles. Durante la inferencia, un Transformer debe almacenar pares clave-valor para cada token en el contexto — ese KV cache crece linealmente con la longitud de la secuencia y es el principal cuello de botella para servir contextos largos. Un SSM mantiene un estado de tamaño fijo sin importar cuántos tokens haya visto. Para un modelo con una ventana de contexto de 128K, esta diferencia es enorme: el SSM usa la misma memoria generando el token 128,001 que al generar el token 1. El throughput de entrenamiento también se beneficia en secuencias largas porque el modo de scan paralelo o convolución escala linealmente con la longitud de la secuencia, versus el escalado cuadratico de la atención completa. Estas ganancias de eficiencia son la razón por la que los SSMs son particularmente atractivos para aplicaciones que necesitan contexto de largo alcance: análisis de documentos, generación de código en repositorios grandes y streaming en tiempo real donde los tokens llegan continuamente.

El problema de la recuperación

Dicho esto, los SSMs tienen limitaciones reales que el entusiasmo tiende a pasar por alto. Los SSMs puros pueden tener dificultades con tareas que requieren recuperación precisa de información anterior en el contexto — el problema de la "aguja en el pajar". Un Transformer puede, en principio, atender directamente a cualquier token pasado a través de sus pesos de atención. Un SSM debe haber comprimido la información relevante en su estado de tamaño fijo, y si no priorizo las cosas correctas cuando proceso originalmente ese token, la información se perdio. Por eso las arquitecturas hibridas — intercalando capas SSM con unas pocas capas de atención — están ganando tracción. Jamba (de AI21) y varios hibridos de investigación han demostrado que se puede obtener la mayor parte de la eficiencia de los SSMs con la precisión de recuperación de la atención usando atención de forma moderada en puntos estrategicos de la red.

La vanguardia

Mamba-3, la generación más reciente, lleva la arquitectura más alla con una formulación multi-entrada multi-salida (MIMO) y estados de valor complejo vía codificaciones de posición rotatorias. La recurrencia usa una regla de integración trapezoidal para mejor estabilidad numerica, y la arquitectura elimina la capa de convolución causal que las versiones anteriores usaban como mecanismo de mezcla de corto alcance. Estos no son ajustes incrementales — cambian el perfil computacional lo suficiente como para que se necesiten kernels personalizados de Triton para obtener rendimiento completo, y el paquete estandar de mamba-ssm en PyPI aun no los incluye. Si estas construyendo sobre SSMs hoy, espera trabajar más cerca del metal de lo que lo harias con un stack de Transformer maduro. Las herramientas están avanzando, pero aun son días tempranos para el despliegue de SSMs en producción.

Conceptos relacionados

← Todos los términos
← Modelo de difusión Modelo de fundamento →
ESC