LM2: Modelos de Gran Memoria

LM2: Large Memory Models

February 9, 2025

Autores: Jikun Kang, Wenqi Wu, Filippos Christianos, Alex J. Chan, Fraser Greenlee, George Thomas, Marvin Purtorab, Andy Toulis

cs.AI

Resumen

Este documento presenta el Modelo de Gran Memoria (LM2), una arquitectura Transformer solo de decodificador mejorada con un módulo de memoria auxiliar que tiene como objetivo abordar las limitaciones de los Transformers estándar en el razonamiento de múltiples pasos, argumentación relacional y síntesis de información distribuida en contextos largos. El LM2 propuesto incorpora un módulo de memoria que actúa como un repositorio de representaciones contextuales, interactuando con tokens de entrada a través de atención cruzada y actualizándose mediante mecanismos de compuertas. Para preservar las capacidades generales de los Transformers, el LM2 mantiene el flujo de información original mientras integra una vía de memoria complementaria. Los resultados experimentales en el banco de pruebas BABILong demuestran que el modelo LM2 supera tanto al modelo RMT con aumento de memoria en un 37.1% como al modelo de línea base Llama-3.2 en un 86.3% en promedio en todas las tareas. El LM2 muestra capacidades excepcionales en inferencia de múltiples saltos, razonamiento numérico y respuesta a preguntas de contexto extenso. En el conjunto de datos MMLU, logra una mejora del 5.0% sobre un modelo vainilla pre-entrenado, demostrando que su módulo de memoria no degrada el rendimiento en tareas generales. Además, en nuestro análisis, exploramos la interpretabilidad de la memoria, la efectividad de los módulos de memoria y el comportamiento en tiempo de prueba. Nuestros hallazgos destacan la importancia de la memoria explícita en mejorar las arquitecturas Transformer.

English

This paper introduces the Large Memory Model (LM2), a decoder-only Transformer architecture enhanced with an auxiliary memory module that aims to address the limitations of standard Transformers in multi-step reasoning, relational argumentation, and synthesizing information distributed over long contexts. The proposed LM2 incorporates a memory module that acts as a contextual representation repository, interacting with input tokens via cross attention and updating through gating mechanisms. To preserve the Transformers general-purpose capabilities, LM2 maintains the original information flow while integrating a complementary memory pathway. Experimental results on the BABILong benchmark demonstrate that the LM2model outperforms both the memory-augmented RMT model by 37.1% and the baseline Llama-3.2 model by 86.3% on average across tasks. LM2 exhibits exceptional capabilities in multi-hop inference, numerical reasoning, and large-context question-answering. On the MMLU dataset, it achieves a 5.0% improvement over a pre-trained vanilla model, demonstrating that its memory module does not degrade performance on general tasks. Further, in our analysis, we explore the memory interpretability, effectiveness of memory modules, and test-time behavior. Our findings emphasize the importance of explicit memory in enhancing Transformer architectures.