LM2: Grote Geheugenmodellen

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

Samenvatting

Dit artikel introduceert het Large Memory Model (LM2), een decoder-only Transformer-architectuur versterkt met een hulpgeheugenmodule die tot doel heeft de beperkingen van standaard Transformers aan te pakken bij multi-step redenering, relationele argumentatie en het synthetiseren van informatie die verspreid is over lange contexten. Het voorgestelde LM2 omvat een geheugenmodule die fungeert als een contextuele representatieopslagplaats, die interageert met invoertokens via kruislingse aandacht en wordt bijgewerkt via gating-mechanismen. Om de algemene mogelijkheden van de Transformers te behouden, handhaaft LM2 de oorspronkelijke informatiestroom terwijl het een aanvullend geheugenpad integreert. Experimentele resultaten op de BABILong benchmark tonen aan dat het LM2-model zowel het geheugenaugmented RMT-model met 37,1% als het basis Llama-3.2-model met 86,3% gemiddeld over taken overtreft. LM2 vertoont uitzonderlijke mogelijkheden in multi-hop redenering, numerieke redenering en vraag-antwoordtaken met grote context. Op de MMLU-dataset behaalt het een verbetering van 5,0% ten opzichte van een voorgeleerd standaardmodel, waaruit blijkt dat de geheugenmodule de prestaties op algemene taken niet vermindert. Verder onderzoeken we in onze analyse de interpreteerbaarheid van het geheugen, de effectiviteit van geheugenmodules en het gedrag op testtijd. Onze bevindingen benadrukken het belang van expliciet geheugen bij het verbeteren van Transformer-architecturen.

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.