MLP Memory: Language Modeling with Retriever-pretrained External Memory

• URL: http://arxiv.org/abs/2508.01832v1
• Date: Sun, 03 Aug 2025 16:40:53 GMT
• Title: MLP Memory: Language Modeling with Retriever-pretrained External Memory
• Authors: Rubin Wei, Jiaqi Cao, Jiarui Wang, Jushi Kai, Qipeng Guo, Bowen Zhou, Zhouhan Lin,
• Abstract summary: We propose to decouple from a decoder using a pretrained, differentiable external memory.<n>Our architecture exhibits strong perplexity and performance on downstream tasks.<n>We demonstrate superior performance on three hallucination benchmarks and nine memory-intensive tasks.
• Score: 26.033369983243624
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: While modern decoder-only LLMs achieve superior performance across various domains, hallucinations have risen to be a common problem in their generated text, hindering their application in knowledge-intensive tasks. Retriever-augmented generation (RAG) offers a solution, but the non-parametric nature of the retriever hinders its deep interaction with LLM. In this work, we propose to decouple memorization from the LLM decoder using a pretrained, differentiable external memory. The external memory is an MLP pretrained by imitating the behavior of a retriever on the entire pretraining dataset. Our resulting architecture, which comprises a transformer decoder and an external MLP memory pretrained on language modeling and retriever imitation respectively, demonstrates strong perplexity and performance on downstream tasks. Experiments show our architecture exhibits steeper power-law scaling with model size, achieving 17.5% and 24.1% improvement on WikiText-103 and Web datasets compared to decoder-only models while benefiting from added training without overfitting. We demonstrate superior performance on three hallucination benchmarks and nine memory-intensive tasks. Additionally, our approach delivers $80\times$ speedup over $k$NN-LM (500M tokens) and $1.3\times$ faster inference than decoder-only models. Unlike $k$NN-LM, which impairs reasoning, our MLP memory improves StrategyQA performance. We will open-source our code and models in the future.

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.