PaperlayerView exact revision on arXiv

Source-linked AI summary

Memory Decoder: A Pretrained, Plug-and-Play Memory for Large Language Models

Jiaqi Cao, Jiarui Wang, Rubin Wei, Qipeng Guo, Kai Chen, Bowen Zhou, Zhouhan Lin

arXiv:2508.09874v2cs.CLcs.AI

TL;DR

Domain adaptation must balance the costs of DAPT against the inference overhead of RAG. Memory Decoder distills non-parametric retrieval behavior into a plug-and-play pretrained memory, improving specialized-domain performance across compatible models while preserving their parameters and general capabilities.

  • Problem

    DAPT requires costly full-parameter training and can cause catastrophic forgetting, while RAG adds inference overhead from nearest-neighbor search and longer contexts.

    Source: §1 · p. 2

  • Method

    Memory Decoder pretrains a small transformer decoder to imitate non-parametric retriever distributions and integrates it with compatible language models through plug-and-play inference.

    Source: §1 · p. 2, §3 · p. 4, §3.1 · p. 5

  • Results

    Across multiple model families and specialized domains, Memory Decoder consistently outperforms parametric adaptation methods and traditional retrieval-augmented approaches.

    Source: §1 · p. 3, §8 · p. 11

  • Takeaways & Limitations

    A single pretrained memory can efficiently transfer domain knowledge across models sharing a tokenizer while preserving base-model capabilities and avoiding retrieval overhead.

    Source: §8 · pp. 11–12

    Pre-training requires datastore searches for kNN supervision, and cross-tokenizer adaptation still requires parameter updates rather than truly zero-shot transfer.

    Source: §9 · p. 12

Abstract

from arXiv · show

Large Language Models (LLMs) have shown strong abilities in general language tasks, yet adapting them to specific domains remains a challenge. Current method like Domain Adaptive Pretraining (DAPT) requires costly full-parameter training and suffers from catastrophic forgetting. Meanwhile, Retrieval-Augmented Generation (RAG) introduces substantial inference latency due to expensive nearest-neighbor searches and longer context. This paper introduces Memory Decoder, a plug-and-play pretrained memory that enables efficient domain adaptation without changing the original model's parameters. Memory Decoder employs a small transformer decoder that learns to imitate the behavior of an external non-parametric retriever. Once trained, Memory Decoder can be seamlessly integrated with any pretrained language model that shares the same tokenizer, requiring no model-specific modifications. Experimental results demonstrate that Memory Decoder enables effective adaptation of various Qwen and Llama models to three distinct specialized domains: biomedicine, finance, and law, reducing perplexity by an average of 6.17 points. Overall, Memory Decoder introduces a novel paradigm centered on a specially pretrained memory component designed for domain-specific adaptation. This memory architecture can be integrated in a plug-and-play manner, consistently enhancing performance across multiple models within the target domain.

1. Introduction

Domain adaptation is important for specialized applications, but DAPT is costly and can cause forgetting while RAG adds inference overhead. Memory Decoder addresses this trade-off with a plug-and-play memory that preserves model parameters and improves performance across models and domains.

  • DAPT requires costly full-parameter training, whereas RAG incurs substantial inference computation from nearest-neighbor search and extended context.Source: §1 · p. 2
  • Memory Decoder is a small pretrained memory that imitates non-parametric retrievers without modifying the original language model parameters.Source: §1 · pp. 2–3
  • A single trained Memory Decoder integrates with compatible language models sharing the same tokenizer without model-specific adaptation or additional training.Source: §1 · pp. 2–3
  • Memory Decoder combines domain-specific performance improvements with minimal inference-latency impact, unlike retrieval-based methods.Source: §1 · pp. 2–3
  • Across biomedicine, finance, and law, a 0.5B Memory Decoder consistently enhanced seven Qwen2.5 models on finance text.Source: §1 · p. 3

2. Background

Domain adaptation is framed as improving next-token prediction on specialized corpora, while kNN-LM provides parameter-preserving adaptation through datastore retrieval. Its effectiveness is limited by substantial inference and storage overhead, motivating a compact parametric alternative.

2.1. Problem Formulation

  • Domain adaptation optimizes a pretrained model’s next-token distribution for a target-domain corpus.Source: §2.1 · p. 3

2.2. Nearest Neighbor Language Models

  • kNN-LM preserves pretrained parameters by retrieving domain-relevant neighbors from a key-value datastore.Source: §2.2 · p. 3
  • During inference, kNN-LM interpolates the pretrained model’s distribution with the distribution constructed from retrieved neighbors.Source: §2.2 · p. 3
  • kNN-LM introduces substantial inference and storage overhead; the Wikitext-103 datastore requires nearly 500GB even for GPT2-small.Source: §2.2 · p. 4

3. Memory Decoder

Memory Decoder distills non-parametric retrieval distributions into a compact decoder and combines its predictions with a base language model at inference. This avoids nearest-neighbor retrieval while retaining plug-and-play domain adaptation.

3. Memory Decoder

  • Memory Decoder has two components: retrieval-distribution alignment during pre-training and interpolation-based plug-and-play inference.Source: §3 · p. 4

3.1. Pre-training

  • Pre-training constructs context–kNN-distribution pairs by building a domain datastore, retrieving neighbors, excluding exact self-retrieval, and caching the resulting distributions.Source: §3.1 · p. 5
  • The training objective combines KL divergence to align Memory Decoder with cached kNN distributions and a language-modeling objective to limit deviation from corpus statistics.Source: §3.1 · p. 5
  • The authors report that JSD provides no advantage for the extremely sparse target distributions, while KL divergence with cross-entropy regularization performs better.Source: §3.1 · p. 5

3.2. Inference

  • At inference, the base model and Memory Decoder process the same context in parallel, and their distributions are interpolated with α controlling domain-knowledge influence.Source: §3.2 · p. 5
  • 1.28× overhead relative to the base model compares with 1.51× for In-Context RAG and 2.17× for kNN-LM.Source: §3.2 · p. 5, Figure 4 · p. 6

4. Experimental Setup

The experiments evaluate Memory Decoder across language modeling, downstream general-capability preservation, cross-model adaptation, and cross-vocabulary transfer. They use specialized corpora, established adaptation baselines, and perplexity or domain-conditional PMI metrics.

  • The evaluation covers WikiText-103 language modeling, nine downstream NLP tasks, cross-model adaptation, and cross-vocabulary adaptation.Source: §4 · p. 6
  • Domain-specific corpora cover biomedical clinical notes, financial news, and legal refugee-status documents.Source: §4 · p. 6
  • Baselines include In-Context RAG, kNN-LM, LoRA with comparable parameter counts, and full-parameter DAPT.Source: §4 · p. 6
  • Table 1 compares perplexity across GPT2 sizes and highlights that Memory Decoder(124M) applied to GPT2-medium outperforms DAPT of GPT2-medium.Source: Table 1 · p. 7
  • Language-modeling, cross-model, and cross-tokenizer experiments use sliding-window perplexity, while downstream tasks use domain-conditional PMI scoring.Source: §4 · p. 7

5. Results

Memory Decoder improves language modeling, downstream zero-shot performance, and domain adaptation across model sizes, architectures, and tokenizer families. Its strongest result is plug-and-play transfer from one memory module to multiple compatible models, with efficient cross-vocabulary adaptation.

5.1. Language Modeling on Wikitext-103

  • A single 124M Memory Decoder consistently improves the GPT2 family across model sizes, outperforming parameter-efficient methods without modifying base-model parameters.It remains competitive with DAPT on larger models and is superior to all adaptation methods on smaller models.Source: §5.1 · p. 7

5.2. Downstream Performance

  • Memory Decoder achieves the highest average score across nine zero-shot tasks while maintaining or improving performance where DAPT suffers catastrophic forgetting.It particularly strengthens textual entailment tasks such as CB and RTE while preserving general language capabilities.Source: §5.2 · p. 8, Table 2 · p. 8

5.3. Cross-Model Adaptation

  • A single 0.5B Memory Decoder enhances Qwen2 and Qwen2.5 models from 0.5B to 72B parameters across biomedical, financial, and legal domains.The gains include substantial perplexity reductions for smaller models and remain substantial for the largest models.Source: §5.3 · p. 8, Table 3 · p. 9

5.4. Cross-Vocabulary Adaptation

  • Qwen2.5-trained Memory Decoder transfers to Llama models after reinitializing only the embedding and language-model head with 10% of the original training budget.It achieves roughly 50% perplexity reductions for Llama3-8B on biomedical and financial domains and generally outperforms LoRA outside legal text.Source: §5.4 · p. 8, Table 4 · p. 9
  • Cross-vocabulary transfer extends Memory Decoder beyond a single tokenizer family, enabling domain knowledge learned on one architecture to support another with minimal additional training.This expands the method’s practical use across diverse model ecosystems.Source: §5.4 · pp. 8–9

6. Analysis

The analysis shows that Memory Decoder combines retrieval-like memorization with parametric language coherence, while its size and hybrid training objective materially affect performance.

6.2. Impact of Memory Decoder Size

  • Larger Memory Decoders consistently improve performance across GPT2 models, with the large variant achieving the best average perplexity.Decoder size can therefore be selected according to computational constraints while preserving the base model’s capabilities.Source: §6.2 · p. 10, Table 6 · p. 10

6.3. Ablation on the pre-training objective

  • Memory Decoder outperforms DAPT-model logit interpolation by an average of 1.90 perplexity points across GPT2 scales.The gains persist at larger scales, supporting the complementary value of the hybrid pre-training objective.Source: §6.3 · pp. 10–11, Table 7 · p. 11

7. Related Work

Prior work uses retrieval, domain pre-training, and parameter-efficient adaptation to specialize language models, but these approaches impose inference, optimization, or model-specific constraints. Memory Decoder is positioned as a differentiable, cross-model memory module that avoids expensive retrieval and broad parameter modification.

Retrieval-Augmented Generation

  • Token-level retrieval supports rare and out-of-domain patterns but incurs substantial inference computation, whereas Memory Decoder provides full-dataset knowledge access without expensive retrieval operations or model-specific datastores.It combines differentiable optimization with access to domain knowledge beyond local contexts.Source: §7 · p. 11

Domain Adaptation

  • DAPT, LoRA, and adapters either require costly domain training or model-specific modifications that limit generalization across architectures.Memory Decoder instead enhances multiple frozen models and supports cross-model and cross-tokenizer adaptation.Source: §7 · p. 11

8. Conclusion

Memory Decoder uses a small pretrained decoder to emulate non-parametric retrieval while adapting language models without changing their parameters. Across model families and domains, it combines retrieval-like performance benefits with preserved general capabilities and modular deployment.

8. Conclusion

  • Memory Decoder pre-trains a small transformer decoder to emulate non-parametric retrievers and adapt compatible language models without modifying their parameters.The approach is reported to outperform both parametric adaptation methods and traditional retrieval-augmented approaches across multiple model families and specialized domains.Source: §8 · p. 11
  • A single pretrained Memory Decoder can enhance models sharing a tokenizer and transfer to different tokenizers and architectures with minimal additional training.This decouples domain expertise from model architecture and reduces resources needed for specialized model development.Source: §8 · pp. 11–12
  • Memory Decoder preserves retrieval-augmented performance benefits while maintaining base-model general capabilities and avoiding catastrophic forgetting often observed with fine-tuning.The resulting framework is presented as modular and efficient for specialized language-model enhancement.Source: §8 · p. 12

9. Limitations

The approach has two acknowledged limitations: one-time kNN-search overhead during pretraining and incomplete zero-shot transfer across tokenizers.

  • Pretraining requires searching key-value datastores for kNN distributions, creating a one-time computational bottleneck that is amortized across adapted models.The overhead occurs during the training pipeline rather than deployment.Source: §9 · p. 12
  • Cross-tokenizer adaptation requires some parameter updates to align embedding spaces, preventing truly zero-shot cross-architecture transfer.Source: §9 · p. 12

A. Interpolation hyperparameter 𝛼of all tasks

Interpolation strength varies by task: larger GPT models generally need less memory augmentation, while downstream tasks center around lower α values and domain-specific tasks require validation tuning.

A.1. Language Modeling on Wikitext-103

  • The general α pattern centers around 0.6 for WikiText-103, with smaller values used for larger GPT models.The trend is attributed to stronger base models requiring less augmentation from memory.Source: §A.1 · p. 17, Table 8 · p. 17

A.2. Downstream Performance

A.3. Cross-Model and Cross-Vocabulary Adaptation

  • Domain-specific language-modeling tasks tune α on validation data by searching over {0.4, 0.6, 0.8, 0.9}.Source: §A.3 · p. 17

B. Analysis of DAPT Performance on Downstream Tasks

The analysis links DAPT evaluation behavior and retrieval-target properties to model scale, domain adaptation, and loss design, identifying KL divergence with cross-entropy regularization as the effective alignment strategy.

B. Analysis of DAPT Performance on Downstream Tasks

C.1. Extreme Sparsity and Concentration

  • k-NN distributions typically assign non-zero probability to only 2–3 tokens, unlike language-model distributions with smooth mass and long tails.Source: §C.1 · p. 18, Figure 5 · p. 19

C.2. Scale-Dependent Behavior

C.3. Domain Adaptation Effects

  • Fine-tuned models produce sharper k-NN distributions, and Memory Decoders trained on them consistently achieve lower perplexity.Source: §C.3 · p. 19

D.1. Failed Approaches

  • Alternative losses underperformed KL divergence: focal loss failed to concentrate distributions, JSD offered no advantage, BiLD underperformed, and explicit sparsity penalties caused instability.Source: §D.1 · p. 19, §D.1.1 · p. 20, §D.1.2 · p. 20

D.2. Why KL Divergence Succeeds

  • Forward KL encourages sparse outputs by penalizing probability mass outside the target support while favoring a few dominant modes.Source: §D.2 · p. 21
  • Cross-entropy regularization anchors predictions to linguistically valid outputs while KL drives sparsity, balancing concentration with semantic coherence.Source: §D.2 · p. 21
Loading 2508.09874v2…