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

• URL: http://arxiv.org/abs/2508.01832v3
• Date: Thu, 23 Oct 2025 05:46:50 GMT
• Title: MLP Memory: A Retriever-Pretrained Memory for Large Language Models
• Authors: Rubin Wei, Jiaqi Cao, Jiarui Wang, Jushi Kai, Qipeng Guo, Bowen Zhou, Zhouhan Lin,
• Abstract summary: NLPRAG Memory is a lightweight parametric module that learns to internalize retrieval patterns without explicit document access.<n>Our architecture integrates this pretrained Memory with Transformer decoders through simple probability, yielding 17.5% and 24.1% scaling gains on WikiText-103 and Web datasets, respectively.
• Score: 39.890098399493986
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Should We Still Pretrain Encoders with Masked Language Modeling? [27.19054714197245]
  Recent evidence suggests that decoder models pretrained with Causal Language Modeling (CLM) can be effectively repurposed as encoders.<n>We train a total of 38 models ranging from 210 million to 1 billion parameters, and conduct over 15,000 fine-tuning and evaluation runs.<n>We find that while training with high-level CLM yields better performance across text representation tasks, CLM-trained models are more data-efficient and demonstrate improved fine-tuning stability.
  arXiv  Detail & Related papers  (2025-07-01T17:45:48Z)
• Maximally-Informative Retrieval for State Space Model Generation [59.954191072042526]
  We introduce Retrieval In-Context Optimization (RICO) to minimize model uncertainty for a particular query at test-time.<n>Unlike traditional retrieval-augmented generation (RAG), which relies on externals for document retrieval, our approach leverages direct feedback from the model.<n>We show that standard top-$k$ retrieval with model gradients can approximate our optimization procedure, and provide connections to the leave-one-out loss.
  arXiv  Detail & Related papers  (2025-06-13T18:08:54Z)
• DRAMA: Diverse Augmentation from Large Language Models to Smaller Dense Retrievers [86.54316283425001]
  Large language models (LLMs) have demonstrated strong effectiveness and robustness while fine-tuned as dense retrievers.<n>LLMs offer better efficiency, but they often fail to generalize effectively with limited supervised fine-tuning data.<n>We introduce DRAMA, a training framework that leverages LLMs to train smaller generalizable dense retrievers.
  arXiv  Detail & Related papers  (2025-02-25T18:59:07Z)
• Democratizing AI: Open-source Scalable LLM Training on GPU-based Supercomputers [65.35142508909892]
  We present a novel four-dimensional hybrid parallel algorithm implemented in a highly scalable, portable, open-source framework called AxoNN.<n>We demonstrate fine-tuning of a 405-billion parameter LLM using AxoNN on Frontier.
  arXiv  Detail & Related papers  (2025-02-12T06:05:52Z)
• Reinforcement Learning for Long-Horizon Interactive LLM Agents [56.9860859585028]
  Interactive digital agents (IDAs) leverage APIs of stateful digital environments to perform tasks in response to user requests.<n>We present a reinforcement learning (RL) approach that trains IDAs directly in their target environments.<n>We derive LOOP, a data- and memory-efficient variant of proximal policy optimization.
  arXiv  Detail & Related papers  (2025-02-03T18:35:42Z)
• ALoRE: Efficient Visual Adaptation via Aggregating Low Rank Experts [71.91042186338163]
  ALoRE is a novel PETL method that reuses the hypercomplex parameterized space constructed by Kronecker product to Aggregate Low Rank Experts.<n>Thanks to the artful design, ALoRE maintains negligible extra parameters and can be effortlessly merged into the frozen backbone.
  arXiv  Detail & Related papers  (2024-12-11T12:31:30Z)
• MOFHEI: Model Optimizing Framework for Fast and Efficient Homomorphically Encrypted Neural Network Inference [0.8388591755871735]
  Homomorphic Encryption (HE) enables us to perform machine learning tasks over encrypted data.<n>We propose MOFHEI, a framework that optimize the model to make HE-based neural network inference, fast and efficient.<n>Our framework achieves up to 98% pruning ratio on LeNet, eliminating up to 93% of the required HE operations for performing PI.
  arXiv  Detail & Related papers  (2024-12-10T22:44:54Z)
• MEMO: Fine-grained Tensor Management For Ultra-long Context LLM Training [24.066283519769968]
  Large Language Models (LLMs) have been trained using extended context lengths to foster more creative applications.<n>We propose MEMO, a novel framework for fine-grained activation memory management.<n>MeMO achieves an average of 1.97x and 1.80x MFU compared to Megatron-LM and DeepSpeed.
  arXiv  Detail & Related papers  (2024-07-16T18:59:49Z)
• SELF-GUIDE: Better Task-Specific Instruction Following via Self-Synthetic Finetuning [70.21358720599821]
  Large language models (LLMs) hold the promise of solving diverse tasks when provided with appropriate natural language prompts. We propose SELF-GUIDE, a multi-stage mechanism in which we synthesize task-specific input-output pairs from the student LLM. We report an absolute improvement of approximately 15% for classification tasks and 18% for generation tasks in the benchmark's metrics.
  arXiv  Detail & Related papers  (2024-07-16T04:41:58Z)
• SHERL: Synthesizing High Accuracy and Efficient Memory for Resource-Limited Transfer Learning [63.93193829913252]
  We propose an innovative METL strategy called SHERL for resource-limited scenarios. In the early route, intermediate outputs are consolidated via an anti-redundancy operation. In the late route, utilizing minimal late pre-trained layers could alleviate the peak demand on memory overhead.
  arXiv  Detail & Related papers  (2024-07-10T10:22:35Z)
• $\text{Memory}^3$: Language Modeling with Explicit Memory [22.572376536612015]
  We equip large language models (LLMs) with explicit memory, a memory format cheaper than model parameters and text retrieval-augmented generation (RAG) As a preliminary proof of concept, we train from scratch a 2.4B LLM, which achieves better performance than much larger LLMs and RAG models. We introduce a memory circuitry theory to support the externalization of knowledge, and present novel techniques including a memory sparsification mechanism that makes storage tractable.
  arXiv  Detail & Related papers  (2024-07-01T11:07:23Z)
• Applying RLAIF for Code Generation with API-usage in Lightweight LLMs [15.366324461797582]
  Reinforcement Learning from AI Feedback (RLAIF) has demonstrated significant potential across various domains. This paper introduces an RLAIF framework for improving the code generation abilities of lightweight (1B parameters) LLMs.
  arXiv  Detail & Related papers  (2024-06-28T17:16:03Z)
• DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models [79.01926242857613]
  Large language models (LLMs) are prone to hallucinations, generating content that deviates from facts seen during pretraining. We propose a simple decoding strategy for reducing hallucinations with pretrained LLMs. We find that this Decoding by Contrasting Layers (DoLa) approach is able to better surface factual knowledge and reduce the generation of incorrect facts.
  arXiv  Detail & Related papers  (2023-09-07T17:45:31Z)
• In-context Autoencoder for Context Compression in a Large Language Model [70.7621953091318]
  We propose the In-context Autoencoder (ICAE) to compress a long context into short compact memory slots. ICAE is first pretrained using both autoencoding and language modeling objectives on massive text data.
  arXiv  Detail & Related papers  (2023-07-13T17:59:21Z)
• LLM-Pruner: On the Structural Pruning of Large Language Models [65.02607075556742]
  Large language models (LLMs) have shown remarkable capabilities in language understanding and generation. We tackle the compression of LLMs within the bound of two constraints: being task-agnostic and minimizing the reliance on the original training dataset. Our method, named LLM-Pruner, adopts structural pruning that selectively removes non-critical coupled structures.
  arXiv  Detail & Related papers  (2023-05-19T12:10:53Z)
• An Efficient Memory-Augmented Transformer for Knowledge-Intensive NLP Tasks [40.81306982129298]
  Parametric and retrieval-augmented models have complementary strengths in terms of computational efficiency and predictive accuracy. We propose the Efficient Memory-Augmented Transformer (EMAT) It encodes external knowledge into a key-value memory and exploits the fast maximum inner product search for memory querying.
  arXiv  Detail & Related papers  (2022-10-30T08:34:49Z)
• Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks [133.93803565077337]
  retrieval-augmented generation models combine pre-trained parametric and non-parametric memory for language generation. We show that RAG models generate more specific, diverse and factual language than a state-of-the-art parametric-only seq2seq baseline.
  arXiv  Detail & Related papers  (2020-05-22T21:34:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.