En los marcos de aprendizaje automático y inteligencia artificial modernos, los transformadores son uno de los componentes más utilizados en diversos 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 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 la kernelización, la compresión de la memoria histórica, la limitación del rango de mezcla de tokens y los enfoques de bajo rango. Recientemente, las Redes Neuronales Recurrentes como los métodos que incluyen Mamba y RWKV han ganado una atención significativa 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 similar a un modelo de espacio de estado fue introducida 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 Visión 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 de trabajo VMamba emplea el modelo base de Mamba para construir modelos jerárquicos similares a ResNet y AlexNet. Sin embargo, ¿es el marco de trabajo de Mamba realmente esencial para las tareas de contexto de reconocimiento visual? La pregunta surge porque el rendimiento de la familia de modelos de Mamba para tareas de visión ha sido decepcionante hasta ahora en comparación con los modelos tradicionales basados en atención y convolucionales.

MambaOut intenta responder si Mamba es idealmente adecuado para tareas con características autoregresivas y de secuencias largas. El marco de trabajo de 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 también no son autoregresivas, muestran características de secuencias largas, lo que lleva al marco de trabajo de MambaOut a plantear la hipótesis de la potencialidad de Mamba para estas tareas. El marco de trabajo de MambaOut se construye apilando bloques de Mamba uno sobre otro mientras se elimina el modelo de espacio de estado, su mezclador de tokens central. Los resultados experimentales apoyan la hipótesis planteada por el marco de trabajo de MambaOut, ya que es capaz de superar a todos los modelos de visión de Mamba en el marco de trabajo 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 de trabajo de MambaOut no puede replicar el rendimiento ofrecido por el modelo de Mamba de última generación, lo que demuestra la potencialidad de la familia de modelos de Mamba para tareas de visión de secuencias largas.

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

MambaOut: ¿Se necesita realmente 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 para secuencias más largas. 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 de Red Neuronal Recurrente como Mamba y RWKV han ganado prominencia debido a su capacidad de entrenamiento paralelo y su rendimiento eficiente en secuencias más largas. Guiados por el rendimiento notable ofrecido por los modelos de Red Neuronal Recurrente, los investigadores están intentando introducir y utilizar la familia de modelos de Mamba en tareas de reconocimiento visual, ya que el mezclador de tokens de los modelos de 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 de trabajo basados en el modelo de espacio de estado para visión realizan de manera decepcionante en tareas de visión reales en comparación con los modelos basados en atención y convolucionales de última generación.

MambaOut es un intento de investigar la naturaleza de la familia de modelos de 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 inherente de Red Neuronal Recurrente. Sin embargo, la mayoría de las tareas de visión no presentan ambas características, y en base a algunos experimentos, MambaOut propone 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 tarea de 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 de Red Neuronal Recurrente del modelo de espacio de estado concluyen que el marco de trabajo de 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 de trabajo de MambaOut en sí, es una serie de modelos de Mamba basados en bloques de Red Neuronal Convolucional con Puerta sin el modelo de espacio de estado, y los resultados experimentales indican que el marco de trabajo de MambaOut es capaz de superar a los modelos de Mamba visuales en tareas de clasificación de imágenes, pero no puede 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 de trabajo de Mamba es un modelo de espacio de estado selectivo que define cuatro parámetros dependientes de la entrada. La propiedad recurrente del marco de trabajo distingue a los modelos de espacio de estado similares a Red Neuronal Recurrente 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 expande 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 de Red Neuronal Recurrente 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 de Mamba no pueden demostrar su fuerza en el manejo de secuencias cortas, un área donde el mecanismo de atención causal se desempeña bien con facilidad. Sin embargo, en escenarios que involucran secuencias largas, el enfoque de atención causal se debilita debido a la complejidad cuadrática. En este escenario, el marco de trabajo de Mamba muestra su eficiencia en la integración de la memoria con la entrada actual, y es capaz de manejar secuencias largas suavemente, lo que indica que la familia de modelos de 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 a los modelos de Mamba manejar eficientemente secuencias largas, introduce una cierta limitación, ya que solo puede acceder a información del token actual y de los tokens 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 completamente 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 completamente 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 de Red Neuronal Recurrente operan inherentemente en modo causal debido a sus propiedades recurrentes. Para resumir, el marco de trabajo de Mamba es adecuado para tareas que involucran el procesamiento de secuencias largas o tareas que requieren el 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 completamente visible permite un rango de mezcla sin restricciones, mientras que el modo causal limita el token actual para 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 puede degradar el rendimiento del modelo potencialmente. En general, el modo completamente visible es adecuado para tareas de comprensión, mientras que el modo causal es más adecuado para tareas autoregresivas. Además, este reclamo está 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 de trabajo de MambaOut experimentalmente. Como se demuestra en la siguiente imagen, el bloque de Mamba se basa en el bloque de Red Neuronal Convolucional con Puerta, y la meta-arquitectura de los bloques de Mamba y Gated CNN se puede tratar como una integración simplificada del mezclador de tokens del marco de trabajo de MetaFormer y un MLP.

El bloque de 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 de Gated CNN y el bloque de Mamba. Además, para mejorar la velocidad práctica, el marco de trabajo de MambaOut solo realiza convolución en profundidad en canales parciales, y como se demuestra en el siguiente algoritmo, la implementación del bloque de 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 consta de 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 ampliación 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 de Mamba, el modelo de MambaOut y otros modelos basados en atención y convolucionales en el conjunto de datos de ImageNet. Como se puede ver, el marco de trabajo de MambaOut sin el modelo de espacio de estado supera consistentemente a los modelos de visión de 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%, un 0,4% más alto que su competidor de Mamba más cercano. Este resultado apoya 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 de trabajo de MambaOut es capaz de superar el rendimiento de algunos modelos de visión de Mamba, todavía no alcanza el rendimiento de los modelos de visión de Mamba de última generación, incluyendo LocalVMamba y VMamba. La disparidad en el rendimiento de MambaOut en comparación con los modelos visuales de última generación enfatiza los beneficios de integrar la familia de modelos de Mamba en tareas de visión 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 última generación y los modelos de visión de Mamba.

Pensamientos finales

La familia de modelos de Mamba parece ser adecuada para tareas que involucran características autoregresivas y de secuencias largas. El marco de trabajo de 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 también no son autoregresivas, muestran características de secuencias largas, lo que lleva al marco de trabajo de MambaOut a plantear la hipótesis de la potencialidad de Mamba para estas tareas. El marco de trabajo de MambaOut se construye apilando bloques de Mamba uno sobre otro mientras se elimina el modelo de espacio de estado, su mezclador de tokens central. Los resultados experimentales apoyan la hipótesis planteada por el marco de trabajo de MambaOut, ya que es capaz de superar a todos los modelos de visión de Mamba en el marco de trabajo 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 de trabajo de MambaOut no puede replicar el rendimiento ofrecido por el modelo de Mamba de última generación, lo que demuestra la potencialidad de la familia de modelos de Mamba para tareas de visión de secuencias largas.