Meta ha publicado recientemente un innovador resultado de investigación: un nuevo tipo de tecnología de capa de memoria que mejora significativamente la precisión fáctica de los grandes modelos de lenguaje (LLM) y permite una expansión sin precedentes en la escala de parámetros. Esta tecnología no solo desafía los métodos tradicionales de expansión de las redes neuronales, sino que también ofrece nuevas direcciones para el diseño de arquitecturas de IA futuras.
El núcleo de esta investigación radica en el uso de un mecanismo de búsqueda de clave-valor entrenable para agregar parámetros adicionales al modelo sin aumentar la carga computacional (FLOPs). La idea central es complementar las capas de alimentación directa (feedforward) de alta densidad computacional con una capa de memoria de activación dispersa, proporcionando así una capacidad especializada para almacenar y recuperar información.
En comparación con las redes densas tradicionales, la capa de memoria es más eficiente en el procesamiento del almacenamiento de información. Por ejemplo, para que un modelo de lenguaje aprenda información simple como nombres y fechas de nacimiento, o capitales de países, la capa de memoria puede hacerlo mediante un simple mecanismo de búsqueda de clave-valor, lo cual es más eficiente que usar una red de alimentación directa.
La principal contribución de esta investigación es la ampliación de la capa de memoria a una escala sin precedentes, alcanzando los 128 mil millones de parámetros. Los resultados experimentales muestran que, en tareas posteriores, los modelos de lenguaje equipados con la capa de memoria mejorada no solo superan a los modelos densos con el doble de carga computacional, sino que también superan a los modelos de expertos mixtos con la misma carga computacional y cantidad de parámetros. La mejora del rendimiento es especialmente notable en tareas fácticas.
Los investigadores de Meta lograron esto reemplazando una o más redes de alimentación directa (FFN) en la red Transformer con la capa de memoria. Este reemplazo mostró una ventaja consistente en diferentes tamaños de modelos base (de 134 millones a 8 mil millones de parámetros) y capacidades de memoria (hasta 128 mil millones de parámetros). Los resultados experimentales muestran que la capa de memoria puede mejorar la precisión fáctica de los modelos de lenguaje en más del 100%, con mejoras significativas también en la codificación y el conocimiento general. En muchos casos, los modelos con capa de memoria alcanzan el rendimiento de modelos densos que requieren cuatro veces más cómputo.
Los investigadores también realizaron varias mejoras en la capa de memoria para superar los desafíos de su aplicación a gran escala:
Mecanismo de búsqueda de clave multiplicativa: Para resolver el cuello de botella de la recuperación de claves de consulta en capas de memoria a gran escala, se utilizaron claves de cuantificación multiplicativa entrenables, evitando así la comparación por pares de cada clave de consulta.
Capa de memoria paralelizada: Para lograr la paralelización de la capa de memoria en entornos multi-GPU, las operaciones de búsqueda y agregación de incrustaciones se distribuyeron en múltiples GPU.
Mecanismo de memoria compartida: Para maximizar el uso compartido de parámetros, se utilizó un grupo de parámetros de memoria compartido entre todas las capas de memoria.
Optimización del rendimiento y la estabilidad: Se utilizaron núcleos CUDA personalizados para optimizar la operación EmbeddingBag, mejorando significativamente la utilización del ancho de banda de memoria. Además, se introdujo un mecanismo de puerta controlado por la entrada con no linealidad silu para mejorar el rendimiento y la estabilidad del entrenamiento.
Los resultados experimentales también revelaron los siguientes hallazgos clave:
El tamaño de la capa de memoria tiene un impacto significativo en el rendimiento: a medida que aumenta el tamaño de la capa de memoria, el rendimiento de la respuesta a preguntas fácticas mejora continuamente.
Varias capas de memoria son mejores que una sola: el uso de varias capas de memoria con parámetros compartidos puede mejorar el rendimiento, pero demasiadas capas de memoria pueden reducirlo. El número óptimo de capas de memoria es tres.
La capa de memoria aprende hechos más rápido: al principio del entrenamiento, los modelos equipados con la capa de memoria muestran una mejora de rendimiento más rápida, lo que indica que la capa de memoria ayuda al modelo a aprender hechos más rápidamente.
La capa de memoria y las capas densas son complementarias: los experimentos muestran que tanto la capa de memoria dispersa como la capa de alimentación directa densa son esenciales.
Para verificar la efectividad de la tecnología de capa de memoria, los investigadores realizaron evaluaciones en varios puntos de referencia, incluyendo:
Respuesta a preguntas fácticas (NaturalQuestions, TriviaQA)
Respuesta a preguntas de múltiples saltos (HotpotQA)
Conocimiento científico y de sentido común (MMLU, HellaSwag, OBQA, PIQA)
Escritura de código (HumanEval, MBPP)
Los resultados muestran que los modelos equipados con la capa de memoria superan a los modelos de referencia en estas pruebas, especialmente en la respuesta a preguntas fácticas, donde la mejora del rendimiento es más significativa.
Esta investigación de Meta no solo proporciona nuevas ideas para la expansión de los modelos de IA, sino que también abre nuevas vías para resolver problemas fácticos y mejorar el rendimiento del modelo. Los investigadores creen que la tecnología de capa de memoria tiene una gran escalabilidad y que en el futuro se podrá aplicar ampliamente en diversas aplicaciones de IA. También señalan que, aunque la capa de memoria aún enfrenta desafíos en la aceleración de hardware, confían en que mediante la investigación y optimización continuas, su rendimiento podrá igualar o incluso superar al de las redes de alimentación directa tradicionales.
Además, el equipo de investigación de Meta espera mejorar aún más el rendimiento de la capa de memoria mediante nuevos métodos de aprendizaje, reducir el olvido y las alucinaciones del modelo, y lograr el aprendizaje continuo.
La publicación de esta investigación sin duda ha inyectado nueva vitalidad al campo de la IA, y nos llena de expectativas sobre el futuro desarrollo de la IA.
Artículo: https://arxiv.org/pdf/2412.09764
