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

Аннотация

Эта статья представляет модель Large Memory (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.