Introducción a la arquitectura Mamba LLM: Un nuevo paradigma en el aprendizaje automático

Los modelos lingüísticos son un tipo de modelo de aprendizaje automático entrenado para realizar una distribución de probabilidad sobre el lenguaje natural. Su arquitectura se compone principalmente de múltiples capas de redes neuronales, como capas recurrentes, capas feedforward, capas de incrustación y capas de atención. Estas capas se combinan para procesar un texto de entrada dado y generar predicciones de salida.

A finales de 2023, investigadores de Carnegie Mellon y de la Universidad de Princeton publicaron un trabajo de investigación que revelaba una nueva arquitectura para Grandes Modelos Lingüísticos (LLM) llamada Mamba. Mamba es una nueva arquitectura de modelos de espacio de estados que se ocupa del modelado de secuencias. Se desarrolló para abordar algunas limitaciones de los modelos de transformador , sobre todo en el procesamiento de secuencias largas, y ha mostrado un rendimiento prometedor.

Exploremos la arquitectura LLM Mamba y su importancia en el aprendizaje automático.

¿Qué es la Mamba?

Mamba es una nueva arquitectura LLM que integra el modelo de secuencia de Espacio de Estado Estructurado (S4) para gestionar secuencias de datos largas. Combinando las mejores características de los modelos recurrentes, convolucionales y de tiempo continuo, el S4 puede simular eficaz y eficientemente las dependencias a largo plazo. Esto le permite manejar datos muestreados irregularmente, tener un contexto ilimitado y mantener la eficiencia computacional durante el entrenamiento y la prueba.

Ampliando el paradigma S4, Mamba aporta varias mejoras dignas de mención, especialmente en el manejo de las operaciones variables en el tiempo. Su arquitectura gira en torno a un mecanismo de selección especial que modifica los parámetros del modelo de espacio de estados estructurado (SSM) en función de la entrada.

Como resultado, Mamba puede filtrar con éxito los datos menos importantes centrándose sólo en la información crucial dentro de las secuencias. Según Wikipedia, "El modelo pasa de un marco invariable en el tiempo a otro variable en el tiempo, lo que repercute tanto en el cálculo como en la eficacia del sistema".

Características principales e innovaciones

La desviación de Mamba de la atención convencional y los bloques MLP la distinguen. Esta simplificación da como resultado un modelo más ligero, más rápido y que escala linealmente con la longitud de la secuencia, lo cual es una hazaña que no ha conseguido ninguno de sus predecesores.

Los componentes clave de Mamba son

• Espacios de Estado Selectivos (SSM): Los modelos recurrentes, que procesan la información selectivamente en función de la entrada actual, son la base de los MSS Mamba. Esto les permite filtrar los datos extraños y concentrarse en la información pertinente, lo que puede dar lugar a un procesamiento más eficaz.
• Arquitectura simplificada: Mamba sustituye los intrincados bloques de atención y MLP de Transformers por un único bloque SSM cohesionado. Con ello se pretende acelerar la inferencia y reducir la complejidad computacional.
• Paralelismo consciente del hardware: Mamba puede funcionar incluso mejor porque utiliza un modo recurrente con un algoritmo paralelo creado especialmente para la eficiencia del hardware.

Otro componente clave es la Invariancia Temporal Lineal (ITL); la ITL es una de las características fundamentales de los modelos S4. Esta característica sugiere que los parámetros del modelo permanecen constantes a lo largo de todos los pasos temporales, manteniendo la coherencia en la dinámica del modelo. Construir modelos de secuencias es más fácil y eficaz con LTI, que es la base de la recurrencia y las convoluciones.

Detalles de la arquitectura de Mamba LLM

La arquitectura de Mamba subraya aún más la importancia de los avances realizados en el aprendizaje automático. Modifica la forma en que los modelos procesan las secuencias introduciendo una capa de modelo de espacio de estados seleccionados (SSM). Esto permite a Mamba hacer dos cosas extremadamente importantes:

1. Céntrate en la información relevante: Mamba puede priorizar los datos más predictivos para la tarea asignando un peso diferente a cada entrada.

2. Adaptarse dinámicamente a las entradas - Gracias a la capacidad del modelo para adaptarse a las entradas, Mamba puede manejar fácilmente diversos trabajos de modelado de secuencias.

Como resultado, Mamba puede manejar secuencias con una eficacia nunca vista, lo que lo convierte en la elección perfecta para tareas que impliquen largas secuencias de datos.

La filosofía de diseño de Mamba se basa en el conocimiento de las capacidades del hardware contemporáneo. Está diseñado para aprovechar al máximo la potencia de cálculo de la GPU, garantizando:

• Uso optimizado de la memoria: Los tiempos de transmisión de datos reducidos y el procesamiento más rápido se consiguen diseñando la expansión de estado de Mamba para que quepa dentro de la memoria de gran ancho de banda (HBM) de las GPU.

• Procesamiento paralelo maximizado: Mamba alcanza un nivel de rendimiento que establece un nuevo punto de referencia para los modelos de secuencias al coordinar sus cálculos con la naturaleza paralela del GPU Computing.

Mamba vs Transformers

La introducción de Transformadores, como el GPT-4, tomó el campo del procesamiento del lenguaje natural (PLN) y estableció puntos de referencia para varias tareas de lenguaje natural. Las secuencias largas han sido durante mucho tiempo una espina en el costado de los transformadores, ya que dificultan significativamente su eficiencia.

En esta deficiencia es donde destaca Mamba. En concreto, mamba puede procesar secuencias largas más rápidamente que los transformadores y lo hace de forma más sencilla gracias a su arquitectura única.

Arquitectura del transformador

Los Transformadores son muy hábiles trabajando con secuencias de datos, como el texto para los modelos lingüísticos. Procesan secuencias completas simultáneamente, a diferencia de los modelos anteriores que procesaban los datos secuencialmente. Esta característica innata les permite captar relaciones intrincadas dentro de los datos.

Utilizan un mecanismo de atención que permite al modelo concentrarse en varios segmentos de la secuencia mientras genera predicciones. Para calcular esta atención se utilizan tres conjuntos de pesos: valores, claves y consultas que se obtienen de los datos de entrada.

Cada elemento de una secuencia se pondera en relación con todos los demás para indicar cuánto peso -o "atención"- debe tener para pronosticar el siguiente elemento de la serie.

Los transformadores constan de dos bloques principales: el descodificador, que crea la salida, y el codificador, que procesa los datos de entrada.

El codificador consta de varias capas, cada una de ellas con dos subcapas: una red básica de avance totalmente conectada en función de la posición y un mecanismo de autoatención multicabezal. Para facilitar el entrenamiento de las redes profundas, se utilizan conexiones residuales y normalización en cada subcapa.

Al igual que el codificador, el descodificador consta de dos capas con dos subcapas, pero además añade una tercera subcapa que se encarga de la atención multicabezal sobre la salida del codificador. La propiedad autorregresiva del descodificador se mantiene debido a su naturaleza secuencial, que limita las predicciones para una posición a tener en cuenta sólo las posiciones anteriores.

Así, Transformers intenta abordar el problema de las secuencias largas utilizando procesos de atención más intrincados, pero Mamba adopta un enfoque diferente.

Arquitectura Mamba

Mamba aprovecha los espacios de estado selectivos. Este método resuelve las ineficiencias informáticas de los Transformers con secuencias largas.

La arquitectura de Mamba hace posible una inferencia más rápida y un escalado lineal de la longitud de las secuencias, creando un nuevo paradigma para el modelado de secuencias que puede resultar más eficaz a medida que las secuencias se hacen más largas.

Como ya hemos profundizado en la arquitectura de Mamba, no entraremos en ella aquí.

He aquí un gráfico de Wikipedia para conceptualizar mejor la comparación entre Mamba y Transformers:

Merece la pena señalar que, a pesar de las muchas ventajas que tienen los SSM sobre los Transformadores, estos últimos pueden manejar secuencias mucho más largas que las que los SSM pueden almacenar en memoria, necesitan muchos menos datos para aprender tareas similares y superan a los SSM en tareas que requieren recuperar o copiar el contexto de entrada, incluso con menos parámetros. 

Primeros pasos con Mamba

Si te interesa jugar con Mamba o aprovecharlo en un proyecto, debes tener lo siguiente:

• Linux
• NVIDIA GPU
• PyTorch 1.12+
• CUDA 11.6+

Para instalar los paquetes necesarios desde el repositorio Mamba, sigue unas sencillas instrucciones de pip:

• [Opción] pip install causal-conv1d>=1.2.0: una implementación eficiente de una capa Conv1d causal simple utilizada dentro del bloque Mamba.
• pip install mamba-ssm: el núcleo del paquete Mamba.

También se puede construir desde el código fuente con pip install . desde este repositorio.

Si las versiones de PyTorch causan problemas de compatibilidad, se puede utilizar pip con el interruptor --no-build-isolation para ayudar. Estos modelos se entrenaron en grandes conjuntos de datos, como los conjuntos de datos Pile y SlimPajama, y se construyeron para satisfacer diversos requisitos informáticos y puntos de referencia de rendimiento.

El modelo Mamba tiene varios niveles de interfaz, pero el módulo principal es el bloque de arquitectura Mamba que envuelve el SSM selectivo.

Una vez instalado todo, se puede utilizar como sigue:

# Source: Mamba Repository
import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

Aplicaciones de la Mamba

La introducción de Mamba LLM supone un cambio potencial significativo en el espacio de la arquitectura LLM. Más rápida, eficiente y escalable, Mamba maneja sin esfuerzo secuencias largas con estándares de alto rendimiento, lo que explica por qué está llamada a desempeñar un papel fundamental en la configuración del futuro de los sistemas sofisticados de IA.

En concreto, la próxima oleada de innovaciones en IA puede surgir de la eficacia y el rendimiento de Mamba, que allana el camino para la creación de modelos y aplicaciones cada vez más complejos. Su influencia potencial es enorme, pues abarca aplicaciones de procesamiento de audio y voz, análisis de textos largos, creación de contenidos, traducción de idiomas en tiempo real y mucho más.

Entre las industrias que esto podría revolucionar se incluyen:

• Sanidad: Mamba puede acelerar el proceso de desarrollo de medicamentos personalizados para la salud analizando rápidamente los datos genéticos.

• Finanzas: Mamba puede desplegarse para analizar las tendencias del mercado a largo plazo, lo que da lugar a previsiones bursátiles más precisas.

• Atención al cliente: Mamba tiene la capacidad de potenciar chatbots que controlan las conversaciones de larga duración, mejorando así las comunicaciones con los clientes.

El camino por delante para Mamba

Mamba destaca como un faro innovador que señala el camino hacia nuevas oportunidades para abordar problemas complejos de modelado de secuencias. Su estreno es un paso hacia sistemas más inteligentes, escalables y eficaces, capaces de comprender y manejar secuencias con una profundidad nunca vista.

Aunque Mamba ha demostrado ser un hito técnico importante, su éxito no depende únicamente de sus capacidades tecnológicas: la investigación colaborativa, los recursos compartidos y las contribuciones de código abierto desempeñan un papel esencial:

• Contribuciones de código abierto: Los modelos más resistentes y adaptables pueden ser el resultado de animar a los investigadores y desarrolladores a contribuir a la base de código Mamba.

• Recursos compartidos: La comunidad puede acelerar el progreso poniendo en común sus conocimientos y recursos y compartiendo modelos preentrenados.

• Investigación en colaboración: Las colaboraciones entre instituciones académicas y empresas pueden ampliar las capacidades de Mamba.

Conclusión

Mamba no sólo proporciona una mejora incremental de los modelos de secuencia actuales, sino que redefine lo que es posible. Con su introducción, la historia de la inteligencia artificial será testigo de un nuevo capítulo en el que las ineficiencias informáticas y las restricciones de longitud de las secuencias quedarán por fin obsoletas.

En los últimos años, hemos asistido a la evolución de la IA desde las RNN a los Transformers y ahora a Mamba, cada paso más cerca de conseguir una IA capaz de un pensamiento profundo y un procesamiento de la información comparables a los humanos. Mamba encarna el espíritu revolucionario que impulsa el campo de la IA con su enfoque de espacio de estados seleccionado y su escalado en tiempo lineal.

Mamba marca el inicio de un prometedor desarrollo de la inteligencia artificial. Se trata de un paradigma diseñado para el futuro y destinado a influir significativamente en la IA con su potencial ilimitado.

Para seguir aprendiendo, consulta

• Cómo crear aplicaciones LLM con Langchain
• Cómo entrenar un LLM con PyTorch