En los marcos modernos de aprendizaje automático e inteligencia artificial, los transformadores son uno de los componentes más utilizados en various dominios, incluyendo la serie GPT y BERT en Procesamiento de Lenguaje Natural, y Transformadores de Visión en tareas de visión por computadora. Aunque incluir transformadores en la arquitectura del modelo da un impulso significativo en el rendimiento del modelo, el módulo de atención en los Transformadores se escala con la longitud de la secuencia de manera cuadrática, lo que lleva a importantes desafíos computacionales. A lo largo de los años, diferentes modelos han explorado diferentes estrategias para abordar los desafíos computacionales, incluyendo métodos como kernelización, compresión de memoria de historia, limitación de rango de mezcla de tokens y enfoques de bajo rango. Recientemente, las Redes Neuronales Recurrentes como los métodos Mamba y RWKV han ganado una gran atención debido a sus resultados prometedores en grandes modelos de lenguaje. 

Mamba, una familia de modelos tiene una arquitectura con un mezclador de tokens de Red Neuronal Recurrente de un modelo de espacio de estado fue introducido recientemente para abordar la complejidad cuadrática de los mecanismos de atención y se aplicó a tareas de visión posteriormente.  Los investigadores ya han explorado formas de incorporar Mamba y SSM o Modelo de Espacio de Estado en tareas de reconocimiento visual, y Vision Mamba que incorpora Mamba para desarrollar modelos de visión isotrópicos similares a los Transformadores de Visión es un gran ejemplo de lo mismo. Por otro lado, LocalMamba incorpora sesgos inductivos locales para mejorar los modelos de visión de Mamba, y el marco VMamba emplea el modelo base Mamba para construir modelos jerárquicos similares a ResNet y AlexNet. Sin embargo, ¿es el marco Mamba realmente esencial para las tareas de contexto de reconocimiento visual? La pregunta surge porque el rendimiento de la familia de modelos Mamba para tareas de visión ha sido poco impresionante hasta ahora en comparación con los modelos basados en atención y convolucionales tradicionales. 

MambaOut intenta responder si Mamba es idealmente adecuado para tareas con características autoregresivas y de secuencias largas. El marco MambaOut plantea la hipótesis de que Mamba no es necesario para tareas de visión, ya que la clasificación de imágenes no se ajusta a características autoregresivas ni de secuencias largas. Aunque las tareas de segmentación y detección no son autoregresivas, muestran características de secuencias largas, lo que lleva al marco MambaOut a plantear la posibilidad de Mamba para estas tareas. El marco MambaOut se construye apilando bloques Mamba uno sobre otro mientras se elimina el modelo de espacio de estado, su mezclador de tokens principal. Los resultados experimentales respaldan la hipótesis planteada por el marco MambaOut, ya que es capaz de superar a todos los modelos de visión Mamba en el marco de clasificación de imágenes ImageNet, lo que indica que Mamba no es necesario para tareas de visión. Por otro lado, para tareas de detección y segmentación, el marco MambaOut no puede replicar el rendimiento ofrecido por el modelo Mamba de estado del arte, lo que demuestra el potencial de la familia de modelos Mamba para tareas visuales de secuencias largas. 

Este artículo tiene como objetivo cubrir el marco MambaOut en profundidad, y exploramos el mecanismo, la metodología, la arquitectura del marco junto con su comparación con marcos de estado del arte. Así que comencemos. 

MambaOut: ¿Es realmente necesario Mamba para la visión?

Con el progreso de las aplicaciones y capacidades de aprendizaje automático, los Transformadores han surgido como la columna vertebral principal para una serie de tareas, impulsando modelos prominentes como Transformadores de Visión, la serie de modelos GPT, BERT, y algunos más. Sin embargo, el mezclador de tokens del transformador incurre en una complejidad cuadrática con respecto a la longitud de la secuencia, y plantea importantes desafíos computacionales. Para abordar este problema, se han introducido numerosos mezcladores de tokens con complejidad lineal con respecto a la longitud del token, como Linformer, Longformer, Performer, Convolución Dinámica y Big Bird. Sin embargo, en tiempos recientes, los modelos similares a Redes Neuronales Recurrentes están ganando prominencia debido a su capacidad de entrenamiento paralelizable y su rendimiento eficiente en secuencias más largas. Guiados por el rendimiento notable de los modelos similares a RNN, los investigadores están intentando introducir y utilizar la familia de modelos Mamba en tareas de reconocimiento visual, ya que el mezclador de tokens de los modelos Mamba es el modelo de espacio de estado estructurado bajo el espíritu de las Redes Neuronales Recurrentes. Sin embargo, los resultados experimentales indican que los marcos basados en el modelo de espacio de estado para la visión funcionan de manera poco impresionante en tareas de visión reales en comparación con los modelos basados en atención y convolucionales de estado del arte. 

MambaOut es un intento de investigar la naturaleza de la familia de modelos Mamba, y resume que Mamba es adecuado para tareas que son autoregresivas o de secuencias largas, ya que el modelo de espacio de estado tiene un mecanismo RNN inherente. Sin embargo, la mayoría de las tareas de visión no presentan ambas características, y sobre la base de algunos experimentos, MambaOut plantea las siguientes dos hipótesis. Primero, el modelo de espacio de estado no es necesario para la clasificación de imágenes, ya que la clasificación de imágenes no se ajusta a características autoregresivas ni de secuencias largas. Segundo, los modelos de espacio de estado pueden ser hipotéticamente beneficiosos para la segmentación de instancias y la segmentación semántica, así como para la detección de objetos, ya que siguen las características de secuencias largas, aunque no son autoregresivas. Los resultados experimentales realizados para analizar el mecanismo similar a Red Neuronal Recurrente del modelo de espacio de estado concluyen que el marco Mamba es adecuado para tareas con características autoregresivas o de secuencias largas, y es innecesario para tareas de clasificación de imágenes. En cuanto al marco MambaOut en sí, se construye apilando bloques Mamba uno sobre otro mientras se elimina el modelo de espacio de estado, y los resultados experimentales indican que el marco MambaOut es capaz de superar a los modelos Mamba en tareas de clasificación de imágenes, pero no logra replicar el rendimiento en tareas de detección y segmentación de imágenes. 

¿Para qué tareas es adecuado Mamba?

El mezclador de tokens del marco Mamba es un modelo de espacio de estado selectivo que define cuatro parámetros dependientes de la entrada. La propiedad recurrente del marco distingue a los modelos de espacio de estado similares a RNN de la atención causal. El estado oculto puede verse como una memoria de tamaño fijo que almacena información histórica. El tamaño fijo significa que la memoria es pérdida, pero también garantiza que la complejidad computacional de integrar la memoria con la entrada actual permanece constante. Por el contrario, las capas de atención causal almacenan todas las claves y valores de los tokens anteriores, y se expanden agregando la clave y el valor del token actual con cada nueva entrada, y esta memoria es sin pérdida, teóricamente. Sin embargo, el tamaño de la memoria crece a medida que se ingresan más tokens, lo que aumenta la complejidad de integrar la memoria con la entrada actual. La diferencia entre los mecanismos de memoria entre la atención causal y los modelos similares a RNN se ilustra en la siguiente figura. 

Dado que la memoria del modelo de espacio de estado es inherentemente pérdida, no alcanza la memoria sin pérdida de la atención causal, y como resultado, los modelos Mamba no pueden demostrar su fuerza en el manejo de secuencias cortas, un área donde el mecanismo de atención causal funciona bien con facilidad. Sin embargo, en escenarios que involucran secuencias largas, el enfoque de atención causal falla debido a la complejidad cuadrática. En este escenario, el marco Mamba muestra su eficiencia en combinar la memoria con la entrada actual, y es capaz de manejar secuencias largas suavemente, lo que indica que la familia de modelos Mamba es adecuada para el procesamiento de secuencias largas. 

También es digno de nota que, por un lado, la naturaleza recurrente del modelo de espacio de estado permite que los modelos Mamba manejen secuencias largas de manera eficiente, pero introduce una cierta limitación, ya que solo puede acceder a información de los tiempos actuales y anteriores, y este tipo de mezcla de tokens se denomina modo causal, y se ilustra en la siguiente figura. Debido a su naturaleza causal, este método es adecuado para tareas de generación autoregresiva. 

El modo fully-visible es adecuado para tareas de comprensión donde el modelo puede acceder a todas las entradas al mismo tiempo. Además, la atención está en modo fully-visible por defecto, y se puede convertir en modo causal fácilmente aplicando máscaras causales a los mapas de atención, y los modelos similares a RNN operan inherentemente en modo causal debido a sus propiedades recurrentes. Para resumir, el marco Mamba es adecuado para tareas que involucran el procesamiento de secuencias largas o tareas que requieren un modo de mezcla de tokens causal.

Tareas de reconocimiento visual, código de mezcla de tokens causal y secuencias muy largas

Como se discutió anteriormente, el modo de mezcla de tokens fully-visible permite un rango de mezcla sin restricciones, mientras que el modo causal limita el token actual a acceder solo a la información de los tokens anteriores. Además, el reconocimiento visual se categoriza como una tarea de comprensión donde el modelo puede ver la imagen completa al mismo tiempo, y esto elimina la necesidad de restricciones en la mezcla de tokens, y imponer restricciones adicionales en la mezcla de tokens puede degradar el rendimiento del modelo potencialmente. Generalmente, el modo fully-visible es adecuado para tareas de comprensión, mientras que el modo causal es más adecuado para tareas autoregresivas. Además, este reclamo se ve respaldado por el hecho de que los modelos BeRT y ViT se utilizan más para tareas de comprensión que para los modelos GPT.

Verificación experimental y resultados

El siguiente paso es verificar las hipótesis propuestas por el marco MambaOut experimentalmente. Como se demuestra en la siguiente imagen, el bloque Mamba se basa en el bloque de Red Neuronal Convolucional con Puerta, y la meta-arquitectura de los bloques Mamba y Gated CNN se puede tratar como una integración simplificada del mezclador de tokens del marco MetaFormer y un MLP. 

El bloque Mamba extiende el bloque de Red Neuronal Convolucional con Puerta con un modelo de espacio de estado adicional, y la presencia de un SSm es lo que distingue el bloque Gated CNN y el bloque Mamba. Además, para mejorar la velocidad práctica, el marco MambaOut solo realiza convolución en canales parciales, y como se demuestra en el siguiente algoritmo, la implementación del bloque Gated CNN es simple, pero efectiva y elegante. 

Tarea de clasificación de imágenes

ImageNet sirve como el benchmark para tareas de clasificación de imágenes, ya que consiste en más de mil clases comunes, más de 1,3 millones de imágenes de entrenamiento y más de 50.000 imágenes de validación. La mejora de datos utilizada para el experimento consiste en recorte de tamaño aleatorio, Mixup, ajuste de color, Borrado Aleatorio, CutMix y Rand Augment. La siguiente tabla resume el rendimiento de la familia de modelos Mamba, el modelo MambaOut y otros modelos basados en atención y convolucionales en el conjunto de datos ImageNet. Como se puede ver, el marco MambaOut sin el modelo de espacio de estado supera consistentemente a los modelos de visión Mamba con SSM en todos los tamaños de modelo. 

Por ejemplo, el modelo MambaOut-Small devuelve una puntuación de precisión superior al 84%, 0,4% más alto que su competidor Mamba más cercano. Este resultado respalda firmemente la primera hipótesis de que afirma que introducir un modelo de espacio de estado para tareas de clasificación de imágenes no es necesario. 

Tareas de detección de objetos y segmentación de instancias

COCO sirve como el benchmark para tareas de detección de objetos y segmentación de instancias. Aunque el marco MambaOut es capaz de superar el rendimiento de algunos modelos de visión Mamba, todavía no logra replicar el rendimiento de los modelos de visión Mamba de estado del arte, incluyendo LocalVMamba y VMamba. La disparidad en el rendimiento de MambaOut contra los modelos visuales de estado del arte enfatiza en los beneficios de integrar la familia de modelos Mamba en tareas visuales de secuencias largas. Sin embargo, es digno de nota que todavía existe una brecha significativa en el rendimiento entre los modelos híbridos de convolución-atención de estado del arte y los modelos visuales Mamba. 

Pensamientos finales

La familia de modelos Mamba parece ser adecuada para tareas que involucran características autoregresivas y de secuencias largas. El marco MambaOut plantea la hipótesis de que Mamba no es necesario para tareas de visión, ya que la clasificación de imágenes no se ajusta a características autoregresivas ni de secuencias largas. Aunque las tareas de segmentación y detección no son autoregresivas, muestran características de secuencias largas, lo que lleva al marco MambaOut a plantear la posibilidad de Mamba para estas tareas. El marco MambaOut se construye apilando bloques Mamba uno sobre otro mientras se elimina el modelo de espacio de estado, su mezclador de tokens principal. Los resultados experimentales respaldan la hipótesis planteada por el marco MambaOut, ya que es capaz de superar a todos los modelos de visión Mamba en el marco de clasificación de imágenes ImageNet, lo que indica que Mamba no es necesario para tareas de visión. Por otro lado, para tareas de detección y segmentación, el marco MambaOut no puede replicar el rendimiento ofrecido por el modelo Mamba de estado del arte, lo que demuestra el potencial de la familia de modelos Mamba para tareas visuales de secuencias largas.