Mamba: el modelo de espacio de estados que desafía a los Transformers

Mamba es un modelo de espacio de estados (State Space Model) que desafía a Transformer

• La IA domina hoy el mundo, y en el centro de todo están los Transformers
• Mamba pertenece a una clase alternativa de modelos llamada modelos de espacio de estados (SSM)
• Mamba ofrece un rendimiento y una escalabilidad similares a los de Transformer, pero puede ejecutarse con longitudes de secuencia largas
• Lo más destacable es que Mamba elimina el “cuello de botella cuadrático” del “mecanismo de atención”, lo que habilita contextos largos
• Mamba se ejecuta hasta 5 veces más rápido que Transformer

Problemas de Transformer: la atención por sí sola podría no ser suficiente

• En Transformer, todos los tokens pueden referenciar tokens anteriores, por lo que el modelo se vuelve más lento a medida que crece el contexto
• Almacenar este KV cache también requiere una complejidad espacial de O(n)
• Existen técnicas para aliviar los cuellos de botella de los Transformers actuales, pero para resolver el problema de raíz se necesita un enfoque distinto

Backbone de modelos fundacionales

• Dos componentes importantes de un buen backbone de arquitectura de ML
  • Comunicación (Communication) entre tokens
  • Cómputo (Computation) dentro de cada token
• Los bloques Transformer están compuestos por atención (Attention) y MLPs
• Mamba usa SSM, inspirados en la teoría de control, para la comunicación, y mantiene proyecciones estilo MLP para el cómputo

Motivación de Mamba: volver a Temple Run

• El estado (state) se refiere a las variables necesarias para determinar el comportamiento futuro de un sistema
• El estado comprime todo lo que se necesita saber del pasado y se transforma en un proceso de decisión de Markov

Discretización: vivir en un mundo cuantizado

• Al proceso de convertir ecuaciones diferenciales en tiempo continuo en ecuaciones en diferencias de tiempo discreto se le llama discretización (discretisation)
• Mamba usa discretización Zero-Order Hold (ZOH)

Entender las matrices de SSM

• Las matrices A, B, C y D cumplen las funciones de transición de estado, mapeo de nuevas entradas al estado, mapeo del estado a la salida del SSM y paso de nuevas entradas a la salida

Eficiencia vs efectividad: Attention is Focus, Selectivity is Prioritisation (la atención es enfoque y la selectividad es priorización)

• Los Transformers son muy efectivos, pero no necesariamente eficientes
• La arquitectura Mamba ofrece una solución que empuja la frontera de Pareto entre eficiencia y efectividad

Mecanismo de selección

• La selectividad (Selectivity) permite convertir cada token en estado según sus necesidades
• Mamba hace que las matrices A, B y C sean funciones de x, volviéndolas dependientes del contexto en lugar de estáticas

Problemas de la selectividad

• Aplicar el mecanismo de selección puede hacer que el cálculo sea más lento que en un SSM no selectivo
• Gracias a la optimización de hardware, Mamba puede ejecutarse más rápido que Transformers de tamaño similar

Aprendizaje automático y economía política: ¿qué tan grande debe ser el estado?

• El trade-off entre eficiencia y efectividad en modelos de secuencia se caracteriza por qué tan bien se comprime el estado
• La representación del estado es importante, y la clave está en comprimirlo de forma selectiva y dinámica

Flujo de información en Transformer vs Mamba

• Transformer aprende a través de los datos de entrenamiento y de contexto
• En Mamba, los datos de entrenamiento y de contexto se comprimen/filtran para poder acceder a ellos

Reemplazo de estado como nuevo paradigma de prompting

• Con modelos como Mamba, se pueden compartir bibliotecas de estados generadas a partir de datos especializados
• Los estados permiten aplicar aprendizaje de contexto infinito en tiempo de inferencia sin backprop

Mamba y la interpretabilidad mecanística

• La interpretabilidad de Mamba se enfoca en entender cómo se mueve la información entre tokens

Lo siguiente para Mamba y los SSM

• Modelos como Mamba podrían destacar en escenarios que requieren contextos muy largos y memoria de largo plazo

Agentes y seguridad en IA

• Los modelos de lenguaje son inherentemente seguros, pero la posibilidad de razonamiento secuencial de largo plazo vuelve a poner sobre la mesa la importancia de la seguridad en IA basada en agentes

La mejor colaboración entre Transformer y Mamba

• Vale la pena combinar el contexto largo de Mamba con la alta resolución de Transformer para secuencias cortas

Opinión de GN⁺

• Mamba resuelve los cuellos de botella de Transformer y presenta una alternativa eficaz para procesar secuencias largas
• Esta tecnología puede ser especialmente útil en áreas donde importan las secuencias largas de datos, como medicina, genética y procesamiento de lenguaje natural
• Se necesita más investigación para confirmar si el mecanismo de selección de Mamba realmente es eficaz
• La selectividad de Mamba podría ayudar a encontrar un equilibrio entre la alta precisión y la eficiencia que ofrecen los Transformers