El desarrollo de modelos de lenguaje grande (LLM, por sus siglas en inglés) creados a partir de modelos transformadores exclusivos para decodificadores ha desempeñado un papel crucial en la transformación del dominio del procesamiento del lenguaje natural (NLP, por sus siglas en inglés), así como en el avance de diversas aplicaciones de aprendizaje profundo, incluidas aprendizaje reforzado, análisis de series temporales, procesamiento de imágenes y mucho más. Sin embargo, a pesar de su escalabilidad y su sólido rendimiento, los LLM creados a partir de modelos de transformadores exclusivamente decodificadores todavía enfrentan importantes deficiencias. Aunque expresivo, el mecanismo de atención en los LLM derivados de transformadores requiere altos recursos computacionales tanto durante la inferencia como durante el entrenamiento, lo que requiere una memoria sustancial para la longitud de la secuencia y los FLOP cuadráticos. Este alto requisito computacional limita la longitud del contexto de los modelos de transformadores, lo que hace que las tareas de generación autorregresiva sean proporcionalmente costosas con la escala, y dificulta el aprendizaje de flujos de datos continuos y la capacidad de un procesamiento de secuencias verdaderamente ilimitado.

En tiempos recientes, Modelos de espacio de estado (SSM) han demostrado capacidades y rendimiento notables, compitiendo con modelos de arquitectura de transformadores en puntos de referencia de modelado a gran escala y al mismo tiempo logrando complejidad de la memoria en función de la longitud de la secuencia y el tiempo lineal. Además, Mamba, un modelo de espacio de estados lanzado recientemente, ha demostrado un rendimiento sobresaliente en una variedad de tareas de modelado de lenguaje y procesamiento de secuencias largas. Al mismo tiempo, los modelos Mixture of Expert (MoE) también han mostrado un rendimiento impresionante al tiempo que reducen significativamente la latencia y los costos computacionales de la inferencia, aunque a expensas de una mayor huella de memoria. Sobre la base de los modelos Mamba y MoE, este artículo analizará BlackMamba, una arquitectura novedosa que combina el modelo espacial estatal Mamba con modelos MoE para aprovechar los beneficios que ofrecen ambos marcos. Los experimentos en BlackMamba han demostrado su capacidad para superar el marco Mamba existente y las líneas base del transformador tanto en el entrenamiento de FLOP como en la inferencia. El rendimiento excepcional del marco BlackMamba muestra que puede combinar eficazmente las capacidades de los marcos Mamba y MoE, ofreciendo una inferencia rápida y rentable desde MoE con generación de complejidad lineal desde Mamba.

Este artículo pretende profundizar en el framework BlackMamba. Exploramos su mecanismo, metodología y arquitectura, comparándolo con frameworks de generación de imágenes y vídeo de vanguardia. Comencemos.

BlackMamba: Introducción al MoE para modelos de espacio estatal

La progresión de los modelos de lenguajes grandes (LLM), en particular aquellos basados ​​en arquitecturas transformadoras exclusivamente decodificadoras, ha influido notablemente en la Procesamiento del lenguaje natural (PNL) campo y se expandió a varias aplicaciones de aprendizaje profundo, incluido el aprendizaje por refuerzo, el análisis de series de tiempo, el procesamiento de imágenes y más. No obstante, a pesar de su escalabilidad y rendimiento sólido, estos LLM basados ​​​​en transformadores solo decodificadores enfrentan desafíos notables. El mecanismo de atención, una característica clave de la tecnología basada en transformadores. LLMs, exige amplios recursos computacionales tanto para la inferencia como para el entrenamiento. Esto implica una necesidad de memoria que crece con la longitud de la secuencia y operaciones computacionales (FLOP) que aumentan cuadráticamente. Estas necesidades computacionales intensivas restringen la longitud del contexto de los modelos, elevan los costos de las tareas de generación autorregresiva a medida que el modelo escala y dificultan su capacidad para aprender de flujos de datos continuos o procesar secuencias de longitud ilimitada de manera eficiente. 

En los últimos años se han realizado importantes esfuerzos para superar estas limitaciones, y la atención se ha centrado en diseñar alternativas arquitectónicas a los modelos canónicos de transformadores de atención densa, siendo los modelos SSM y MoE las arquitecturas candidatas más prometedoras. La principal ventaja de preferir los modelos de espacio de estados sobre los modelos de arquitectura de transformadores es la complejidad computacional lineal con respecto a la longitud de la secuencia de entrada que ofrecen los SSM, en contraposición a la complejidad cuadrática de los transformadores. En teoría, la complejidad computacional lineal con respecto a la longitud de la secuencia de entrada permite a los modelos de espacio de estados procesar secuencias más grandes que los modelos de arquitectura de transformadores para un presupuesto determinado de FLOPS u operaciones de punto flotante por segundo, y mantener la generación autorregresiva constante en cómputo sin una caché KV. Los modelos de espacio de estados desarrollados recientemente, como Mamba, RetNet y algunos otros, han demostrado una inferencia y un entrenamiento eficientes de secuencias largas, junto con un rendimiento competitivo en tareas de modelado de lenguajes frente a transformadores con propiedades de escalado similares. Por otro lado, las arquitecturas de modelos de mezcla de expertos están ganando popularidad como alternativa a los transformadores densos, ya que facilitan una reducción significativa de los errores de inferencia y entrenamiento, esenciales para lograr una calidad comparable a la de un modelo denso. Los modelos MoE (mezcla de expertos) funcionan activando solo una selección dispersa de los parámetros totales durante una sola pasada hacia adelante. Utilizan una función de enrutamiento para determinar qué "expertos" entran en acción según el contexto dado. Este enfoque crea una separación entre el coste computacional de la inferencia y el número total de parámetros, lo que permite un mejor rendimiento con un presupuesto de inferencia fijo, aunque con un mayor número de parámetros y un mayor requerimiento de memoria.

Este avance en la arquitectura ofrece beneficios notables sobre los transformadores tradicionales y representa una dirección interesante para un mayor desarrollo. Postulamos que la integración de estas mejoras en un modelo combinado Mamba-MoE podría acelerar significativamente las capacidades y la eficiencia del modelado del lenguaje más allá de la de los modelos transformadores estándar. Las ventajas previstas de una arquitectura Mamba-MoE en comparación con un modelo de transformador denso tradicional incluyen:

Tipo de serpiente venenosa: Logra una complejidad computacional lineal en relación con la longitud de la secuencia de entrada tanto para las fases de entrenamiento como de inferencia. Permite que se produzca una generación autorregresiva en un período de tiempo constante y con un uso de memoria constante.

Ministerio de Educación: Ofrece una velocidad de inferencia y una eficiencia computacional de entrenamiento comparables a un modelo de referencia denso y más pequeño, al tiempo que mantiene un nivel de calidad del modelo que rivaliza con el de un modelo con un número equivalente de parámetros como la versión más densa.

Dicho esto, es esencial afirmar que los modelos de arquitectura de transformadores siguen siendo de última generación y han demostrado un rendimiento consistente y notable en tareas de modelado de lenguaje y tareas de procesamiento de secuencias. En esencia, la arquitectura del transformador emplea autoatención que realiza una comparación cuadrática de todas las similitudes del producto escalar entre las incrustaciones de diferentes tokens en una secuencia y realiza un mapa lineal a un vector de salida. El modelo de transformador consta de bloques de autoatención apilados entre MLP o bloques de perceptrón multicapa que, además, constan de un MLP de dos capas con una función de activación determinada. 

BlackMamba: Arquitectura y Metodología

Modelos de espacio de estado

Los modelos de espacio de estados pertenecen al grupo de modelos de secuencia con complejidad lineal con respecto a la longitud de la secuencia de entrada. La arquitectura de los modelos de espacio de estados se alinea más con las redes neuronales recurrentes y las redes neuronales convolucionales que con la arquitectura basada en la atención, y está inspirada en un sistema dinámico continuo que mapea una función unidimensional a través de un espacio latente implícito. Un sistema dinámico lineal hace que los cálculos paralelos sean eficientes mediante un escaneo asociativo o convolucional. En escenarios prácticos, la naturaleza recurrente de los modelos de espacio de estados ha sido la razón por la que aún no se han adoptado en hardware de IA altamente paralelo, como las GPU. Sin embargo, la aparición de MUS como RWKV y Mamba han utilizado núcleos de escaneo paralelo para asignar operaciones recurrentes de manera eficiente a las GPU, facilitando así el entrenamiento de arquitecturas novedosas con una eficiencia comparable a las logradas por los modelos de transformadores. 

La complejidad cuadrática inherente en relación con la longitud de la secuencia dentro de los transformadores es una limitación bien conocida que impide el razonamiento y la comprensión en contextos muy largos. Innovaciones recientes han introducido la idea de ampliar la duración del contexto, lo que permite entrenar a los transformadores en una escala factible antes de aplicarlos a contextos mucho más largos durante la inferencia. A pesar de estos avances, el proceso de inferencia todavía exige una cantidad considerable de recursos computacionales y memoria, especialmente para mantener la caché de valores clave (KV), lo que lo convierte en una tarea que consume muchos recursos. Los esfuerzos de investigación recientes se han centrado en mejorar las capacidades expresivas de los modelos de espacio de estados mediante la incorporación de mecanismos de activación dependientes de la entrada, similares a las matrices de consulta, clave y valor (QKV) que se encuentran en los mecanismos de atención. 

Estos esfuerzos tienen como objetivo preservar la progresión inherentemente lineal de la recursividad del espacio de estados, lo que permite una ejecución eficiente mediante convolución o un proceso de escaneo selectivo. Este enfoque reduce significativamente la disparidad de rendimiento con los transformadores en aplicaciones prácticas. Entre estos avances, Mamba se destaca como un modelo de espacio de estados que refleja los objetivos de investigaciones anteriores, mostrando impresionantes niveles de rendimiento comparables a los transformadores en escalas de hasta 2.8 millones de parámetros. Lo logra aplicando activación dependiente de la entrada a las entradas de la recursividad del modelo de espacio de estados (SSM), al mismo tiempo que garantiza un cálculo eficiente mediante el uso de núcleos de escaneo selectivo personalizados.

Mezcla de modelos expertos

Los modelos de mezcla de expertos (MoE) logran una separación entre el costo de inferencia y el recuento total de parámetros activando selectivamente los parámetros durante el paso hacia adelante. En lugar de utilizar todos los parámetros, estos modelos dirigen tokens a expertos específicos en perceptrones multicapa (MLP). Idealmente, cada experto está diseñado para procesar un tipo particular de entrada, con un mecanismo de enrutamiento, esencialmente una red neuronal compacta, que determina el experto más adecuado para cada token. Este enfoque tiene como objetivo preservar el poder expresivo integral de un modelo con un número equivalente de parámetros en una configuración más densa, pero con demandas computacionales considerablemente reducidas. Por lo general, el enrutador es un mapeo de las capas lineales desde tokens hasta índices de expertos, siendo cada experto simplemente un perceptrón multicapa transformador estándar. Sin embargo, los desarrolladores aún tienen que descubrir el método de entrenamiento óptimo para el enrutador, ya que el problema de asignación de expertos no es diferenciable, y los modelos de Mezcla de Expertos a menudo tienen dificultades con el equilibrio de carga y la estabilidad del entrenamiento entre diferentes expertos para lograr la eficiencia del hardware. 

Arquitectura

En esencia, BlackMamba emplea un modelo de transformador estándar que consta de bloques MLP entrelazados y bloques de atención agregados en secuencia a lo largo de un flujo residual. Ahora, la mayoría de los modelos Mixture of Expert simplemente reemplazan los bloques de perceptrón multicapa con una capa experta enrutada. Por otro lado, el marco BlackMamba no solo reemplaza el bloque de perceptrón multicapa en el transformador con una capa experta enrutada, sino que también reemplaza la capa de atención con una capa del Modelo de espacio de estado de Mamba. La arquitectura del marco BlackMamba se demuestra en la siguiente figura. 

Entrenamiento y conjunto de datos

El modelo BlackMamba se entrena con más de 300 mil millones de tokens en un conjunto de datos personalizado y utiliza la función de activación SwiGLU para los perceptrones multicapa expertos. El marco se entrena con 8 expertos, un número que los desarrolladores consideraron el equilibrio adecuado entre la huella de memoria y el costo de inferencia del modelo. El conjunto de datos personalizado utilizado para entrenar el marco BlackMamba consta de una combinación de conjuntos de datos de código abierto ya existentes, incluidos Starcoder, SlimPajama, Pile y más. La siguiente tabla demuestra los pesos de cada uno de los conjuntos de datos utilizados para entrenar el marco BlackMamba. En total, hay 1.8 billones de tokens en el conjunto de datos. 

BlackMamba: Resultados

Para garantizar una comparación justa entre Mamba y BlackMamba, los desarrolladores han entrenado ambos modelos con los mismos parámetros de entrenamiento con los mismos datos de entrenamiento. El marco BlackMamba es capaz de superar a los modelos Mamba y Transformer para un tamaño de modelo de paso directo idéntico en el momento de la inferencia, así como para entrenar operaciones de punto flotante por segundo. La siguiente figura demuestra el tiempo necesario para generar una secuencia de una longitud determinada de forma autorregresiva a partir de un mensaje inicial de un token en función de la longitud de la secuencia. 

Además, los beneficios de latencia de los modelos Mixture of Expert y Mamba se combinan en el marco BlackMamba, lo que da como resultado tiempos de inferencia significativamente más rápidos en comparación con los modelos transformadores, modelos Mamba puros y modelos MoE. Además, la ventaja de inferencia del marco BlackMamba es directamente proporcional a la longitud de las secuencias, lo que hace que BlackMamba sea extremadamente eficaz en la generación de secuencias largas. A continuación, la siguiente figura ilustra la cantidad de tokens asignados a los modelos BlackMamba con 340 millones y 640 millones de parámetros respectivamente. Como puede verse, la mayoría de las capas demuestran un alto nivel de equilibrio experto como resultado del algoritmo Sinkhorn mejorado implementado por los modelos BlackMamba. 

La siguiente tabla cubre los puntajes de evaluación del marco BlackMamba en comparación con una variedad de modelos de lenguaje previamente entrenados de código abierto. Como se puede observar, el marco BlackMamba puede competir y superar a la mayoría de los marcos en todas las líneas de base. Además, vale la pena señalar que los modelos que superan a BlackMamba tienen un número considerablemente mayor de parámetros y la brecha en el rendimiento es mínima, lo que indica la capacidad del marco BlackMamba con menos parámetros. 

Conclusión

En este artículo, hemos hablado de BlackMamba, una arquitectura novedosa que combina el modelo espacial de estados de Mamba con una combinación de modelos expertos para aprovechar los beneficios que ofrecen ambos marcos. Los experimentos con BlackMamba han demostrado que supera el marco Mamba existente y las líneas base del transformador tanto en el entrenamiento de FLOP como en la inferencia. El rendimiento excepcional del marco BlackMamba demuestra que es capaz de heredar y combinar las capacidades de los marcos Mamba y MoE excepcionalmente bien, ya que combina la inferencia rápida y barata de MoE con la generación de complejidad lineal de Mamba. Hemos hablado sobre cómo la arquitectura del marco BlackMamba puede superar los modelos de lenguaje grande fuertemente entrenados, el marco Mamba existente y los modelos mixtos de expertos en términos de entrenamiento de FLOP y costo de inferencia. Además, el marco BlackMamba también hereda los FLOP de generación y el entrenamiento reducido de los modelos Mixture of Expert y el marco Mamba simultáneamente.