LM2: 대용량 메모리 모델

LM2: Large Memory Models

February 9, 2025

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

cs.AI

초록

본 논문은 대용량 메모리 모델(LM2)을 소개합니다. LM2는 어떤 표준 Transformer의 한계를 극복하기 위해 보조 메모리 모듈이 강화된 디코더 전용 Transformer 아키텍처로, 다단계 추론, 관계적 논증, 그리고 긴 맥락에 분산된 정보를 종합하는 데 목적을 두고 있습니다. 제안된 LM2는 입력 토큰과 상호 작용하며 게이팅 메커니즘을 통해 업데이트되는 문맥적 표현 저장소 역할을 하는 메모리 모듈을 통합하였습니다. Transformer의 일반적인 기능을 유지하면서 보충적인 메모리 경로를 통합함으로써 LM2는 원래의 정보 흐름을 유지합니다. BABILong 벤치마크에서의 실험 결과는 LM2 모델이 작업별로 평균적으로 기억 증강 RMT 모델보다 37.1% 및 기준선 Llama-3.2 모델보다 86.3% 우수한 성능을 보여준다는 것을 입증합니다. LM2는 다단계 추론, 숫자 추론, 그리고 대규모 맥락 질의응답에서 뛰어난 능력을 보여줍니다. MMLU 데이터셋에서는 사전 훈련된 바닐라 모델 대비 5.0% 향상을 달성하여, 일반적인 작업에서의 성능 저하가 없음을 입증합니다. 더불어, 우리의 분석에서는 메모리 해석 가능성, 메모리 모듈의 효과성, 그리고 테스트 시 동작을 탐구합니다. 우리의 결과는 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.