LM2 : Modèles à Grande Mémoire

LM2: Large Memory Models

February 9, 2025

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

cs.AI

Résumé

Cet article présente le Modèle de Grande Mémoire (LM2), une architecture Transformer à décodeur uniquement, améliorée avec un module mémoire auxiliaire qui vise à résoudre les limitations des Transformers standard dans le raisonnement multi-étapes, l'argumentation relationnelle et la synthèse d'informations réparties sur de longs contextes. Le LM2 proposé intègre un module mémoire agissant comme un dépôt de représentations contextuelles, interagissant avec les jetons d'entrée via une attention croisée et se mettant à jour à travers des mécanismes de filtrage. Pour préserver les capacités générales des Transformers, le LM2 maintient le flux d'informations d'origine tout en intégrant un chemin mémoire complémentaire. Les résultats expérimentaux sur le banc d'essai BABILong démontrent que le modèle LM2 surpasse à la fois le modèle RMT augmenté de mémoire de 37,1% et le modèle de base Llama-3.2 de 86,3% en moyenne sur l'ensemble des tâches. Le LM2 présente des capacités exceptionnelles en inférence multi-sauts, raisonnement numérique et question-réponse sur de grands contextes. Sur l'ensemble de données MMLU, il obtient une amélioration de 5,0% par rapport à un modèle vanille pré-entraîné, démontrant que son module mémoire ne dégrade pas les performances sur des tâches générales. De plus, dans notre analyse, nous explorons l'interprétabilité de la mémoire, l'efficacité des modules mémoire et le comportement au moment des tests. Nos résultats soulignent l'importance de la mémoire explicite dans l'amélioration des architectures 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.