El desarrollo de Modelos de Lenguaje Grande (LLM) construidos a partir de modelos de transformador solo decodificador ha jugado un papel crucial en la transformación del dominio de Procesamiento de Lenguaje Natural (NLP) y ha avanzado en diversas aplicaciones de aprendizaje profundo, incluyendo aprendizaje de refuerzo, análisis de series temporales, procesamiento de imágenes y mucho más. Sin embargo, a pesar de su escalabilidad y fuerte rendimiento, los LLM construidos a partir de modelos de transformador solo decodificador aún enfrentan importantes limitaciones. Aunque son expresivos, el mecanismo de atención en los LLM derivados de transformadores requiere recursos computacionales intensivos durante la inferencia y el entrenamiento, lo que requiere una gran cantidad de memoria para la longitud de la secuencia y operaciones (FLOPs) que aumentan cuadráticamente. Esta alta demanda computacional limita la longitud del contexto de los modelos de transformador, hace que las tareas de generación autoregresiva sean proporcionalmente costosas con la escala y obstaculiza la capacidad de los modelos para aprender de flujos de datos continuos y procesar secuencias de longitud ilimitada de manera eficiente.

En tiempos recientes, Modelos de Espacio de Estado (MEE) han demostrado capacidades y rendimiento notables, compitiendo con modelos de arquitectura de transformador en benchmarks de modelado a gran escala mientras logran una complejidad de memoria como función de la longitud de la secuencia y tiempo lineal. Además, Mamba, un MEE recientemente lanzado, ha mostrado un rendimiento destacado en una variedad de tareas de modelado de lenguaje y procesamiento de secuencias largas. Al mismo tiempo, los modelos de Mezcla de Expertos (MoE) también han demostrado un rendimiento impresionante mientras reducen significativamente la latencia y los costos computacionales de inferencia, aunque a expensas de una huella de memoria más grande. Basándose en Mamba y MoE, este artículo discutirá BlackMamba, una arquitectura novedosa que combina el Modelo de Espacio de Estado Mamba con modelos MoE para aprovechar los beneficios ofrecidos por ambos marcos. Los experimentos en BlackMamba han demostrado su capacidad para superar al marco de trabajo Mamba existente y a las líneas base de transformadores en ambos entrenamiento FLOPs e inferencia. El rendimiento excepcional del marco de trabajo BlackMamba muestra que puede combinar efectivamente las capacidades de los marcos de trabajo Mamba y MoE, ofreciendo inferencia rápida y rentable desde MoE con generación de complejidad lineal desde Mamba.

Este artículo tiene como objetivo cubrir el marco de trabajo BlackMamba en profundidad. Exploramos el mecanismo, la metodología y la arquitectura del marco de trabajo, junto con su comparación con marcos de trabajo de generación de imágenes y videos de última generación. Comencemos.

BlackMamba: Introducción a MoE para Modelos de Espacio de Estado

El progreso de los Modelos de Lenguaje Grande (LLM), particularmente aquellos basados en arquitecturas de transformador solo decodificador, ha influido notablemente en el campo de Procesamiento de Lenguaje Natural (NLP) y se ha expandido a diversas aplicaciones de aprendizaje profundo, incluyendo aprendizaje de refuerzo, análisis de series temporales, procesamiento de imágenes y más allá. Sin embargo, a pesar de su escalabilidad y rendimiento robusto, estos LLM basados en transformadores solo decodificadores enfrentan desafíos notables. El mecanismo de atención, una característica clave de los LLM basados en transformadores, demanda recursos computacionales intensivos para la inferencia y el entrenamiento. Esto implica una necesidad de memoria que crece con la longitud de la secuencia y operaciones (FLOPs) que aumentan cuadráticamente. Tales demandas computacionales intensivas limitan la longitud del contexto de los modelos, elevan los costos de las tareas de generación autoregresiva a medida que el modelo se escala y obstaculizan la capacidad de los modelos para aprender de flujos de datos continuos o procesar secuencias de longitud ilimitada de manera eficiente.

Se han realizado esfuerzos significativos en los últimos años para superar estas limitaciones, y la atención se ha desviado hacia la creación de alternativas arquitectónicas a los modelos de transformador densos con MEE y MoE siendo las arquitecturas de modelos más prometedoras. El beneficio clave obtenido al favorecer los Modelos de Espacio de Estado sobre los modelos de arquitectura de transformador es la complejidad computacional lineal con respecto a la longitud de la secuencia de entrada ofrecida por los MEE en lugar de la complejidad cuadrática ofrecida por los transformadores. Teóricamente, la complejidad computacional lineal con respecto a la longitud de la secuencia de entrada permite que los Modelos de Espacio de Estado procesen secuencias más largas que los modelos de arquitectura de transformador para un presupuesto de FLOPS o operaciones de punto flotante por segundo determinado, y rinden la generación autoregresiva constante en cómputo sin una caché KV. Los MEE recientemente desarrollados, incluyendo Mamba, RetNet y algunos otros, han demostrado una inferencia y entrenamiento de secuencias largas eficientes, junto con un rendimiento competitivo en tareas de modelado de lenguaje comparable a los transformadores con propiedades de escalabilidad similares. Por otro lado, los modelos de Mezcla de Expertos han ganado popularidad como una alternativa a los transformadores densos, ya que facilitan una reducción significativa en los FLOPs de inferencia y entrenamiento esenciales para lograr una calidad comparable a un modelo denso. Los modelos MoE operan activando solo una selección esparsa de los parámetros totales durante un solo paso hacia adelante. Utilizan una función de enrutamiento para determinar qué “expertos” se llaman a acción en función del contexto dado. Este enfoque crea una separación entre el costo computacional de la inferencia y la cantidad total de parámetros, lo que permite un mejor rendimiento dentro de un presupuesto de inferencia fijo, aunque con un aumento en la cantidad de parámetros y una mayor huella de memoria.

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

Mamba: Logra una complejidad computacional lineal relativa a la longitud de la secuencia de entrada para ambas fases de entrenamiento e inferencia. Permite que la generación autoregresiva ocurra en un marco de tiempo constante y con uso de memoria constante.

MoE: Ofrece una velocidad de inferencia y eficiencia computacional de entrenamiento comparable a un modelo de referencia denso más pequeño, mientras mantiene un nivel de calidad de modelo que rivaliza con el de un modelo con un número equivalente de parámetros al del modelo denso.

Con eso dicho, es esencial afirmar que los modelos de arquitectura de transformador aún son los mejores, y han demostrado un rendimiento consistente y notable en tareas de modelado de lenguaje y procesamiento de secuencias. En su núcleo, la arquitectura de transformador emplea autoatención que realiza una comparación de todo a todo cuadrática de las similitudes del producto de puntos entre las incrustaciones de diferentes tokens en una secuencia, y realiza un mapeo lineal a un vector de salida. El modelo de transformador consiste en bloques de autoatención apilados entre bloques de MLP o Perceptrón Multicapa que consisten en 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 Estado 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 Estado se alinea más con las Redes Neuronales Recurrentes y las Redes Neuronales Convolucionales que con la arquitectura basada en 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 utilizando un escaneo asociativo o de convolución. En escenarios prácticos, la naturaleza recurrente de los Modelos de Espacio de Estado ha sido la razón por la que aún no se ha adoptado en hardware de AI altamente paralelo como las GPU. Sin embargo, el surgimiento de MEE como RWKV y Mamba han utilizado kernels de escaneo paralelo para mapear operaciones recurrentes de manera eficiente a las GPU, lo que facilita el entrenamiento de arquitecturas novedosas con una eficiencia comparable a la de los modelos de transformador.

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 obstaculiza el razonamiento y la comprensión sobre contextos muy largos. Las innovaciones recientes han introducido la idea de extender la longitud del contexto, lo que permite que los transformadores se entrenen en una escala factible antes de aplicarse a contextos mucho más largos durante la inferencia. A pesar de estos avances, el proceso de inferencia aún requiere una cantidad considerable de recursos computacionales y memoria, especialmente para mantener la caché de Clave-Valor (KV), lo que lo convierte en una empresa intensiva en recursos. Los esfuerzos de investigación recientes se han centrado en mejorar las capacidades expresivas de los modelos de espacio de estado mediante la incorporación de mecanismos de puerta dependientes de la entrada, similares a las matrices de Consulta, Clave, Valor (QKV) encontradas en los mecanismos de atención.

Estos esfuerzos apuntan a preservar la progresión lineal inherente de la recursividad del espacio de estado, lo que permite una ejecución eficiente a través de la convolución o un proceso de escaneo selectivo. Este enfoque reduce significativamente la brecha de rendimiento con los transformadores en aplicaciones prácticas. Entre estos avances, Mamba se destaca como un modelo de espacio de estado que refleja los objetivos de la investigación previa, mostrando niveles de rendimiento impresionantes comparables a los de los transformadores a escalas de hasta 2.8 mil millones de parámetros. Logra esto aplicando una puerta dependiente de la entrada a las entradas de la recursividad del modelo de espacio de estado (MEE), asegurando al mismo tiempo un cálculo eficiente a través del uso de kernels de escaneo selectivo personalizados.

Modelos de Mezcla de Expertos

Los modelos de Mezcla de Expertos (MoE) logran una separación entre el costo de inferencia y la cantidad total de parámetros al activar selectivamente los parámetros durante el paso hacia adelante. En lugar de utilizar todos los parámetros, estos modelos dirigen los tokens a expertos de Multicapa Perceptrón (MLP) específicos. Idealmente, cada experto está diseñado para procesar un tipo particular de entrada, con un mecanismo de enrutamiento, esencialmente una red neural compacta, que determina el experto más adecuado para cada token. Este enfoque apunta a 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 significativamente reducidas. Normalmente, el enrutador es un mapeo de las capas lineales desde los tokens hasta los índices de expertos, con cada experto siendo simplemente un MLP de transformador estándar. Sin embargo, los desarrolladores aún no han determinado 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 luchan con el equilibrio de carga y la estabilidad del entrenamiento entre diferentes expertos para la eficiencia del hardware.

Arquitectura

En su núcleo, BlackMamba emplea un modelo de transformador estándar que consiste en bloques de MLP intercalados con bloques de atención agregados en secuencia a lo largo de un flujo residual. Ahora, la mayoría de los modelos de Mezcla de Expertos simplemente reemplazan los bloques de MLP con una capa de expertos enrutada. Por otro lado, el marco de trabajo BlackMamba no solo reemplaza el bloque de MLP en el transformador con una capa de expertos enrutada, sino que también reemplaza la capa de atención con una capa de Modelo de Espacio de Estado Mamba. La arquitectura del marco de trabajo BlackMamba se demuestra en la siguiente figura.

Entrenamiento y Conjunto de Datos

El modelo BlackMamba se entrena en 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 MLP de expertos. El marco de trabajo se entrena con 8 expertos, un número que los desarrolladores encontraron que es 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 de trabajo BlackMamba consiste en una mezcla de conjuntos de datos de código abierto existentes, incluyendo 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 de trabajo BlackMamba. En general, 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 en los mismos datos de entrenamiento. El marco de trabajo BlackMamba puede superar tanto a Mamba como a los modelos de transformador en tamaño de modelo de paso hacia adelante idéntico en el tiempo de inferencia, así como en el entrenamiento de operaciones de punto flotante por segundo. La siguiente figura demuestra el tiempo necesario para generar una secuencia de una longitud determinada de manera autoregresiva a partir de una señal de inicio de un token como función de la longitud de la secuencia.

Además, los beneficios de latencia de ambos modelos de Mezcla de Expertos y Mamba se combinan en el marco de trabajo BlackMamba, lo que resulta en tiempos de inferencia significativamente más rápidos en comparación con los modelos de transformador, los modelos de Mamba puros y los modelos de MoE. Además, la ventaja de inferencia del marco de trabajo BlackMamba es directamente proporcional a las longitudes de la secuencia, lo que hace que BlackMamba sea extremadamente efectivo 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 se puede ver, la mayoría de las capas demuestran un alto nivel de equilibrio de expertos como resultado del algoritmo Sinkhorn mejorado implementado por los modelos BlackMamba.

La siguiente tabla cubre las puntuaciones de evaluación del marco de trabajo BlackMamba en comparación con una serie de modelos de lenguaje preentrenados de código abierto. Como se puede observar, el marco de trabajo BlackMamba puede competir y superar a la mayoría de los marcos en todas las líneas base. Además, es digno de destacar que los modelos que superan a BlackMamba tienen un número significativamente mayor de parámetros, y la brecha en el rendimiento es mínima, lo que indica la capacidad del marco de trabajo BlackMamba con menos parámetros.

Pensamientos Finales

En este artículo, hemos hablado sobre BlackMamba, una arquitectura novedosa que combina el Modelo de Espacio de Estado Mamba con modelos de Mezcla de Expertos para aprovechar los beneficios ofrecidos por ambos marcos. Los experimentos en BlackMamba han demostrado su capacidad para superar al marco de trabajo Mamba existente y a las líneas base de transformadores en ambos entrenamiento FLOPs e inferencia. El rendimiento excepcional del marco de trabajo BlackMamba muestra que puede combinar efectivamente las capacidades de los marcos de trabajo Mamba y MoE, ofreciendo inferencia rápida y rentable desde MoE con generación de complejidad lineal desde Mamba. Hemos hablado sobre cómo la arquitectura del marco de trabajo BlackMamba puede superar a los Modelos de Lenguaje Grande entrenados, al marco de trabajo Mamba existente y a los modelos de Mezcla de Expertos en términos de entrenamiento FLOPs y costo de inferencia. Además, el marco de trabajo BlackMamba también hereda la generación FLOPs y el entrenamiento reducido de ambos modelos de Mezcla de Expertos y Mamba simultáneamente.