LM2: Modelli con Ampia Memoria

LM2: Large Memory Models

February 9, 2025

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

cs.AI

Abstract

Questo articolo introduce il Modello di Grande Memoria (LM2), un'architettura Transformer solo decoder potenziata con un modulo di memoria ausiliario che mira ad affrontare le limitazioni dei Transformer standard nella ragionamento a più passaggi, nell'argomentazione relazionale e nella sintesi delle informazioni distribuite su contesti lunghi. Il LM2 proposto incorpora un modulo di memoria che funge da repository di rappresentazioni contestuali, interagendo con i token di input tramite attenzione incrociata e aggiornandosi attraverso meccanismi di gating. Per preservare le capacità generiche dei Transformer, LM2 mantiene il flusso di informazioni originale integrando un percorso di memoria complementare. I risultati sperimentali sul benchmark BABILong dimostrano che il modello LM2 supera sia il modello RMT potenziato dalla memoria del 37,1% che il modello di base Llama-3.2 del 86,3% in media su diverse attività. LM2 mostra capacità eccezionali nell'inferenza a più salti, nel ragionamento numerico e nel questionario con contesti ampi. Sul dataset MMLU, ottiene un miglioramento del 5,0% rispetto a un modello vaniglia preaddestrato, dimostrando che il suo modulo di memoria non degrada le prestazioni su compiti generici. Inoltre, nella nostra analisi, esploriamo l'interpretabilità della memoria, l'efficacia dei moduli di memoria e il comportamento al momento del test. Le nostre conclusioni sottolineano l'importanza della memoria esplicita nel potenziare le architetture 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.