Renovación de las redes de memoria a largo plazo: XLSTM para la IA de próxima generación

Las LSTM, o redes de memoria a corto plazo, existen desde hace mucho tiempo. Se han aplicado para bastantes tareas relacionadas con secuencias, como la generación y traducción de texto o incluso la generación de títulos de imágenes.

Su inconveniente ha sido que no se podían paralelizar para utilizar las potentes GPU modernas. Esta limitación allanó el camino para la aparición de transformadores que aprovechan las GPU para una paralelización masiva del entrenamiento y la inferencia.

Si ahora intentamos renovar y paralelizar los LSTM, ¿pueden convertirse en herramientas para crear LLM de próxima generación?

Esta es exactamente la pregunta respondida por el artículo " XLSM - Extended Long Short-term Memory Networks", que significa memoria larga y corta " extendida ". Lo hacen proponiendo dos bloques novedosos en la arquitectura, a saber, sLSTM y mLSTM.

Entonces, profundicemos en los bloques sLSTM y mLSTM propuestos en este documento y veamos cómo podemos apilarlos para desarrollar la arquitectura XLSTM.

Explicación visual

Si usted es alguien como yo y desea que se le expliquen visualmente los XLSTM, consulte el video de YouTube en este artículo:

Actualización de LSTM

Una de las primeras redes diseñadas para abordar datos secuenciales es la red neuronal recurrente.

Red neuronal recurrente

Utiliza conexiones recurrentes en su arquitectura con x como entrada y o como salida. Si lo desplegamos, podemos visualizarlo como una secuencia de operaciones que ocurren en las marcas de tiempo, t-1, t y t+1 . Un inconveniente importante de los RNN fue el problema del gradiente que desaparece, donde el gradiente llega a cero cuando apilamos demasiados bloques.

Se propusieron LSTM, o redes de memoria a corto plazo, para superar los gradientes que desaparecen mediante la introducción de estados celulares y mecanismos de activación de la red.

Una ilustración simplificada de LSTM

Los estados celulares, c , son recuerdos a largo plazo que viven en varias marcas de tiempo. Los estados ocultos, h , son recuerdos a corto plazo que se transmiten de un paso en el tiempo a otro. Y, por supuesto, tenemos las entradas, z , de la secuencia de entrada.

Tres puertas tienen funciones en forma de S. La puerta del olvido utiliza una función sigmoidea para decidir qué información olvidar en la memoria a largo plazo. La puerta de entrada también utiliza una función sigmoidea para procesar la entrada y la agrega a la salida de la puerta de olvido. Esta operación de suma tiene un término elegante llamado carrusel de errores constantes en el artículo XLSTM y la literatura académica. Esta operación de suma es lo que aborda el problema de los gradientes de fuga que se encuentran en los RNN. Luego, la salida c_t es procesada por la puerta de salida , que generalmente es una función tanh que conduce a la salida de estado oculto h_t que se pasa al siguiente paso.

Con estas operaciones, hemos analizado las dos ecuaciones principales de LSTM que son las de c_t y h_t .

Inconveniente 1: Revisar las decisiones de almacenamiento

Uno de los principales inconvenientes de los LSTM es su incapacidad para revisar las decisiones de almacenamiento. Lo que esto significa es que a medida que aumenta la longitud de la secuencia, el modelo debería poder decidir si mantiene o no información pasada en la memoria.

Por ejemplo, si tomamos esta frase, “Tom fue a la tienda. Compró algunas bebidas” y compárelo con “Tom fue a la tienda a comprar algunos alimentos que incluían zanahorias, cebollas, plátanos, manzanas, naranjas, café y pan. También compró algunas bebidas”. Para cada palabra nueva, como plátanos o manzanas, el modelo tiene que revisar constantemente si debe mantener la palabra pasada "Tom" en su memoria. Este es un gran desafío para los LSTM y surge de la función sigmoidea de su puerta de olvido.

La función sigmoidea versus la función exponencial. El sigmoide se aplana hacia los extremos, pero el exponencial sigue aumentando.

Entonces, si tomamos una puerta de olvido, se compone de una función sigmoidea, que tiene una curva en forma de S que se aplana hacia el final. Esto indica que a medida que avanzamos hacia valores más altos de entrada, la decisión sobre qué olvidar y qué conservar en la memoria se vuelve bastante desafiante. Pero si usamos una función exponencial en su lugar, entonces el juego cambia y a medida que aumentamos los valores de la entrada, obtenemos un rango más amplio para las salidas. Esto, a su vez, indica que los LSTM pueden mejorar en la revisión de las decisiones de almacenamiento.

Solución 1: sLSTM

Entonces, la solución propuesta en este trabajo son los bloques del sistema. Si volvemos a la ecuación clásica de LSTM que representa el estado de la celda, como vimos antes, es función de la puerta de olvido y de las puertas de entrada.

Estas puertas, a su vez, están compuestas de funciones sigmoideas. Entonces, ¿qué pasa si reemplazamos estas funciones sigmoideas con funciones exponenciales? Las nuevas puertas f_t e i_t ahora se convierten en exp(f_t) y exp(i_t), y esa es prácticamente la modificación principal para crear el bloque sLSTM.

A diferencia de la función sigmoidea, que comprime las entradas para que estén en un rango fijo, las funciones exponenciales tienden a aumentar su valor a medida que aumenta la entrada, y no normaliza naturalmente la salida para que se encuentre entre, digamos, 0 y 1 como la función sigmoidea. función.

Por lo tanto, necesitamos introducir un nuevo estado de normalizador , que es una función de las puertas de entrada y olvido. Podemos considerarlo como un promedio móvil de un valor de normalización.

Usamos los valores de normalización calculados para normalizar la salida o el nuevo estado oculto.

Si bien la normalización se ocupa de los estados ocultos, para evitar que el exponencial explote las puertas de entrada y olvido, necesitamos introducir un estabilizador . Viene en forma de funciones logarítmicas para contrarrestar el efecto de las exponenciales e introducir estabilidad. Entonces, el estado del estabilizador es el máximo del registro de la puerta de olvido y la salida de la puerta de entrada. Restamos estos valores del estabilizador de la entrada y nos olvidamos de las puertas para estabilizarlos.

Inconveniente 2: memoria y paralelización

El segundo inconveniente de los LSTM es la falta de paralelización. Los LSTM fueron diseñados para manejar datos secuenciales, lo que significa que necesitan el resultado del procesamiento de la entrada anterior en la secuencia para procesar la entrada actual en la secuencia. Este inconveniente particular impide la paralelización y fue el culpable del inicio de la era Transformers.

La solución propuesta en este artículo son los nuevos bloques mLSTM. Entonces, veámoslos a continuación.

Solución: mLSTM

El siguiente bloque de construcción de XLSTM es el bloque mLSTM, donde m significa memoria. Volvamos nuevamente a la ecuación clásica de LSTM para ver cuál es su inconveniente. Podemos ver que el estado de la celda c_t es un escalar. Esto significa que solo tratamos con 1 número a la vez cuando tenemos el lujo de las GPU modernas con al menos 12 GB de memoria.

El bloque mLSTM introduce matrices en lugar de escalares para los estados de la celda. Volviendo a nuestra ecuación clásica de LSTM, ¿qué pasa si reemplazamos c_t con una matriz C *_t* para que el estado de la celda ahora se convierta en C *_t* mayúscula para indicar matrices y los estados de la celda se puedan recuperar no solo mediante una puerta i_t sino almacenando pares clave-valor que son vectores. cuyos valores se pueden recuperar mediante consultas que sean vectores de la misma dimensión.

Para que suene familiar la terminología del transformador, han introducido clave y valor aquí para formar esta matriz.

XLSTM

Con esa información sobre sLSTM y mLSTM, profundicemos en la arquitectura detallada de XLSTM.

sLSTM

Una mirada detallada al bloque sLSTM

Cuando se trata de sLSTM, utilizamos proyecciones posteriores. Entonces, la entrada pasa primero a través de capas de convolución causal con una función de activación rápida. La salida de estas capas luego se alimenta a través de una capa lineal de bloque diagonal con cuatro bloques diagonales o "cabezas". La salida de estos luego pasa a través del bloque sLSTM con cuatro cabezales. Finalmente, la salida se proyecta hacia arriba usando una capa MLP cerrada con activación GeLU y hacia abajo usando una función MLP cerrada.

mLSTM

Pasando a los detalles del bloque mLSTM, utilizamos proyecciones previas. Lo que significa que la entrada se proyecta primero hacia arriba con un factor de proyección de 2. Una de las salidas de proyección va al mLSTM y la otra va a la puerta de salida. La entrada al bloque mLSTM pasa por una convolución causal y luego a través de matrices de proyección diagonal de bloque de tamaño de bloque 4, que generan la consulta, la clave y el valor que el bloque mLSTM utiliza fácilmente.

Arquitectura XLSTM

Finalmente, podemos apilar los dos tipos de bloques para formar la arquitectura LSTM extendida. Entonces, los bloques de color gris oscuro son los bloques mLSTM y los de color gris claro son los bloques sLSTM.

En términos de ventajas, el artículo menciona que las redes XLSTM tienen una complejidad computacional lineal y una complejidad de memoria constante con respecto a la longitud de la secuencia.

Evaluaciones

Los autores se han entrenado en el conjunto de datos SlimPajama para compararlo con otros métodos basados en transformadores como LLAMA y métodos basados en el espacio de estados como MAMBA. Usaron esta notación de xLSTM a:b, donde a es la cantidad de bloques mLSTM y b es la cantidad de bloques sLSTM en la pila.

En términos de precisión, informan precisiones relativas escalando las precisiones entre 0 y 1, donde 0 es aleatorio y 1 es perfecto.

Las evaluaciones indican que XLSTM se desempeña mejor en tareas como Parity, mientras que Llama y Mamba se desempeñan mal.

De los resultados, algo que es de particular interés es la tarea de paridad, donde los transformadores o los modelos de espacio de estados tienden a tener problemas sin mezcla de memoria o seguimiento de estados. Podemos ver que en este tipo de tarea, xLSTM alcanza una precisión de 1 cuando usamos los bloques sLSTM y mLSTM juntos.

También han realizado algunos estudios de ablación para demostrar la solidez de los XLSTM. Son fáciles de entender en el documento. Además, este artículo trata más sobre las novedades arquitectónicas de XLSTM, por lo que no entraré aquí en los resultados experimentales.

Gritar

Si te gustó este artículo, ¿por qué no me sigues? Gorjeo ¿Dónde comparto actualizaciones de investigaciones de los mejores laboratorios de IA todos los días de la semana?

También suscríbete a mi Canal de Youtube donde explico conceptos y artículos de IA visualmente.

Conclusión

Espero que este artículo simplifique y facilite la comprensión de la arquitectura XLSTM, por qué la necesitamos y cómo podrían superar a los transformadores en un futuro próximo.

Esperemos y veamos qué tienen reservado. Te veré en mi próximo...