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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Wei, Rubin, Cao, Jiaqi, Wang, Jiarui, Kai, Jushi, Guo, Qipeng, Zhou, Bowen, Lin, Zhouhan
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912930580135936
author Wei, Rubin
Cao, Jiaqi
Wang, Jiarui
Kai, Jushi
Guo, Qipeng
Zhou, Bowen
Lin, Zhouhan
author_facet Wei, Rubin
Cao, Jiaqi
Wang, Jiarui
Kai, Jushi
Guo, Qipeng
Zhou, Bowen
Lin, Zhouhan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2508_01832
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle MLP Memory: A Retriever-Pretrained Memory for Large Language Models
Wei, Rubin
Cao, Jiaqi
Wang, Jiarui
Kai, Jushi
Guo, Qipeng
Zhou, Bowen
Lin, Zhouhan
Computation and Language
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.
title MLP Memory: A Retriever-Pretrained Memory for Large Language Models
topic Computation and Language
url https://arxiv.org/abs/2508.01832