给Transformer加个记忆模块，让模型更擅长长文本推理。
LM2: Large Memory Models
- 用交叉注意力和门控机制构建可更新的记忆体，与原模型并行运行。
- 在BABILong上比Llama-3.2高86.3%，多跳推理能力显著提升。
- 记忆模块不降低通用性能，适合需要长期推理的任务。
本文提出大型记忆模型（LM2），一种基于解码器的Transformer架构，通过引入辅助记忆模块来解决标准Transformer在多步推理、关系论证及长上下文信息整合方面的局限性。该记忆模块作为上下文表征仓库，通过交叉注意力与输入标记交互，并利用门控机制进行更新。为保持Transformer的通用性，LM2在保留原始信息流的同时，集成了一条互补的记忆路径。在BABILong基准测试中，LM2模型在各任务平均表现上较记忆增强型RMT模型高出37.1%，较基线Llama-3.2模型高出86.3%。其在多跳推理、数值推理和大上下文问答方面表现出色。在MMLU数据集上，相较预训练的纯模型提升5.0%，表明记忆模块未损害通用任务性能。进一步分析揭示了记忆可解释性、模块有效性及推理过程中的行为特征，强调显式记忆对提升Transformer架构的重要性。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。