MLP Memory: A Retriever-Pretrained Memory for Large Language Models

TL;DR

MLP Memory is a lightweight, pretrained parametric module that internalizes retrieval patterns, improving language model accuracy and efficiency without explicit document access, outperforming traditional retrieval-augmented methods.

Contribution

We introduce MLP Memory, a novel pretrained parametric memory that imitates retrieval behavior, bridging the gap between retrieval-based and fine-tuning methods in language models.

Findings

Achieves 17.5 ext{ and }24.1 ext{ extbackslash }% scaling gains on datasets

Improves question-answering benchmarks by 12.3 ext{ extbackslash }%

Reduces hallucinations by up to 10 points on HaluEval

Abstract

Modern approaches to enhancing Large Language Models' factual accuracy and knowledge utilization face a fundamental trade-off: non-parametric retrieval-augmented generation (RAG) provides flexible access to external knowledge but suffers from high inference latency and shallow integration, while parametric fine-tuning methods like LoRA risk catastrophic forgetting and degraded general capabilities. In this work, we propose MLP Memory, a lightweight parametric module that learns to internalize retrieval patterns without explicit document access. By pretraining an MLP to imitate a $k$NN retriever's behavior on the entire pretraining dataset, we create a differentiable memory component that captures the benefits of retrieval-based knowledge access in a fully parametric form. Our architecture integrates this pretrained MLP Memory with Transformer decoders through simple probability interpolation, yielding 17.5\% and 24.1\% scaling gains on WikiText-103 and Web datasets, respectively. It further achieves 12.3\% relative improvement on five question-answering benchmarks and 5.2 points absolute gain across nine general NLP tasks, while reducing hallucinations by up to 10 points on HaluEval. Moreover, MLP Memory delivers 2.5$\times$ faster inference than RAG with superior accuracy. Our findings show that learning retrieval patterns parametrically bridges the gap between efficient inference and effective knowledge access, offering a practical alternative to both RAG and fine-tuning approaches.