Inteligencia artificial
Mamba: redefiniendo el modelado de secuencias y superando la arquitectura de los transformadores
En este artículo sobre Mamba, exploraremos cómo este innovador modelo de espacio de estados (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 del lenguaje, la genómica y el análisis de audio. Su modelado de secuencia de tiempo lineal con espacios de estado selectivos garantiza un rendimiento excepcional en estas diversas modalidades.
Profundizaremos en la capacidad de Mamba para superar los desafíos computacionales que enfrentan los Transformers tradicionales, especialmente con secuencias largas. Su enfoque selectivo en modelos de espacio de estados permite una inferencia más rápida y un escalado lineal con longitud de secuencia, lo que mejora significativamente el rendimiento.
La singularidad de Mamba radica en su rápida capacidad de procesamiento, su capa SSM selectiva y su diseño compatible con hardware inspirado en FlashAttention. Estas características permiten a Mamba superar a muchos modelos existentes, incluidos aquellos basados en el enfoque del transformador, lo que lo convierte en un avance notable en el aprendizaje automático.
Transformadores vs. Mamba
Los transformadores, como GPT-4, han establecido puntos de referencia en el procesamiento del lenguaje natural. Sin embargo, su eficiencia disminuye con secuencias más largas. Aquí es donde Mamba da un paso adelante, con su capacidad para procesar secuencias largas de manera más eficiente y su arquitectura única que simplifica todo el proceso.
Transformadores expertos en el manejo de secuencias de datos, como texto para modelos de lenguaje. A diferencia de los modelos anteriores que procesaban datos de forma secuencial, Transformers procesa secuencias completas simultáneamente, lo que les permite capturar relaciones complejas dentro de los datos.
Utilizan un mecanismo de atención, que permite que el modelo se centre en diferentes partes de la secuencia al realizar predicciones.
Esta atención se calcula utilizando tres conjuntos de ponderaciones: consultas, claves y valores, derivados de los datos de entrada. Cada elemento de una secuencia se compara con todos los demás, lo que proporciona un peso que indica la importancia o "atención" que cada elemento debe recibir al predecir el siguiente elemento de 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 varias capas, cada una de las cuales contiene dos subcapas: un mecanismo de autoatención de múltiples cabezales y una red de alimentación directa simple, totalmente conectada en función de la posición. 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 atención de múltiples cabezales sobre la salida del codificador. La naturaleza secuencial del decodificador garantiza que las predicciones para una posición sólo puedan considerar posiciones anteriores, preservando la propiedad autorregresiva.
A diferencia de Transformers, el modelo Mamba adopta un enfoque diferente. Mientras que Transformers aborda el problema de las secuencias largas mediante el uso de mecanismos de atención más complejos, Mamba utiliza espacios de estado selectivos, lo que proporciona una visión más computacional.
Aquí hay una descripción general de alto nivel de cómo funciona un transformador:
- Procesamiento de entrada: Los transformadores primero codifican los datos de entrada en un formato que el modelo pueda comprender, a menudo utilizando incrustaciones que también incorporan la posición de cada elemento en la secuencia.
- Mecanismo de atención: En esencia, el mecanismo de atención calcula una puntuación que representa cuánto enfoque poner en otras partes de la secuencia de entrada al comprender un elemento actual.
- Arquitectura codificador-decodificador: El modelo de transformador se compone 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 de la entrada por parte del modelo.
- Atención multicabezal: Tanto dentro del codificador como del decodificador, la atención de múltiples cabezas permite que el modelo preste atención simultáneamente a diferentes partes de la secuencia desde diferentes espacios de representación, mejorando su capacidad para aprender de diversos contextos.
- Redes de retroalimentación por posición: Después de la atención, una red neuronal simple procesa la salida de cada posición por separado e idénticamente. Esto se combina con la entrada a través de una conexión residual y seguido de la normalización de capas.
- Generación de salida: Luego, el decodificador predice una secuencia de salida, influenciada por el contexto del codificador y lo que ha generado hasta el momento.
La capacidad del transformador para manejar secuencias en paralelo y su robusto mecanismo de atención lo hacen poderoso para tareas como traducción y generación de texto.
Por el contrario, el modelo Mamba opera de manera diferente al utilizar espacios de estados selectivos para procesar secuencias. Este enfoque aborda la ineficiencia computacional en Transformers cuando se trata de secuencias largas. El diseño de Mamba permite una inferencia más rápida y 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 la atención tradicional y los bloques MLP. Esta simplificación conduce a un modelo más ligero y rápido que escala linealmente con la longitud de la secuencia, una hazaña incomparable con sus predecesores.
Las características clave de Mamba incluyen:
- MSE selectivos: Estos permiten a Mamba filtrar información irrelevante y centrarse en datos relevantes, mejorando su manejo de secuencias. Esta selectividad es crucial para un razonamiento eficaz basado en contenidos.
- Algoritmo consciente del hardware: Mamba utiliza un algoritmo paralelo optimizado para hardware moderno, especialmente GPU. Este diseño permite un cálculo más rápido 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, incluidos el lenguaje, el audio y la genómica, sobresaliendo tanto en tareas previas al entrenamiento como en tareas específicas de dominio. Por ejemplo, en el modelado de lenguaje, Mamba iguala o supera el rendimiento de modelos Transformer más grandes.
El código de Mamba y los modelos previamente entrenados están disponibles abiertamente para uso comunitario en GitHub.
Las tareas de copia estándar son sencillas para modelos lineales. Los cabezales de inducción y copia selectiva requieren una memoria dinámica y consciente del contenido para los LLM.
Los modelos de espacio de estados estructurados (S4) han surgido recientemente como una clase prometedora de modelos de secuencia, que abarcan rasgos de RNN, CNN y modelos clásicos de espacio de estados. Los modelos S4 se inspiran en sistemas continuos, específicamente un tipo de sistema que mapea funciones o secuencias unidimensionales a través de un estado latente implícito. En el contexto del aprendizaje profundo, representan una innovación significativa, ya que proporcionan una nueva metodología para diseñar modelos de secuencia que sean eficientes y altamente adaptables.
La dinámica de los modelos S4
MUS (S4) Este es el modelo básico de espacio de estados estructurado. Se necesita una secuencia x y produce una salida y usando parámetros aprendidos A, B, Cy un parámetro de retardo Δ. La transformación implica discretizar los parámetros (convertir funciones continuas en discretas) y aplicar la operación SSM, que es invariante en el tiempo, lo que significa que no cambia en 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, permitiendo a los modelos S4 mantener una conexión con sistemas de tiempo continuo. Esto dota a los modelos de propiedades adicionales, como la invariancia de resolución, y garantiza una normalización adecuada, mejorando la estabilidad y el rendimiento del modelo. La discretización también establece paralelismos con los mecanismos de activación que se encuentran en los RNN, que son fundamentales para gestionar el flujo de información a través de la red.
Invariancia de tiempo lineal (LTI)
Una característica central de los modelos S4 es su invariancia temporal lineal. Esta propiedad implica que la dinámica del modelo permanece consistente a lo largo del tiempo, con los parámetros fijos para todos los pasos de tiempo. LTI es una piedra angular de la recurrencia y las convoluciones y ofrece un marco simplificado pero potente para crear modelos de secuencia.
Superar las limitaciones fundamentales
El marco S4 ha estado tradicionalmente limitado por su naturaleza LTI, lo que plantea desafíos en el modelado de datos que requieren dinámicas adaptativas. El reciente artículo de investigación presenta un enfoque que supera estas limitaciones mediante la introducción de parámetros que varían en el tiempo, eliminando así la restricción de LTI. Esto permite que los modelos S4 manejen un conjunto más diverso de secuencias y tareas, ampliando significativamente su aplicabilidad.
El término "modelo de espacio de estados" cubre ampliamente cualquier proceso recurrente que implique 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 SSM estructurados, se refieren a una clase específica de modelos que se han optimizado para un cálculo eficiente y al mismo tiempo conservan la capacidad de modelar secuencias complejas.
Los modelos S4 se pueden integrar en arquitecturas de redes neuronales de un extremo a otro, funcionando como transformaciones de secuencia independientes. Pueden verse como análogas a las capas de convolución en las CNN, proporcionando la columna vertebral 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 sostiene que un aspecto fundamental del modelado de secuencias es la compresión del contexto a un estado manejable. Los modelos que pueden centrarse selectivamente en las entradas o filtrarlas proporcionan un medio más eficaz para mantener este estado comprimido, lo que lleva a modelos de secuencia más eficientes y potentes. Esta selectividad es vital para que los modelos controlen de forma adaptativa cómo fluye la información a lo largo de la dimensión de secuencia, una capacidad esencial para manejar tareas complejas en el modelado del 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 antes era inalcanzable con modelos invariantes en el tiempo. Esto da como resultado SSM que varían en el tiempo y que ya no pueden usar convoluciones para un cálculo 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 entrada a los parámetros. B y Cy un parámetro de retardo Δ. Esto permite que el modelo se centre 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 forma variable en el tiempo mediante una operación de escaneo, que procesa elementos secuencialmente, ajustando el enfoque dinámicamente en 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 de audio, superando a los modelos anteriores en tareas complejas que requieren dependencias de largo alcance. Su versatilidad también se destaca en las evaluaciones de disparo cero en múltiples tareas, estableciendo un nuevo estándar para dichos 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, PyTorch 1.12+ y CUDA 11.6+. La instalación implica comandos pip simples para instalar los paquetes necesarios desde el repositorio de mamba. Si surgen problemas de compatibilidad con las versiones de PyTorch, puede ser útil utilizar el indicador –no-build-isolation con pip. Estos modelos, entrenados en conjuntos de datos extensos como Pile y SlimPajama, están diseñados para satisfacer diversas necesidades computacionales y puntos de referencia de rendimiento.
Mamba ofrece diferentes niveles de interfaces, desde la capa SSM selectiva hasta el bloque Mamba y estructuras completas de modelos de lenguaje. El bloque Mamba, que es el módulo principal de la arquitectura, utiliza una capa Conv1d causal y se puede integrar fácilmente en diseños de redes neuronales. El ejemplo de uso proporcionado en Python demuestra la creación de instancias de un modelo Mamba y el procesamiento de datos a través de él, destacando la simplicidad y flexibilidad del sistema.
Los modelos Mamba previamente entrenados están disponibles en Abrazando la cara, con tamaños que van desde 130 millones a 2.8 mil millones de parámetros, entrenados en el extenso conjunto de datos Pile y el conjunto de datos SlimPajama. Estos modelos están diseñados para cumplir con diversos requisitos computacionales y de rendimiento, adhiriéndose a los estándares dimensionales de GPT-3. Los usuarios pueden esperar un alto rendimiento y precisión de estos modelos, lo que convierte a Mamba en una opción competitiva para diversas aplicaciones, incluidas, entre otras, el modelado de lenguajes.
El impacto de Mamba
Mamba representa un paso adelante en el modelado de secuencias y ofrece una poderosa alternativa a las arquitecturas Transformer para procesar datos con gran densidad de información. Su diseño se alinea con las demandas del hardware moderno, optimizando tanto el uso de la memoria como las capacidades de procesamiento paralelo. La disponibilidad de código abierto del código base de Mamba y sus modelos previamente entrenados lo convierte en una herramienta accesible y sólida para investigadores y desarrolladores en el campo de la IA y el aprendizaje profundo.
