Modelos y plataformas de IA
Mamba: Redefiniendo el modelado de secuencias y superando la arquitectura de Transformadores
Añade Unite.AI a tus fuentes preferidas en Google
En este artículo sobre Mamba, exploraremos cómo este innovador modelo de espacio de estado (SSM) revoluciona el modelado de secuencias. Desarrollado por Albert Gu y Tri Dao, Mamba se distingue por su eficiencia en el procesamiento de secuencias complejas en campos como el procesamiento de lenguaje, la genómica y el análisis de audio. Su modelado de secuencias lineales en el tiempo con espacios de estado selectivos garantiza un rendimiento excepcional en estas diversas modalidades.
Nos sumergiremos en la capacidad de Mamba para superar los desafíos computacionales que enfrentan los Transformadores tradicionales, especialmente con secuencias largas. Su enfoque selectivo en los modelos de espacio de estado permite una inferencia más rápida y una escalabilidad lineal con la longitud de la secuencia, lo que mejora significativamente el rendimiento.
La singularidad de Mamba radica en su capacidad de procesamiento rápido, su capa de SSM selectiva y su diseño inspirado en FlashAttention, que permite a Mamba superar a muchos modelos existentes, incluidos aquellos basados en el enfoque de transformadores, lo que lo convierte en un avance notable en el aprendizaje automático.
Transformadores vs Mamba
Los Transformadores, como GPT-4, han establecido estándares en el procesamiento de lenguaje natural. Sin embargo, su eficiencia disminuye con secuencias más largas. Es aquí donde Mamba se adelanta, con su capacidad para procesar secuencias largas de manera más eficiente y su arquitectura única que simplifica todo el proceso.
Los Transformadores son hábiles para manejar secuencias de datos, como texto para modelos de lenguaje. A diferencia de los modelos anteriores que procesaban los datos secuencialmente, los Transformadores procesan las secuencias enteras simultáneamente, lo que les permite capturar relaciones complejas dentro de los datos.
Utilizan un mecanismo de atención, que permite al modelo enfocarse en diferentes partes de la secuencia al hacer predicciones.
Esta atención se calcula utilizando tres conjuntos de pesos: consultas, claves y valores, derivados de los datos de entrada. Cada elemento en una secuencia se compara con todos los demás elementos, proporcionando un peso que indica la importancia, o “atención”, que cada elemento debe recibir al predecir el siguiente elemento en la secuencia.
Los Transformadores mantienen dos bloques principales: el codificador, que procesa los datos de entrada, y el decodificador, que genera la salida. El codificador consta de múltiples capas, cada una con dos subcapas: un mecanismo de autoatención multiencabezado y una red neuronal fully connected simple y position-wise. La normalización y las conexiones residuales se utilizan en cada subcapa para ayudar en el entrenamiento de redes profundas.
El decodificador también tiene capas con dos subcapas similares al codificador, pero agrega una tercera subcapa que realiza una atención multiencabezado sobre la salida del codificador. La naturaleza secuencial del decodificador garantiza que las predicciones para una posición solo puedan considerar posiciones anteriores, preservando la propiedad autoregresiva.
En contraste con los Transformadores, el modelo Mamba adopta un enfoque diferente. Mientras que los Transformadores abordan el problema de las secuencias largas utilizando mecanismos de atención más complejos, Mamba utiliza espacios de estado selectivos, lo que proporciona una forma más eficiente de procesar las secuencias.
Aquí hay una visión general de alto nivel de cómo funciona un transformador:
- Procesamiento de entrada: Los Transformadores codifican primero los datos de entrada en un formato que el modelo pueda entender, a menudo utilizando incrustaciones que también incorporan la posición de cada elemento en la secuencia.
- Mecanismo de atención: En su núcleo, el mecanismo de atención calcula una puntuación que representa cuánta atención poner en otras partes de la secuencia de entrada cuando se entiende un elemento actual.
- Arquitectura codificador-decodificador: El modelo de transformador consta de un codificador para procesar la entrada y un decodificador para generar la salida. Cada uno consta de múltiples capas que refinan la comprensión del modelo de la entrada.
- Atención multiencabezado: Dentro de ambos el codificador y el decodificador, la atención multiencabezado permite al modelo atender simultáneamente a diferentes partes de la secuencia desde diferentes espacios de representación, mejorando su capacidad para aprender de contextos diversos.
- Redes neuronales position-wise: Después de la atención, una red neuronal simple procesa la salida de cada posición de manera separada e idéntica. Esto se combina con la entrada a través de una conexión residual y se sigue de una normalización de capa.
- Generación de salida: El decodificador luego predice una secuencia de salida, influenciada por el contexto del codificador y lo que ha generado hasta ahora.
La capacidad del transformador para manejar secuencias en paralelo y su robusto mecanismo de atención lo hacen poderoso para tareas como la traducción y la generación de texto.
En contraste, el modelo Mamba opera de manera diferente al utilizar espacios de estado selectivos para procesar secuencias. Este enfoque aborda la ineficiencia computacional en los Transformadores al tratar con secuencias largas. El diseño de Mamba permite una inferencia más rápida y se escala linealmente con la longitud de la secuencia, estableciendo un nuevo paradigma para el modelado de secuencias que podría ser más eficiente, especialmente a medida que las secuencias se vuelven cada vez más largas.
Mamba
Lo que hace que Mamba sea verdaderamente único es su alejamiento de los bloques de atención y MLP tradicionales. Esta simplificación conduce a un modelo más ligero y rápido que se escala linealmente con la longitud de la secuencia, una hazaña sin precedentes para sus predecesores.
Las características clave de Mamba incluyen:
- SSM selectivos: Estos permiten a Mamba filtrar información irrelevante y centrarse en datos relevantes, mejorando su manejo de secuencias. Esta selectividad es crucial para el razonamiento basado en contenido eficiente.
- Algoritmo consciente del hardware: Mamba utiliza un algoritmo paralelo que se optimiza para el hardware moderno, especialmente las GPU. Este diseño permite una computación más rápida y reduce los requisitos de memoria en comparación con los modelos tradicionales.
- Arquitectura simplificada: Al integrar SSM selectivos y eliminar los bloques de atención y MLP, Mamba ofrece una estructura más simple y homogénea. Esto conduce a una mejor escalabilidad y rendimiento.
Mamba ha demostrado un rendimiento superior en varios dominios, incluyendo lenguaje, audio y genómica, destacándose tanto en la preentrenación como en tareas específicas del dominio. Por ejemplo, en el modelado de lenguaje, Mamba iguala o supera el rendimiento de modelos de Transformadores más grandes.
El código y los modelos preentrenados de Mamba están disponibles abiertamente para la comunidad en GitHub.
Tareas de copia estándar son fáciles para los modelos lineales. La copia selectiva y las cabezas de inducción requieren memoria dinámica y consciente de contenido para los LLM.
Los modelos de espacio de estado estructurados (S4) han surgido recientemente como una clase prometedora de modelos de secuencias, que abarcan rasgos de las RNN, las CNN y los modelos de espacio de estado clásicos. Los modelos S4 derivan su inspiración de sistemas continuos, específicamente de un tipo de sistema que mapea funciones unidimensionales o secuencias a través de un estado latente implícito. En el contexto del aprendizaje profundo, representan una innovación significativa, proporcionando una nueva metodología para diseñar modelos de secuencias que sean eficientes y altamente adaptables.
La dinámica de los modelos S4
SSM (S4) Este es el modelo de espacio de estado estructurado básico. Toma una secuencia x y produce una salida y utilizando parámetros aprendidos A, B, C y un parámetro de retraso Δ. La transformación implica discretizar los parámetros (convirtiendo funciones continuas en discretas) y aplicar la operación SSM, que es invariante en el tiempo, lo que significa que no cambia a lo largo de los diferentes pasos de tiempo.
La importancia de la discretización
La discretización es un proceso clave que transforma los parámetros continuos en discretos a través de fórmulas fijas, lo que permite a los modelos S4 mantener una conexión con los sistemas de tiempo continuo. Esto dota a los modelos de propiedades adicionales, como la invarianza de resolución, y garantiza una normalización adecuada, lo que mejora la estabilidad y el rendimiento del modelo. La discretización también establece paralelismos con los mecanismos de puerta encontrados en las RNN, que son críticos para gestionar el flujo de información a lo largo de la dimensión de la secuencia.
Invarianza lineal en el tiempo (LTI)
Una característica fundamental de los modelos S4 es su invarianza lineal en el tiempo. Esta propiedad implica que la dinámica del modelo permanece consistente en el tiempo, con los parámetros fijos para todos los pasos de tiempo. La LTI es una piedra angular de la recurrencia y las convoluciones, ofreciendo un marco simplificado pero poderoso para construir modelos de secuencias.
Superar las limitaciones fundamentales
El marco S4 ha estado tradicionalmente limitado por su naturaleza LTI, lo que plantea desafíos en la modelización de datos que requieren dinámicas adaptativas. El reciente artículo de investigación presenta un enfoque que supera estas limitaciones al introducir parámetros que varían con el tiempo, lo que elimina la restricción de LTI. Esto permite a los modelos S4 manejar un conjunto más diverso de secuencias y tareas, expandiendo significativamente su aplicabilidad.
El término ‘modelo de espacio de estado’ abarca cualquier proceso recurrente que involucre un estado latente y se ha utilizado para describir varios conceptos en múltiples disciplinas. En el contexto del aprendizaje profundo, los modelos S4, o espacios de estado estructurados, se refieren a una clase específica de modelos que se han optimizado para una computación eficiente mientras retienen la capacidad de modelar secuencias complejas.
Los modelos S4 se pueden integrar en arquitecturas de redes neuronales de extremo a extremo, funcionando como transformaciones de secuencias independientes. Pueden considerarse análogos a las capas de convolución en las CNN, proporcionando la base para el modelado de secuencias en una variedad de arquitecturas de redes neuronales.
SSM vs SSM + Selección
Motivación para la selectividad en el modelado de secuencias
SSM estructurados
El artículo argumenta que un aspecto fundamental del modelado de secuencias es la compresión del contexto en un estado manejable. Los modelos que pueden enfocarse selectivamente en los datos de entrada o filtrarlos proporcionan una forma más efectiva de mantener este estado comprimido, lo que conduce a modelos de secuencias más eficientes y poderosos. Esta selectividad es vital para que los modelos puedan controlar adaptativamente cómo fluye la información a lo largo de la dimensión de la secuencia, una capacidad esencial para manejar tareas complejas en el modelado de lenguaje y más allá.
Los SSM selectivos mejoran los SSM convencionales al permitir que sus parámetros dependan de la entrada, lo que introduce un grado de adaptabilidad que anteriormente no era alcanzable con modelos de tiempo invariante. Esto da como resultado SSM que varían con el tiempo y que ya no pueden utilizar convoluciones para la computación eficiente, sino que dependen de un mecanismo de recurrencia lineal, una desviación significativa de los modelos tradicionales.
SSM + Selección (S6) Esta variante incluye un mecanismo de selección, agregando dependencia de la entrada a los parámetros B y C, y un parámetro de retraso Δ. Esto permite al modelo enfocarse selectivamente en ciertas partes de la secuencia de entrada x. Los parámetros se discretizan teniendo en cuenta la selección, y la operación SSM se aplica de manera que varía con el tiempo utilizando una operación de exploración, que procesa los elementos secuencialmente, ajustando el enfoque dinámicamente con el tiempo.
Aspectos destacados del rendimiento de Mamba
Mamba es el mejor de su clase en cada resultado de evaluación
En términos de rendimiento, Mamba sobresale tanto en velocidad de inferencia como en precisión. Su diseño permite una mejor utilización de contextos más largos, lo que se demuestra tanto en el modelado de ADN como en el modelado de audio, superando a modelos anteriores en tareas complejas que requieren dependencias de largo alcance. Su versatilidad también se destaca en evaluaciones de disparo cero en múltiples tareas, estableciendo un nuevo estándar para este tipo de modelos en términos de eficiencia y escalabilidad.
Empezando con Mamba
Para aquellos interesados en aprovechar Mamba, los requisitos técnicos incluyen un sistema operativo Linux, una GPU NVIDIA (NVDA ), PyTorch 1.12+ y CUDA 11.6+. La instalación implica simples comandos pip para instalar los paquetes necesarios desde el repositorio de Mamba. Si surgen problemas de compatibilidad con versiones de PyTorch, utilizar la bandera –no-build-isolation con pip puede ayudar. Estos modelos, entrenados en extensos conjuntos de datos como el Pile y el conjunto de datos SlimPajama, están diseñados para satisfacer diversas necesidades computacionales y estándares de rendimiento.
Mamba ofrece diferentes niveles de interfaces, desde la capa de SSM selectiva hasta el bloque de Mamba y las estructuras de modelo de lenguaje completas. El bloque de Mamba, que es el módulo principal de la arquitectura, utiliza una capa Conv1d causal y puede integrarse fácilmente en diseños de redes neuronales. El ejemplo de uso proporcionado en Python demuestra cómo instanciar un modelo Mamba y procesar datos a través de él, destacando la simplicidad y flexibilidad del sistema.
Los modelos preentrenados de Mamba están disponibles en Hugging Face, con tamaños que van desde 130M hasta 2.8B parámetros, entrenados en el extenso conjunto de datos Pile y el conjunto de datos SlimPajama. Estos modelos están diseñados para satisfacer diversas necesidades computacionales y de rendimiento, ajustándose a los estándares dimensionales de GPT-3. Los usuarios pueden esperar un alto rendimiento y precisión de estos modelos, lo que hace que Mamba sea una opción competitiva para diversas aplicaciones, incluyendo pero no limitado a modelado de lenguaje.
El impacto de Mamba
Mamba representa un salto adelante en el modelado de secuencias, ofreciendo una poderosa alternativa a las arquitecturas de Transformadores para el procesamiento de datos densos en información. Su diseño se alinea con las demandas del hardware moderno, optimizando tanto el uso de memoria como las capacidades de procesamiento paralelo. La disponibilidad de código abierto de Mamba y sus modelos preentrenados lo convierten en una herramienta accesible y robusta para investigadores y desarrolladores en el campo del aprendizaje automático y el aprendizaje profundo.
