Research map / Large language models / Retrieval and knowledge
Memory Decoder: A Pretrained, Plug-and-Play Memory for Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains a small 0.5B transformer decoder to imitate kNN retriever distributions over a domain corpus, then interpolates it with frozen LLMs of the same tokenizer for plug-and-play domain adaptation without retrieval at inference.
Problem
Domain adaptation of LLMs is costly with full DAPT (catastrophic forgetting, per-model training) while RAG and kNN-LM add heavy inference latency and storage overhead from datastores and nearest-neighbor search.
Approach
Pretrain a compact parametric transformer decoder with a KL distribution-alignment loss (plus LM loss) to match kNN retriever distributions built from a domain datastore; at inference it is simply interpolated with any same-tokenizer LLM's output distribution, with no retrieval.
Key results
- On WikiText-103, a 124M Memory Decoder reduces average perplexity from 18.34 (base) to 12.01, versus 12.20 for full DAPT, and applied to GPT2-medium it outperforms DAPT on GPT2-medium (12.25 vs 12.78 on GPT2-medium).
- A single 0.5B Memory Decoder lowers domain perplexity for Qwen2.5-72B from 5.85 to 3.46 average across Bio/Fin/Law, and for Qwen2-0.5B from 14.88 to 4.05.
- Inference overhead is 1.28x of the base model versus 1.51x for In-Context RAG and 2.17x for kNN-LM; average perplexity reduction of 6.17 points reported in abstract.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| augment GPT2-small with Memory Decoder | GPT2-small · WikiText-103 | perplexity: 24.89 -> 13.36 | Helped |
| augment GPT2-medium with Memory Decoder | GPT2-medium · WikiText-103 | perplexity: 18.29 -> 12.25 | Helped |
| augment GPT2-large with Memory Decoder | GPT2-large · WikiText-103 | perplexity: 15.80 -> 11.53 | Helped |
| augment GPT2-xl with Memory Decoder | GPT2-xl · WikiText-103 | perplexity: 14.39 -> 10.93 | Helped |
| augment GPT2-small with Memory Decoder | GPT2-small · Downstream tasks (average) | score: 67.45 -> 69.79 | Helped |
| augment Qwen2.5-7B with Memory Decoder | Qwen2.5-7B · Biomedicine | perplexity: 8.19 -> 3.57 | Helped |
| augment Qwen2.5-7B with Memory Decoder | Qwen2.5-7B · Finance | perplexity: 8.61 -> 3.42 | Helped |
| augment Qwen2.5-7B with Memory Decoder | Qwen2.5-7B · Law | perplexity: 5.94 -> 4.01 | Helped |
| transfer Memory Decoder to Llama3-8B | Llama3-8B · Biomedicine | perplexity: 7.95 -> 3.92 | Helped |
| transfer Memory Decoder to Llama3-8B | Llama3-8B · Finance | perplexity: 8.63 -> 4.32 | Helped |
| transfer Memory Decoder to Llama3-8B | Llama3-8B · Law | perplexity: 5.96 -> 4.46; LoRA: 4.12 | Mixed |
| increase Memory Decoder size from 124M to 345M (ablation) | Memory Decoder · WikiText-103 | average perplexity: 12.01 -> 11.26 | Helped |
Surprising findings
- DAPT causes severe drops on several downstream tasks (HYP, Yahoo), which the authors partly attribute to DCPMI evaluation artifacts.
- kNN distributions become extremely sparse for larger models (GPT-2-Large top-1 probability ~93.48%), and fine-tuned models yield sharper kNN distributions.
Limitations
- Pre-training requires costly kNN datastore searches to produce distillation targets, though this is one-time per domain.
- Cross-tokenizer transfer still requires some parameter updates (embedding and LM head re-initialized and trained), so it is not zero-shot.
- Downstream DCPMI results for DAPT may be affected by evaluation artifacts, making comparisons partly dependent on the evaluation method.
Experimental setup
Datasets: WikiText-103 MIMIC-III financial news Asylex SST-2 MR CR RT HYP CB RTE AGN
Benchmarks: WikiText-103 SST-2 MR CR RT HYP CB RTE AGN Yahoo MIMIC-III financial news
Metrics: perplexity accuracy latency overhead
Baselines: In-Context RAG kNN-LM DAPT LoRA
Builds on: kNN-LM Domain Adaptive Pretraining LoRA kNN-prompt
Closest papers by meaning
- Online Adaptation of Language Models with a Memory of Amortized Contexts (NeurIPS 2024 · similarity 0.89)
- Memory-Augmenting Decoder-Only Language Models through Encoders (Student Abstract) (AAAI 2024 · similarity 0.89)
- Peripheral Memory for LLMs: Integration of Sequential Memory Banks with Adaptive Querying (ICML 2025 · similarity 0.88)
- Reference Trustable Decoding: A Training-Free Augmentation Paradigm for Large Language Models (NeurIPS 2024 · similarity 0.88)
- Adaptive Draft-Verification for Efficient Large Language Model Decoding (AAAI 2025 · similarity 0.88)
- Generative Adapter: Contextualizing Language Models in Parameters with A Single Forward Pass (ICLR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.