Memory Decoder: A Pretrained, Plug-and-Play Memory for Large Language Models
Authors: Jiaqi Cao
Jiarui Wang
Rubin Wei
Qipeng Guo
Kai Chen
Bowen Zhou
Zhouhan Lin
Venue: NeurIPS 2025 poster ()
Keywords: large language models
domain adaptation
retrieval-augmented generation
Primary area: deep_learning

Abstract: 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 \textit{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.

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 1:2t
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University
Shanghai AI Laboratory TTsinghua University
SJTU Paris Elite Institute of Technology
maximus cao@outlook. com, 2in.hhouhann@mmii.coom
code model available at http:p//gghabbooeooooooa/aaooaae
Abstract
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 train- ing and suffers from catastrophic forgetting. Meanwhile, Retrieval-Augmentt Generation (RAG) introduces substantial inference latency due expensive nearest- neighbor searches and longer context. This paper introduces Memory Decoder, a plug-and and-play pretrained memory that enables efficient domain adaptation without changing the original model's 's parameters. Memory Decoder employs a small trans- former decoder that learns imitate the behavior of an external non-paramen parameerii retriever. Once trained, Memory Decoder can be seamlessly integrated with any pre- trained 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-specinf adaptation. This memory architecture can be integrated a gllg-and-pil manner, consistently enhancing performance across multiple models within the target domain.
Introduction
Large Language Models (LLMs) have demonstrated remarkable capabilities across a wide range of natural language processing tasks (Grattafiori et al. 2024 Yang et al. 2024 Liu al. 2024 Guo et al. (025). Pretrained on vast corpora of general text data, LLMs have revolutionized how we approach language understanding and generation tasks. However, despite their impressive general capabilities, adapting LLMs to perform optimally in specific domains remains significant challenge. Domain-specific adaptation is crucial for applications in specialized fields such biomedicine, finance, and law (Chen et al.) 2023 Liu et al. 20235; Colombo et al., 2024), where domain expertise and terminology are essential for accurate and reliable performance.
Domain adaptation for pretrained language models has traditionally followed several approaches, each with distinct advantages and limitations. Domain Adaptive Pre- Pre-Training (DAPT) in- volves continued pre-training of the LLM on domain-specific corpora (Gururangan et al. 2020).
"Equal contribution.
Corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

DAPT
RAG
MemDec (Ours)
Finance LLM
Qmen2-0n.58
Owenz-0L.5
Owenz-A.51
Pre-train
Finance
Finance
MemDee
Omen2-7H
Owenz-78 7B
(weca-778
Retrieve
Interpolate
Pre-train
Owen2-7288
Owen2-7218
Qwen2-7211
Pre-train
Figure I: Comparison of domain adaptation approaches. DAPT (left) requires separate pre-training for each model size, modifying original parameters. RAG (middle) maintains model parameters but requires expensive retrieval from external datastores during inference. Memory Decoder (right) offers plug-and-play solution where single pretrained memory component be interpolated with models of different sizes, avoiding both parameter modification and retrieval overhead.
While effective, this approach suffers from substantial com- putational costs associated with full-parameter training, es- pecially as model sizes continue to grow into billions of pa- rameters. Furthermore, adapting multiple models to the same domain requires separate training runs for each model, lead- ing to resource inefficiency. Even with successful DAPT im- plementation, these models often encounter catastrophic for- getting, where the adaptation process diminishes the model's s general capabilities (Kirkpatrick et al. 2017 Ven van de et al. 2024).
0.30
Qwen2.5
LoRA
14
Reriesuts
73
120
Parameter Count (8)
Retrieval-Augeee Augmenteddd Generation (RAG) offers an alternative approach by enhancing model outputs with relevant retrieved information (Lewis et al., 2020 Izacard et al. (003)). While this method preserves the original model parameters, intro- duces substantial computation overhead during inference due to expensive nearest neighbor (kNN) searches across large datastores and extended context (He et al., 2021).
Figure 2: Perplexity comparison of Qwen2.5 models augmented by Mem- ory Decoder and LoRA adapter of the same param count on the finance do- main.
These two approaches present fundamental dilemma domain adaptation: DAPT requires costly training procedures and cannot efficiently adapt multiple models to the same domain, while RAG introduces significant computation and storage overhead during inference. This inherent trade-off between the plug-and-pll nature of RAG and the inference efficiency of DAPT highlights the research gap for solution offering both adaptability across models and computational efficiency during deployment. To address this challenge, we propose Memory Decoder (MemDec), a plug and- play pretrained memory designed for efficient domain adaptation of large language models without modifying their parameters. Our approach draws inspiration from retrieval-based methods like kNN- LM (Khandelwal al. 2019), but overcomes their limitations through different paradigm. Rather than building and searching model-specifie datastores during inference, Memory Decoder employs small transformer decoder that is specially pretrained to imitate the behavior of non- non-parametric retrievers by aligning its output distribution with the ones of non-parametric retrievers. Figure illustrates how our approach differs fundamentally from both DAPT and RAG.
The key innovation of our approach lies in its plug-and-play and-play functionality: once trained, a single Memory Decoder can be seamlessly integrated with any large language model that shares the same tokenizer, without requiring model-specific adaptations or additional training. This architectural design enables immediate deployment across diverse model architectures, significantly reducing the computational resources needed for domain adaptation pre-training. Furthermore, unlike RAG methods, Memory Decoder achieves domain-specific performance improvements with minimal impact inference latency, combining versatility with computational efficiency.

---

Experimental results across three specialized domains (biomedicine, finance, and law) and multiple model architectures demonstrate the versatility of Memory Decoder. As shown Figure Z the same Memory Decoder with only 0.5B parameters consistently enhances performance across seven different models from the Qwen2.5 model family on the finance domain. Our comprehensive analysis confirms that Memory Decoder successfully preserves the advantages of hon-parametric approaches while eliminating their computational overhead, establishing new paradigm for efficient domain adaptation of LLMs.
Our contributions can be summarized as follows:
We introduce Memory Decoder, a plug-and-play pretrained memory that enables efficient domain adaptation for large language models without modifying their original parameters.
We present the first approach that replaces traditional non con-parametric retricvers with compact parametric model, achieving superior performance while eliminating costly retrieval operations during inference.
We demonstrate Memory Decoder's generalizability, where single domain domain-specific pre- trained memory can be seamlessly integrated across all models with the same tokenizer.
Background
2.1 Problem Formulation
Domain adaptation aims to enhance pretrained language model's performance on specialized text. Formally, given a pretrained model MMMм with parameters 0 and a domain corpus D domainn the goal is to optimize the next-token prediction distribution peLM ar; 0) for the target domain. Here, x == (T1, $2, 11) represents the context sequence and y/s denotes the target token.
2.2 Nearest Neighbor Language Models
The k-nearest neighbor language model (KNN-LM) (Khandelwal et al. 2019) enables non- non- con-parametric domain adaptation without modifying the pretrained model's parameters.
For a domain corpus, kNN-LM first constructs a key-value datastore:
(K, V) {(0(x;), v/) (z4,, y4) Doomain
(I)
where o(-) extracts hidden representations from the pretrained model.
During inference, for context æ, computes k = <(x), retrieves k-nearest neighbors, and constructs probability distribution:
PINN(M (r) x
Σ Lu exp(-d(ki, kt)/T)
(2)
The final prediction interpolates between the pretrained model and kNN distributions:
PANN-PLM((6(17) = À ( + (1 - X) PPLM():
While effective, kNN-LM introduces substantial computational and storage overhead during infer- ence. For instance, the Wikitext- 103 datastore requires nearly 500GB storage for GPT2-sma model (He et al. (02TT. These limitations motivate our Memory Decoder, a compact parametric model pretrained mimic retrieval behavior while eliminating the need for large datastores.
3 Memory Decoder
In this section, we present Memory Decoder (MemDec), glug-and-play pretrained memory designed for efficient domain adaptation of large language models. Our method consists of two primary com- ponents: a specialized pre-trate pre-training procedure that aligns the output distribution of Memory Decoder with those of non- non-parann parameeriii retrievers (Section 3.1 and an efficient inference mechanism that enables plug-and-pla, domain adaptation (Section 3.2 As illustrated Figure 3 Memory Decoder first learns to mimic non-parametric retrieval distributions during pre-training (upper part), then seamlessly integrates with any compatible language model during inference (lower part), eliminating the computational overhead associated with datastore maintenance and nearest neighbor search.

---

Memory Decoder Pretraining
Nun-parametric Retriever
Misk
Besus
Cook
Pre-train Corpus
Distribution Aligument
Zinia slock rooe after statement from
Memory Decoder
Mask
Bezon
Cook
Memory Decoder Inference
Mask
Cook
Question
Musk
Memory Decoder
Bizes
Musk
Cook
Bezna
Cook
During pre-training, Memory Decoder learns align its output distributions with those generated by non-parametric retrievers through distribution alignment loss. Lower 3.2) During inference, Memory Decoder processes input in parallel with the base LLM, and their distributions are interpolated to produce domain-enhanced predictions without retrieval overhead.
Figure 3: Overview of Memory Decoder architecture. Uppers 3.1
3.1 Pre-training
Our primary goal during pre-training is to enable Memory Decoder MMem produce probability distributions that closely resemble those generated by non-parametric retrievers when encountering the same context. This approach effectively encodes the domain knowledge captured in large key-value datastores into the parameters of our compact model.
Data Construction Since we require non-parametric distributions as supervision signals, we construct training pairs of (I;, PaNN (=a:) in advance to enable efficient pre-training. Here, represents the input context and PANN (x:) denotes the probability distribution generated by the non- parametric retriever for that context. First, build key-value datastore (K, V) {(((z)) v4) (x,, (4.,3) ] using our domain-specific corpus, where <(.) extracts hidden representations from specific layer of the pretrained model. For each context x; in the corpus, we then perform k-nearest neighbor search against this datastore to identify similar contexts. To avoid trivial self-retrieval that would contaminate the learning signal,
exclude the top-1 neighbor where its key exactly matches the query key. Finally, we compute the non non-parametric distribution PaNN(-12; (-12,) for each context using the retrieved neighbors and cache these context-distribution pairs for training.
Pre-training Objective Unlike traditional language modeling with single-label targets, kNN distributions offer richer supervision signals by capturing the diversity of plausible continuations in the domain (Xu et al. (see Appendix for detailed analysis kNN distributions). Through extensive experimentation, we have identified that a hybrid objective yields optimal performance.
Our approach centers Distribution Alignment Loss that minimizes the KI. divergence (Van Erven, Harremos 2014) between Memory Decoder's 's output distribution and the cached kNN distributions for each sample:
(x) KL PMcm (-1x:)
(4)
To prevent excessive deviation from the underlying corpus distribution, integrate complementary standard Language Modeling objective (Zhang, Sabuncu 2018):
(zM(2.) log PM (v:
(5)

---

2.17x
1.51x
Lafency 10)
1.28x
1.0x
3-
KNN-LM
In-Context RAG
Base LM
150k
175k
Figure 4: Inference latency comparison across domain adaptation methods. These measurements were conducted on Qwen2.5-1.5B (Yang et al. 2024 for biomedicine domain text, augmented by 0.5B Memory Decoder.
The final loss function balances these two objectives through hyperparameter 8: [(x;) = 3 (x.) + (1 3) (z;)
(6)
Previous failed attempts to learn kNN distributions and conjecture why vanilla KL divergence with cross-entropy regularization succeeds are detailed in Appendix D
3.2 Inference
Once pretrained, Memory Decoder exhibits key Blug-and-play capability that allows it to adapt any language model with a compatible tokenizer to the target domain via simple interpolation. During inference, both the pretrained language model MPLM and Memory Decoder MMem process the same input context in parallel, and their output distributions are interpolated:
PMem-PL (ye
(yelzr) (zz)
PMem (yr (x) (1 a:) PPLM
where a E [0, 1] controls the influence of domain-specific knowledge.
approaches that introduce substantial latency from nearest neighbor search and extended context processing, Memory Decoder requires only a single forward pass through relatively small transformer decoder. As demonstrated in Figure 4 method achieves significant improvements in inference efficiency compared to alternative domain adaptation techniques. With just 1.28x overhead relative the base model, Memory Decoder substantially outperforms both In-Context RAG (Ram et al. 2023) (1.51x) and KNN-LM (Khandelwal et al. 2019) (2. 17x..x). This computational advantage, combined with Memory Decoder's model-agnoste design, makes our approach particularly valuable for production environments where both performance and efficiency are critical considerations.
Unlike traditional
Experimental Setup
Overview We evaluate Memory Decoder across four complementary settings: (1) Language
modeling on WikiText- 103 (55.1) to demonstrate effectiveness across GPT-2 model scales; (2)
Downstream tasks (55. to verify preservation of general capabilities during domain adaptation; (3)
Cross-model adaptation (5.3) showing single Memory Decoder enhancing Qwen models from 0.5B
to 72B parameters; (4) Cross-vocabulary adaptation ($5.4) demonstrating efficient transfer between
tokenizer families; These experiments establish Memory Decoder as a versatile, plug-and-play solution for efficient domain adaptation across diverse architectures and applications.
Datasets For language modeling experiments, we use Wikitext- 103 (Merity et al. 2016), a standard benchmark containing over 100M tokens from Wikipedia articles. For downstream evaluation,
5

---

GPT2-small
GPT2-med
GPT2-large
GPT2-xl
24.89
18.29
15.80
14.39
Non-parumetric methods
+ +In-Contens RAG
18.46
14.01
12.09
11.21
+KNN-LM
15.62
12.95
12.21
11.30
Parametric methods
+DAPT 14.76
12.78
11.10
10.16
+LoRA
18.63
13.88
11.77
10.67
+MemDec 13.36
12.25
11.53
10.93
Table Perplexity comparison of domain adaptation methods across GPT2 model sizes on Wikitext- 103. The best performing results are highlighted in bold, while the second- best results are underlined. Notably, applying our Memory Decoder( (12MM) to GGP2-medium(3455)
outperforms DAPT of GPT2-medium(345M demonstrating the effectiveness of our approach in capturing domain knowl- edge without modifying original parameters.
following the kNN-prompt framework, we assess performance across nine NLP tasks: sentiment
analysis (SST2 (Socher et al.] 2013), MR (Pang, Lee 2005b), CR (Hu, Liu, 2004). RT (Pang, Lee 2005ah), textual entailment (HYP P (Kiesel et al. 2219)) CB (De Marneffe et al. | 2019), RTE (Dagan al. 2010)), and text classification (AGN (Zhang al., (015a), Yahoo (Zhang et al. 2015b)). For domain-specific adaptation, we utilize three specialized corpora: (1) biomedical text from MIMIC- III (Johnson et al. 2 clinical notes covering over 46,000 patients, (2) financial news (Liu et al. 2023a) from April 2024 to October 2024, and (3) legal text from the Asylex corpus (Barale et al. containing 59, documents of refugee status determination in Canada from 1996 2022.
2023
Baselines We compare Memory Decoder against several established domain adaptation methods: In-Context RAG (Ram et al. 2023), which implements BM25 retriever that processes 32 query tokens, with retrieval occurring every tokens. kNN-LM (Khandelwal et al. 2019), configured with interpolation parameter > = 0.25 and temperature settings of r = 1 for GPT-2 small and medium, and 7 13 for large and xl models. LoRA (Hu et al. 2022), applied MLP layers, with rank adjusted for each model to achieve parameter counts comparable to Memory Decoder. Domain Adaptive Pretraining(DAP)) (Gururangan et al. 2020), which involves complete retraining of all model parameters on the domain-specifte corpus.
query, key, value and
Training Details We conduct our experiments on an 8xA800 80GB GPU setup. For language modeling and downstream evaluations, we use GPT2-xl modee(fee on wikitext) to build the key-value datastore and non non-maramttric distributions for training, and continue training on GPT2-small mfnel(finetueed on wikitext) with learning rate 1e-3. For cross-model adaptation, we use Qwen2.5-L.5B (Yang et al. 2024) to build the datastore, and continue training on Qwen2.5-0.5B with learning rate lc-4. For cross- coss-vocabulary adaptation, we use Llama3.2-1D 1B (Grattafiori et al. 2024) to build the datastore, and continue training on the Memory Decoder trained from cross-model experiments, with its embedding layer and language model head re-initialized. All experiments use training budget equivalent to the computational cost of training a 7B parameter model for epoch, with DAPT and LoRA baselines using the same maximum training FLOPS but early stopped to prevent overfitting. The training hyperparameter is set to 0.5 across all tasks.
Evaluation Metrics For language modeling, cross cross-model, and cross-tokenizer experiments, we use sliding window perplexity. Following Baevski, Auli (2018), in each test example, the context length set to 1024 where only the latter 512 tokens are scored. For downstream evaluation, following methodology from Shi et al. (2022), we report results using the domain-conditional PMI scoring rule (Holtzman et a0.21 The interpolation hyperparameter a is tuned on the validation split of each task following Khandelwal et al. (2019), see more details in Appendix A

---

SST2
MR
CR
RT
HYP
CB
RTE
AGN
Yahoo
Avg
base
81.98
78.40
84,40
76.54
63.75
41.07
52.70
78.79
49,40
67.45
Non-parametric methods
+kNN-LM
81.98
77.95
83.80
77.95
64.14
39.28
$2.70
77.73
49.63
67.24
Parametric methods
+DAPT
83.52
80.15
80.45
77.39
36.04
50.00
51.26
64.31
24.40
60.84
+LoRA
80.88
76.90
83.95
76.07
64.14
39.28
53.79
81.06
49.46
67.28
+MemDec
82.43
78.35
84.35
77.30
64.15
57.14
55.24
79.80
49.37
69.79
Table 2: Performance on nine diverse NLP tasks including sentiment analysis, textual entailment, and text classification.
Results
5.1 Language Modeling Wikitext-100
Table demonstrates the exceptional effectiveness of Memory Decoder across all GPT2 model sizes. A single Memory Decoder with only 124M parameters consistently enhances the entire GPT2 family, showcasing its plug-and-play capability regardless of base model size. For smaller models, our approach delivers superior results compared all adaptation methods-notabbl maintaining an advantage for GPT2-medium despite utilizing only one third of the parameters. Even when applied larger models where DAPT has inherent advantages due to full model updates, Memory Decoder remains highly competitive while consistently outperforming all other parameter-efficient methods without modifying any original parameters. These results validate that a small parametric decoder can effectively capture the benefits of non- con-paramttric retrieval while eliminating computational overhead.
5.2
Downstream Performance
Table 2 reveals Memory Decoder's ability enhance domain adaptation while preserving general language capabilities in zero-shot evaluation settings. Unlike DAPT, which suffers catastrophic forgetting on several tasks (particularly HYP and Yahoo where performance drops by nearly half; Appendix B for detailed analysis), Memory Decoder maintains or improves performance across all evaluated tasks. Our approach achieves the highest average score across all nine tasks, outperforming the base model, kNN-LM, and LoRA while demonstrating particular strength on textual entailment tasks like CB and RTE. These results validate key advantage of our plug-and-play architecture: by keeping the original model parameters intact while augmenting them with domain knowledge, Memory Decoder achieves domain adaptation without sacrificing general capabilities. Importantly, all experiments are conducted in zero-shot setting, and our method should be viewed as orthogonal to in-context learning approaches.
5.3 Cross-Model Adaptation
Table demonstrates Memory Decoder's exceptional plug-and-plas capabilities across diverse model sizes and architectures. single Memory Decoder (0.5B parameters) consistently enhances performance across all models in both the Qwen2 and Qwen2.5 families, spanning from 0.5B to 72B parameters. For smaller models like Qwen2-0.5B, our approach achieves dramatic perplexity reductions- deductions-ttrassformmngg transforming Nomain-speciin and financial text. Even for the largest models in the family, Memory Decoder provides substantial improvements, demonstrating that larievee-- knowledge remains valuable regardless of model scale. These results validate Memory Decoder's 's core strength: a single pretrained memory component enhance multiple models sharing the same tokenizer, providing efficient domain adaptation that scales from the smallest to the largest models while consistently outperforming existing approaches.
performance to state-of-the-ar results on both biomedical
5.4 Cross-Vocabulary Adaptation
Table @emonstraate Memory Decoder's ability to generalize across different tokenizers and model architectures. By re-initializing only the embedding layer and language model head of our Qwen2.5- trained Memory Decoder, we successfully adapt it to the Llama model family with just 10% of the original training budget. This efficient transfer enables substantial performance improvements across
7

---

Model
Río
Fin
Law
Avg
Owenn2
Family
Qwen2-0.5B
18.41
16.00
10.23
14.88
+LoltA
Avg
+MemDec
3,75
3,84
4.57
4.05
Llama3 Family
Owen2-L5B
Llama3-8B
+LoRA
5.73
7.37
4.84
5.98
+LoRt
4.38
5.68
4.12
4.73
+MemDDe
3,61
4.32
Qwen2-7B
8,36
8,31
5.92
7.53
Llamas-70B
5.92
6.87
4.90
5.90
+LoRA
4.47
5.64
4.02
4.71
*MoonDee
3.39
3.38
4.00
3.66
Llamat
Famnly
Qwen2-72B
Liama3.1-88
+MemDDe
3
3.20
3.69
3.45
+LoRA
4.38
5.72
4,10
4.73
Owen2.5
Family
Owen2.5-0,5B
17.01
16.04
9.86
14.30
1lama3.1-708
5.85
6,68
4.89
5.81
+LoRA
7.02
5.75
7.55
*MenDDe
3.74
3.87
4.57
4.06
Llamat
Fammly
Qwen2.5-1.5B
11.20
7.42
Llama322-18
+LoRA
5.59
7.50
4.82
5.97
+LoRA
5.97
7.83
5.21
6.34
+MemDer
3,67
3.61
4.29
3.86
MeemDee
4.06
4.85
5.11
4.67
Qwen2.331B
9.70
9.83
6,68
8.74
Llama3-2-3B
9.83
9.70
6.83
8.79
+LoRA
5.07
6.71
4.45
5.41
+oRR
5.11
6.55
4.59
5.42
+MemDee
3,63
3,52
4.16
3.77
+MemDee
3,99
4.45
4.76
4,40
QwenZ.5.7B
K.19
K.61
5.94
7.58
+LoIA
4.03
5.31
3,81
4.38
Table 4: Cross-vocabulary adaptation results
+MemDee
3,57
3,42
4.01
3.67
demonstrating efficient knowledge transfer be- tween model families. Memory Decoder trained
Qwen2.5.14B
7.01
7.60
5.35
6.65
+MemDer
351
3.31
3,86
3,56
Qwen2.--3B
6.65
7:38
3.18
6.40
on Qwen2.5 can be adapted to Llama models with minimal additional training (10% of original budget), achieving substantial perplexity reduc- tions across all Llama variants and consistently outperforming LoRA in biomedical and financial domains.
+MemDer
3.48
3.29
3.81
3,53
Qwen2.3-72
3.90
6.80
4.84
3.85
+MemDee
3.44
3.23
3.70
3.46
Table 3: Cross-model adaptation results across three specialized domains. A single 0.5B Mem- ory Decoder enhances models ranging from 0.5B to 72B parameters.
all Llama variants. For Llama3-8B, Memory Decoder achieves roughly 50% perplexity reduction on both biomedical and financial domains. Similar improvements extend to the Llama3.1 and Llama3.2 families, with our method consistently outperforming LoRA on biomedical and financial domains, though showing room for improvement legal text. These findings illustrate Memory Decoder's versatility beyond single tokenizer family, demonstrating that domain knowledge learned from one architecture can be efficiently transferred to another with minimal additional training. This capability expands the practical utility of our approach, offering a streamlined path to domain adaptation across diverse model ecosystems.
6
Analysis
6.1
Case Study:
Bridging Parametric and Non-Parametrie Methods
Memory Decoder fundamentally learns to compress the knowledge stored in large non-parametric datastores into a compact parametric model, combining the memorization capabilities retrieval methods with the efficiency and generalization of parametric approaches. To validate this hypothesis, we conducted case studies on WikiText- 103 examining how different methods assign probabilities to specific tokens.
As shown in Table 5 Memory Decoder exhibits two crucial capabilities:
Long-tail Knowledge: For factual information like "Jacobi" and "1 "1906". Memory Decoder assigns dramatically higher probabilities than the base model (68.9) 94% vs. 0.% and 98.65% vs. 1.57%), successfully capturing the memorization benefits of non-parametric methods.
Semantic Coherence: For function words and logical continuations like "on" and "C", Memory Decoder maintains probabilities closer to the base model rather than following kNN-LM's lower probabilities, demonstrating its ability to preserve coherent language modeling capabilities that pure retrieval methods sacrifice.
These observations confirm that Memory Decoder occupies unique position: it enhances memoriza- tion of domain-specific and long-tail knowledge like non non-parametric methods, while maintaining the generalization and reasoning capabilities inherent to parametric models.

---

Long-tail Knowledge Learning
Context (target token underlined)
MemDec
kNN
Base L.M
starred alongside actors Mark Strong and Derek
68.94%
9.39%
0.12%
Jacobi
The launch of HMS Dreadnought in 1906 by the
98.65 %
40.62%
1.57%
Royal Navy raised the stakes
Semantic Coherence and Reasoning
Context (target token underlined) In 2000 Boulter had guest-starring role on the tele- vision series The Bill
MemDec
kNN
Base LM
40.11%
8.07%
45.51 %
...three tank squadrons for special overseas operations, known as ', 'B' and C Special Service Squadrons
50.10%
10.76%
63.04% %
Table 5: Probability assignments for specific tokens by different methods. Orange section Memory Decoder excels at capturing long-tail factual knowledge, assigning dramatically higher probabilities than the base model. Cyan section For semantic coherence, Memory Decoder intelligently balances between KNN-LM and base model probabilities, preserving linguistic fluency.
GPT2-small
GPT2-medium
GPT2-large
GPT2-xl
Avg
Base
24.89
18.29
15.80
14.39
18.34
DAPT
14.76
12.78
11.10
10.16
12.20
MemDec-small (1aal((1144)) (124M)
13.36
12.25
11.53
10.93
12.01
MemDec- (345M)
12.08
11.59
10.92
10.43
11.26
MemDec-large (774M)
11.67
11.23
10.83
10.28
11.00
Table 6: Performance comparison with different Memory Decoder sizes. Even small Memory Decoders achieve competitive performance with full-parameter DAPT while maintaining plug-and- play capability.
6.2 Impact of Memory Decoder Size
Table examines how Memory Decoder size affects performance across the GPT2 family. As Memory Decoder size increases, performance consistently improves across all base models, with the large variant achieving the best average perplexity. These results validate that Memory Decoder provides an efficient alternative to full model fine-tuning: practitioners can choose the decoder size based on their computational constraints while maintaining the crucial advantage of preserving the original model's capabilities.
6.3 Ablation on the pre-training objective
We compare Memory Decoder against logit interpolation with DAPT model. Table 7 shows Memory Decoder consistently outperforms DAPT interpolation across all GPT2 scales, with an average gain of 1.90 perplexity points. Notably, the gains persist even at larger scales, confirming that our hybrid pre-training objective provides complementary value beyond what standard language modeling objectives can achieve, regardless of model size.
Related Work
Retrieval-Augmentee Generation Retrieval-Augmented Generation (RAG) enhances language models by incorporating knowledge from external sources, with retrieval granularity ranging from documents (Chen al.] 2017) to passages (Guu al., 2020 Lewis et al.) [2020] Izacard al. 2023) to tokens (Khandelwal et al. 2019 He et al.] 2021] Min et al. 2022] Yogatama et al. (2021) Teven- retrieval achieves superior performance for rare patterns and out-of-domain scenarios but introduces substantial computation overhead during inference. While non-differentiable retrieval mechanisms prevent end-to-end optimization and memory token approaches (Chevalier et al. (023) enable differentiable access but are limited to local contexts, Memory Decoder provides both differentiable

---

Base Model
Baseline PPL
-DAPT-mmal
MemDec-small
GPT2-small 24.89
13.95
13.36 (-2.59)
GPT2-medium
18.29
14.26
12.25 (-2.01)
GPT2-large
15,80
13.13
11.53 (-1.60)
GPT2-xl
14.39
12.30
10,93 (-1.37)
Average
18.34
13.91
12.01 (-1.0)
Table 7: Memory Decoder vs. DAPT model interpolation WikiText- 103, Both use 124M parameter models with different training objectives.
optimization and full-dataset knowledge access without expensive retrieval operations model- specific datastores.
Domain Adaptation
Domain adaptation techniques have evolved from domain-specific pre-training
(SciBERT Beltagy et al., 2019). BioBERT (Lee et al. (020), ClinicalBERT (Huang et al.) 2019))
to parameter-effciemn methods like LoRA (Hu et al. (022) and adapters (Wang et al. 2020 Diao
ct al. 2021 2023). However, these approaches require model-specific modifications, preventing generalization across architectures. Memory Decoder addresses this limitation by providing domain-specific memory module that enhances multiple frozen language models without parameter modifications, enabling cross-model adaptation within tokenizer families and efficient cross-tokenizer transfer with minimal additional training.
Conclusion
In this paper, introduced Memory Decoder, novel plug-and-play approach for domain adaptation of large language models. By pre-training a small transformer decoder to emulate the behavior of non-parametric retrievers, Memory Decoder effectively adapts any compatible language model to specific domain without modifying its parameters. Our comprehensive experiments across multiple model families and specialized domains demonstrate that Memory Decoder consistently outperforms approaches.
both parametric adaptation methods and traditional retrieval-augmente
The key innovation of Memory Decoder lies in its versatility and efficiency. A single pretrained Memory Decoder can seamlessly enhance any model that shares the same tokenizer, and with minimal additional training, can be adapted to models with different tokenizers and architectures. This capability enables efficient domain adaptation across model families, dramatically reducing the resources typically required for specialized model development. Our results confirm that Memory Decoder preserves the performance benefits of retrieval-augmentee methods while maintaining the general capabilities the base models, avoiding the catastrophic forgetting often observed with fine-tuning approaches.
paradigm for domain adaptation that fundamentally reimagines how we specialize language models for particular domains. By decoupling domain expertise from model architecture through a pretrained memory component, our approach creates more modular, efficient, and accessible framework for enhancing language model performance in specialized fields.
Memory Decoder introduces
Limitations
While Memory Decoder demonstrates significant advantages for domain adaptation, we acknowledge several limitations in our current approach. The pre- treiiiiinng phase requires searching in key-value datastores to obtain kNN distributions as training signals, introducing computational overhead during the training process. Although this cost is incurred only once per domain and is amortized across all adapted models, remains bottleneck in the pipeline. Additionally, while cross-tokenizer adaptation requires minimal training compared to training from scratch, it still necessitates some parameter updates to align embedding spaces, preventing truly zero-shot cross-ares archttcrcee transfer.
10

---

Acknowledgement
This work is sponsored by the National Key Research and Development Program of China (No. 20232D0121402), Shanghai Fundamental Research Program for General AI Models (No. 20238H2DZ2025110011, and the National Natural Science Foundation of China (NSFC) grant (No.62576211).
References
Baevski Alexei, Auli Michael. Adaptive imput representations for neural language modeling // arxiv preprint arXiv:1809.1088: 2018.
Barale Claire, Rovatsos Michuel, Bhuta Nehal. Automated refugee case analysis: An nlp pipeline for supporting legal practitioners arXiv preprint arXiv:2305. 155333 2023.
Beltagy Lo Kyle, Cohan Arman. SciBERT: pretrained language model for scientific arXiv preprint
arXiv:1903. 10676. 2019.
Chen Dunqi, Fisch Adam, Weston Jason, Bordes Antoine. Reading wikipedia answer open-domain questions arxiv preprint 104.0000000001111... 2017.
Chen Yirong, Wang Zhenyu, Xing Xiaofen, Zhipei, Fang Kai, Wang Junhong, Sihang, Wu Jieling, Liu Qi, Xu Xiangmin, others. Bianque: Balancing the questioning and suggestion ability of health Ilms with multi-turn health conversations polished by chatgpt // arXiv preprint arXiv:2310. 2023.
Cheng Daixuan, Huang Shaohan, Wei Furu. Adapting large language models via reading comprehension // The Twelfth International Conference on Learning Representations. 2023.
Chevalier Alexis, Wettig Alexander, Ajith Anirudh, Chen Danqi Adapting language models to compress contexts /| arXiv preprint arXiv:2305. 14788. 2023.
Colombo Pierre, Pires Telmo Pessoa, Boudiaf Malik, Culver Dominic, Melo Rui, Corro Caio, Martins Andre FT, Esposito Fabrizio, Raposo Vera Lúcia, Morgado Sofia, others Saullm-7b: A pioneering large language model for law // arXiv preprint Xii::2403.008833 2024.
Dagan Ido, Dolan Bill, Magnini Bernardo, Roth Dan. Recognizing textual entailment: Rational, evaluation and approaches-crratum // Natural Language Engineering. 2010. 16, 105-105.
De Marneffe Matri-Caaterrne,, Simons Mandy, Tonhauser Judith. The commitmentbank: Investigating projection naturally occurring discourse // proceedings of Sinn und Bedeutung. 23, 2. 2019. 107-124.
Diao Shizhe, Xu Rujjia, Su Hongjin, Jiang Yilei, Song Yan, Zhang Tong. Taming pre-train trainanied language models with n-gram representations low-resource domain adaptation Proceedings the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers). 2021. 3336-3349.
Diao Shizhe, Xu Tianyang, Xu Ruijia, Wang Jiawei, Zhang Tong. Mixture-of-lomain---ddeeer Decoupling and injecting domain knowledge to pre-trained language models memories // arXiv preprint arXiv:2306.05406. 2023.
Grattafiori Aaron, Dubey Abhimanyu, Jauhri Abhinav, Pandey Abhinav, Kadian Abhishek, Al-Dahle Ahmad, Letman Aiesha, Mathur Akhil, Schelten Alan, Vaughan Alex, others. The llama 3 herd of models // arXiv preprint arXiv:2407.21170 2024.
Guo Daya Yang Dejian, Zhang Haowei, Song Junxiao, Zhang Ruoyu, Xu Rumxin, Thu Qihao, Ma Shirong, Wang Peiyi, Bi Xiao, others. . Deepseek-rl: Incentivizing reasoning capability in 11ms via reinforcement learning arXiv preprint acXiv:2501. 2025.
Gururangan Suchin, Marasović Ana, Swayamdipta Swabha, Kyle, Beltagy Iz, Downey Doug, Smith Noah Don't stop pretraining: Adapt language models to domains and tasks /| arxiv preprint arXiv:2004. 10964. 2020.
Gum Kelvin, Lee Kenton, Ting Zora, Pasupat Panupong, Chang Mingwei. Retrieval augmented language model pre-training // International conference on machine learning. 2020. 3929-3933
He Junxian, Neubig Graham, Berg-Kiie Kirrkkrrrkkkkkk Taylor. Efficient nearest neighbor language models // arXiv preprint arXiv:2109.04D. (0.01099.222.... 2021.

---

Holtzman Ari, West Peter, Shwartz, Vered, Choi Yejin, Zettlemoyer Luke, Surface form competition: Why the highest probability answer isn't always right II arXiv preprint arXiv:2104.08315. 2021.
Hu Edward J, Shen Yelong, Wallis Phillip, Allen-Zhu Zeyuan, Li Yuanzhi, Wang Shean, Wang Lu, Chen Weizhu, others. Lora: Low-rank adaptation of large language models. // ICLR. 2022. 1, 3.
Hu Minging, Liu Bing. Mining and summarizing customer reviews // Proceedings of the Tenth ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. New York, NY, USA: Association for Computing Machinery, 2004. 168 177. (KDD '04).
Huang Kexin, Altosaar Juan, Ranganath Rajesh. Clinicalbert: Modeling clinical notes and predicting hospital readmission arXiv preprint arXiv: 1904.05342. 2019.
Izacard Gautier, Lewis Patrick, Lomeli Maria, Hosseini Lueas, Petroni Fabio, Schick Timo, Dwivedi- Yu Jane, Joulin Armand, Riedel Sebastian, Grave Edouard. Atlas: Few-shot learning with retrieval augmented language models // Journal of Machine Learning Research. 2023. 24, 251. 1-43.
Johnson Alistair EW, Pollard Tom J, Shen Lu, Lehman Li-wei Feng Mengling, Ghassemi Mohanmad, Moody Benjamin, Szolovits Peter, Anthony Celi Leo, Mark Roger G. MIMIC-III, freely accessible critical care database // Scientific data. 2016. 3, 1. 1-9.
Khandelwal Urvashi, Levy Omer, Jurafsky Dan, Zettlemoyer Luke, Lewis Mike. Generalization through memorization: Nearest neighbor language models // arXiv preprint arXiv:1911.00111 2019
Kiesel Johannes, Mestre Maria, Shukla Rishahh, Vincent Emmamuel, Adineh Payam, Corney David, Stein Benno, Potthast Martin. SemFval-2019 Task 4: Hyperpartisan News Detection // Praceedings of the 13th International Workshop on Semantic Evaluation. Minneapolis, Minnesota, USA: Association for Computational Linguisties. VI 2019. 829-839.
Kirkpatrick James, Pasconu Razvan, Rabinowitz Neil, Veness Joel, Desjardins Guillaume, Rusu Andrei A, Milan Kieran, Quan John, Ramalho Tiago, Grabska- Barwinska Agnieszka, others Overcoming catastrophic forgetting in neural networks // Proceedings the national academy of sciences. 2017. 114, 13. 3521 -3526.
Lee Jinhyuk, Yoon Wonjin, Kim Sungdong, Kim Donghyeon, Kim Sunkyu, So Chan Ho, Kang Jaewoo. BioBERT: pre-trained biomedical language representation model for biomedical text mining // Bioinformaties. 2020. 36,4 4. 1234-1240.
Lewis Patrick, Perez Ethan, Piktus Aleksandra, Petroni Fabio, Karpukhin Vladimir, Goyal Naman, Kiittler Heinrich, Lewis Mike, Yih Wen-tau, Rocktäschel Tim, others Retrieval-auggeente generation for knowledge- intensive nlp tasks I/ Advances in neural information processing systems. 2020. 33. 9459-9474.
Liu Aixin, Feng Bei, Xue Bing, Wang Bingxuan, Wu Bochao, Lu Chengda, Zhao Chenggang, Deng Chengqi, Thang Chenyu, Ruan Chong, others. Deepseek-v3 technical report // arXiv preprint arXiv:2412. 194477. 2024.
Liu Xiao- Yang. Wang Guoxuan, Yang Hongyang, Zha Daochen. Data-centric FinGPT: Democratizing Internet- scale Data for Financial Large Language Models // NeurIPS Workshop on Instruction Tuning and Instruction Following. 2023a.
Liu Xiao-Yang Yang, Wang Guoxuan, Yang Hongyang, Zha Daochen. Fingpt: Democratizing internet-scale data for financial large language models Il arXiv preprint sxXiv:2307.10483 2023b.
Merity Stephen, Xiong Caiming, Bradbury James, Socher Richard. Pointer Sentinel Mixture Models. 2016.
Min Sewon, Shi Weijia Lewis Mike, Chen Xilun, Yih Wen-tau, Hajishirzi Hannaneh, Zettlemoyer Luke. Nonpara- metric masked language modeling // arXiv preprint arXiv:2212.01339 2022.
Pang Bo, Lee Lillian. Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales // Proceedings of the ACL. 2005a.
Pang Bo, Lee Lillian. Seeing stars: exploiting class relationships for sentiment categorization with respect to rating scales // Proceedings of the 43rd Annual Meeting on Association for Computational Linguistics. USA: Association for Computational Linguistics, 2005b. 115 124. (ACL '05).
Ram Ori, Levine Yoav, Dalmedigos Ity Muhlgay Dor, Shashua Amnon, Leyton-Brac Brown Kevin, Shoham Yoav, In- context retrieval-augnennnn augmented language models // Transactions of the Association for Computational Linguistics. 2023. 11. 1316-1331.
Shi Weijia, Michael Julian, Gururangan Suchin, Zenteeaaer Luke. Nearest neighbor zero-shot inference // Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, 2022.
12

---

Socher Richurd, Perelygin Alex, Wu Jean, Chuang Jason, Monning Christopher D., Andrew, Potts Christopher. Recursive Deep Models for Semantic Compositionality Over Sentiment Treebank Proceedings of the 2013 Conference on Empirical Methods Natural Language Processing. Seattle, Washington, USA: Association for Computational Linguistics, X 2013. 1631-1642.
Van Erven Tim, Harremos Peter. Rényi divergence and Kullback-Leibee divergence // IEEE Transactions on Information Theory. 2014. 60, 3797-3820.
Ven Gido M van Soures Nicholas, Kudithipudi Dhireesha. Continual learning and catastrophic forgetting arXiv preprint arXiv:2403.05 05175.222 2024.
Wang Ruize, Tang Duyu, Duan Nan, Wei Zhungyu, Huang Xuanjing, Cao Guihong, Jiang Daxin, Zhou Ming, ers. , K-adapter: Infusing knowledge into pre-trained models with adapters // arXiv preprint arXiv:2002.01800 2020.
Xu Frank F, Alon Uri, Neubig Graham. Why do nearest neighbor language models work? International Conference on Machine Learning. 2023. 38325-38341
Yang An, Yang Baosong, Zhang Beichen, Hui Binyuan, Zheng Bo, Yu Bowen, Li Chengyuan, Liu Dayiheng, Huang Fei, Wei Haoran, others Qwen2. 5 technical report // arXiv preprint a.Xiv2212.15115. 2024.
Yogatama Duni, Masson d' d'Autume Cyprien de, Kong Lingpeng. Adaptive semiparametric language models Transactions of the Association for Computational Linguistics. 2021. 9. 362-373.
Zhang Xiang, Zhao Junbo, LeCun Yamn. Character-level convolutional networks for text classification Advances in neural information processing systems. 2015a. 28.
Zhang Xiang, Zhao Junbo, LeCun Yann. Character-level convolutional networks for text classification Advances in neural information processing systems. 2015b. 28.
Thang Thila, Sabuncu Mert. Generalized cross entropy loss for training deep neural networks with noisy lahels If Advances in neural information processing systems. 2018. 31.
13

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our 3 reflects the paper's main contribution.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Please section
14

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't ' acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: Our method doesn't require thcoretical result.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Detailed training hyperparameters is provided in
Guidelines:
The answer NA means that the paper does not include experiments.
15

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Code and checkpoint has been released https://githie com/ LUMIA JUMIA-Grae Saaad/GGrroo//eeorreeeeeeee
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.cc/public/eee. www.deddeee.oeuuummmcecom ssi ssionyolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
16

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
Please refer to Section 4
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: All experiments of downstream task evaluation and perplexity evaluation are deterministic.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-simma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please refer to
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
17

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We have ensured anonymity in our paper.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: Our work mainly focused on particular field of NLP that doesn't 't have much societal impacts in general.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Justification: Our work domain adaptation doesn't 't require safeguards.
Guidelines:
18

---

The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Please refer the citations in Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- ce-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Currently, our work doesn't contain assets that will be uploaded. Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
19

---

Justification:
The experiments in our work don't don'trew require interaction with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The experiments in our work don't 't require interaction with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: Our works doesn't use LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy for what should or should be described.
20

---

A Interpolation hyperparameter a of all tasks
A.1 Language Modeling Wikitext-103
For language modeling on WikiText- (section p.1). we use the following a values for different GPT-2 model sizes:
Model
GPT-2-small
0.80
GPT-22medium
0.60
GPT-2-large
0.55
GPT-2-xl 0.55
Table 8: Interpolation hyperparameter a for GPT-2 models on WikiText-103.
The trend of smaller for larger GPT models aligns with intuition intuition-stronger base models require less augmentation from the memory component. The general pattern centers around a=0.6, confirming robust default choice.
A.2 Downstream Performance
Table Presentsssts the optimal values for downstream tasks in section 5.2
Task
SST-2
0.30
MR
0.30
CR
0.05
RT
0.20
HYP
0.20
CB
0.30
RTE
0.60
AGN
0.20
Yahoo
0.20
Table 9: Optimal interpolation hyperparameter for downstream tasks.
The general pattern centers around o==033 which is consistent with the findings in Shi et al, (2022).
A.3 Cross-Model and Cross- Vocabulary Adaptation
For domain-specific language modeling tasks (section 5.3 and 5.4 we tune on the validation set by searching over (0.4, 0.6, 0.8, 0.9}.
B Analysis of DAPT Performance on Downstream Tasks
Previous work has shown that domain-adaptive pretraining can adversely affect a model's prompting ability (Cheng et al. 2023 Our experiments reveal that this effect is particularly pronounced when using domain-conditionn PMI (DCPMI) scoring for evaluation, especially on tasks where label verbalizers overlap with the pretraining domain vocabulary.
As shown in Table TO, while direct language modeling evaluation reveals only modest performance drops with DAPT, the DCPMI scores show dramatic degradation for smaller models. This discrepancy arises because we employ fuzzy verbalizers following Shi et al. (2022). and the label spaces for Yahoo and AGN tasks contain terms (e.g., "politics," "technology" that appear frequently in WikiText- 103. When DAPT increases the domain probability for these terms, causes the conditional PMI scores to drop substantially, as the denominator in the DCPMI calculation becomes inflated.
The results for GPT-2-xl demonstrate that larger models exhibit greater robustness to this evaluation artifact, maintaining relatively stable DCPMI scores after domain adaptation. This suggests that
21

---

Model
Yahoo (LM)
Yahoo (DCPMI)
HYP (LM)
HYP (DCPMI)
Avg
GPT-2-small
0.466
0.495
0.639
0.638
0.559
+DAPT
0,429
0.244
0.608
0.361
0.410
-0.037
-0.251
-0.031
-0.277
-0.149
GPT-2-xI
0,520
0,499
0.628
0,609
0.564
+DAPT
0,490
0.491
0.624
0.618
0.556
A
-0.030
-0.008
-0.004
+0.009
-0,008
Table Comparison of standard language modeling (LM) scores versus domain-ccnditional PMI (DCPMI) scores for DAPT models on Yahoo and HYP tasks.
the apparent failure of DAPT on certain downstream tasks is partly an artifact the evaluation methodology rather than fundamental limitation of the approach, though the phenomenon highlights an important interaction between domain adaptation and prompt-based evaluation methods.
C Characteristics of k--N Distributions
C.1 Extreme Sparsity and Concentration
k-NN distributions exhibit fundamentally different characteristics from standard language model out- puts, While LM distributions maintain smooth probability mass across vocabulary with extensive long tails, K--N distributions demonstrate extreme sparsity-typically assigning non-zero probabilities to only 2-3 tokens from 50,257 dimesiaaae vocabulary.
Top16 Probability Comparison
KNN
LM
Mcmueec
0.5
Probabiity
0.3
0.1
Probability Rank Position
Figure 5: Probability distributions from k-NN retrieval, standard LM, and Memory Decoder for GPT-2-Large. The k-NN distribution shows extreme sparsity with concentrated probability mass.
This concentration emerges from two factors: (1) the hard constraint selecting only k nearest neighbors eliminates low-probability candidates, and (2) high-dimensional embedding spaces (e.g., 1280 dimensions for GPT-2-Large) amplify distance relationships through the curse dimensionality, causing nearest neighbors to dominate disproportionatelly
C.2 Scale-Dependent Behavior
Model scale dramatically affects k-NN distribution quality. GPT-2-small (124M) produces dis- tributions marginally different from its LM outputs (top-1 probability 50%), while GPT-2-Large
22

---

(1.5B) generates radically sparse distributions with 93.48% average top-1 probability a 67% relative increase over its baseline.
KNN Distribution Comparison (Top16)
Model Stoee
Smalll Model
Midd Mozel
Large Hodel
Pababiitt
0.6
04
0.0
10
13
Probability Rank Position
Highest)
Figure 6: k-NN distribution sparsity across model scales. Despite identical retrieval parameters (k 1024), larger models produce substantially sparser distributions.
Base Model
Baseline
PPL with MemDec from:
PPL
Small
Medium
Large
GPT-2-Small
24.89
14.01
13.80
13.77
GPT-2-Medium
18.29
12.88
12.74
12.70
GPT-2-Large 15.80
12.05
11.95
11.93
Table 11: Perplexity with Memory Decoder (125M) trained using k-NN distributions from different source models
Larger models benefit from: (1) higher-dimeenioonll spaces where distance concentration intensifies, and (2) superior contextual representations that better disambiguate polysemous tokens and preserve semantic distinctions, leading more coherent nearest neighbor retrievals.
C.3 Domain Adaptation Effects
Fine-tuned models produce sharper k-NN distributions than base models. Domain adaptation creates specialized embedding clusters with reduced intra-cluster variance and increased inter-cluster separation, leading to more decisive retrievals. Memory Decoders trained with fine sene-tnndd distributions consistently achieve lower perplexity, validating that domain domain-adapted retrieval targets.
adapted representations provide superior
D Alternative Loss Functions for Imitating k-NN Distributions
D.1
Failed Approaches
While KL divergence (combined with cross-entropy regularization) successfully matches k-NN distributions, we systematically evaluated several alternative loss functions that all demonstrated inferior performance:
D.1.1
Focal Loss
Adapted from object detection to handle class imbalance through gradient rescaling:
C = =-
Σ ja(1
po puNN) log (o)
(1 (1)) PANN (i)) log(1 po (i))
(8)
23

---

With a = 0.5 and = 2. focal loss theoretically emphasizes hard-to-classify sparse regions but
failed to achieve sufficient distribution concentration in practice.
D.1.2 Jensen-Shannon Divergence
A symmetric alternative to KL divergence:
ISD || Q) PKL (P M) (Ql(Q1) (I M), M (
(9)
Despite avoiding the directional bias of KL divergence, JSD provided no advantage for our extremely sparse target distributions.
D.1.3 Bi-directional Logits Difference (BiLD)
BiLD focuses on relative rankings by computing pairwise differences among top-k logits:
LBILD DxI[Pid KNN IIPeat + (por KNN (
(10)
While theoretically suited for distributions where relative ordering matters more than exact probabili- ties, BiLD consistently underperformed standard KL divergence.
D.1.4 Explicit Sparsity Penalty
Direct penalization of non-zero predictions in zero-probability regions:
Sparss fKL +
Σ
(p)))((==)) Po()), 0.01
(11)
This approach created training instability without meaningfully improving output sparsity.
D.2 Why KL Divergence Succeeds
The superior performance of KI. divergence (with cross-entropy regularization) for matching k-NN distributions likely stems from its unique mathematical properties that align with the retrieval-based nature of the target:
Asymmetric penalty structure: KL divergence DxL(PIIQ) = P()) log Q heavily penalizes
P(i)
P(t
placing probability mass where the target has none (when P(t) R O but Q(i) > 0), while being more forgiving of missing mass where the target has some. This asymmetry naturally encourages sparsity-the model learns to aggressively zero out predictions outside the k-NN support.
Mode-seeking behavior: The forward KL divergence (Pl())))) is inherently mode-sceking, preferring to capture a few probbbbbbayility modes rather than covering the entire distribution. For k-NN distributions with 2-3 dominant modes, this bias perfectly matches the desired behavior, unlike symmetric losses (ISD) or mode-covering alternatives.
Information-theeretie optimality: KL divergence directly minimizes the expected encoding length difference between distributions. For k-NN distributions that encode "retrieval-award uncertainty," KL naturally preserves the information structure maintaninng both the sharp peaks (high retrieval confidence) and the specific ranking among top candidates that emerges from the datastore 's empirical distribution.
The cross cross-entroes cntsopy regularization component anchors the model to linguisticallyy valid outputs, pre- venting collapse to degenerate solutions while the KL term drives sparsity. This combination uniquely balances the competing demands of extreme concentration and semantic coherence.
24


Reviews:
[{"review_id": "iTD2RxIGZf", "reviewer": "Reviewer_hqPH", "summary": "The authors present an approach to augmenting language models with non-parametric/neural adapters that are trained on domain specific data. The training procedure constructs pairs based on a kNN search to find relevant documents based on hidden representations of the pre-trained model. The pre-training objective aligns the memory decoder's output distribution and the kNN distributions for each sample; the overall loss an interpolation of the KL and LM losses. The authors show improved average performance on in-context learning and domain adaptation benchmarks in context, along with ablation studies on the pre-training objective, along with the setup's ability to adapt across different vocabularies.", "questions": "- What is the computational cost and dataset size of the pre-training KNN dataset? How many samples are needed for training? A brief mention of the cost of this for large datasets (even if it pales in comparison to costs incurred during inference) would help.\n- How much does the size of the transformer decoder storing memory affect performance for in-context tasks? Ablations would help elucidate where some of the in-context learning gains are coming from.\n- What's the detailed setup for the full parameter finetuning baseline? There's a width breadth of finetuning approaches with varying levels of effectiveness beyond simply retraining the whole model on the domain-specific corpus.\n- Consider also citing [Eyuboglu et al. 2025](https://arxiv.org/abs/2506.06266) in recent work along with other architectures around memory-based augmentations to architectures such as [Zhang et al.](https://arxiv.org/pdf/2405.06394).\n- What are the capacity limitations of the decoder based on model size? Does performance show degradation with larger-sized in-context learning benchmarks?", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The presented approach is novel, simple, and not particularly resource intensive to train.\n- The manuscript is well-written and easily to understand.\n- The proposed approach shows results that broadly improve performance and give balanced accuracy across relevant benchmarks.\n- The authors ablate some of their proposed approach, including the choice of loss function.\n\nWeaknesses:\n- Ablations top of the decoder architecture itself would elucidate more about the particular attributes of the proposed approach that improve overall performance; this would strengthen the credibility of the approach and give some predictive intuition around how it might scale.\n- The scale of the experiments is not particularly large per larger-scale in-context learning tasks. The authors might consider evaluating on tasks such as [Arora et al.](https://direct.mit.edu/tacl/article/doi/10.1162/tacl_a_00580/117168) or on other larger-scale QA benchmarks such as [Yang et al.](https://arxiv.org/abs/1809.09600) that might require larger-scale domain knowledge.", "was_revised": true, "final_justification": "See my recent responses. The adjustment in my score is due to stolid results from larger in-context learning benchmarks, the ablation studies demonstrating the efficiency of the approach, and the results on other QA tasks being promising.", "rebuttal": "Thank you for the thoughtful feedback. We greatly appreciate the reviewer's recognition of our novel approach and strong empirical results.\n\n> Weakness 1: Architecture choice\n\nWe appreciate the reviewer's suggestion for architecture ablations. Our choice of transformer decoder was driven by **our training objective: learning to predict kNN distributions from context tokens autoregressively**. Transformer decoders have proven exceptional at modeling complex probability distributions, which is precisely what we need for mimicking retriever behavior.\n\nOur results validate this choice—with just 117M parameters, our decoder successfully captures retrieval patterns that would otherwise require searching through massive datastores (e.g., 500GB for Wikitext-103). While we acknowledge that other architectures like specialized memory networks might also be well-suited for distribution imitation, we focused this initial work on demonstrating that pre-trained memory decoders are viable and effective across diverse domains and model scales. We view architectural exploration as a valuable direction for future work.\n\n> Weakness 2: Performance on Large-scale QA benchmarks\n\nWe thank the reviewer for this excellent suggestion. Following the reviewer's recommendation, we conducted additional large-scale experiments on knowledge-intensive tasks, specifically evaluating on Natural Questions (NQ) and HotpotQA benchmarks. We closely followed the experimental setup from [1] to ensure fair comparison.\n\nInterestingly, [1] found that kNN-LM actually **hurts** performance on these reasoning-intensive tasks, despite the answer strings existing in the non-parametric datastore.\n\nTo address this challenge, we trained a 1B Memory Decoder on the same **large heterogeneous unlabeled corpus (513M tokens)** used in the previous study. Our results demonstrate that Memory Decoder successfully overcomes the limitations of traditional retrieval methods:\n\n| Method | NQ | HotpotQA |\n|--------|------|----------|\n| Llama3-8B | 23.64 | 25.14 |\n| + kNN-LM | 24.00 (+0.36) | 24.48 (-0.66) |\n| + Memory Decoder (1B) | **28.01** (+4.37) | **27.72** (+2.58) |\n\nWe show that **Memory Decoder achieves the best of both worlds**: it provides substantial improvements on both benchmarks (+4.37 on NQ, +2.58 on HotpotQA) while avoiding the reasoning degradation observed with kNN-LM. This demonstrates that our approach successfully enhances memorization capabilities through learned retrieval patterns without compromising the model's reasoning abilities—a key limitation of traditional retrieval methods.\n\n> Question 1: Computational Cost and Dataset size\n\nWe appreciate the reviewer's question about training costs. Computationally, Memory Decoder requires FLOPs equivalent to 1 epoch of training a 7B model—negligible compared to DAPT needed for each model from 0.5B to 72B.\n\nFor dataset size, we used ~100M tokens for specialized domains and successfully scaled to **500M tokens** for a large heterogeneous corpus (Wikipedia + CC-NEWS) in response to Weakness 2, demonstrating effective handling of larger datasets.\n\nWhile experiments beyond 500M tokens would provide additional insights, our current evaluation comprehensively demonstrates Memory Decoder's effectiveness across diverse domains and scales. This work establishes a fundamentally novel approach—training plug-and-play memory components to imitate retrieval behavior—opening new research directions in domain adaptation. We view larger-scale exploration as valuable future work as this new paradigm matures.\n\n> Question 2: Ablation on the size of Memory Decoder\n\nWe appreciate the reviewer's question and would like to clarify that the downstream tasks in our paper follow zero-shot evaluation protocols from [2]. Our Memory Decoder is designed to **enhance learning from unlabeled domain corpora**, which is why we primarily compare against DAPT throughout the paper. Our approach should be viewed as orthogonal to methods like in-context learning and supervised fine-tuning. For additional experiments on in-context learning capabilities, please see our response to Question 5.\n\nTo address the reviewer's suggestion, we conducted ablations on Memory Decoder size:\n\n| Model | GPT2-small | GPT2-med | GPT2-large | GPT2-xl | Avg |\n|-------|------------|----------|------------|---------|------|\n| Base | 24.89 | 18.29 | 15.80 | 14.39 | 18.34 |\n| DAPT | 15.47 | 12.78 | 11.10 | 10.16 | 12.38 |\n| +MemDec(117M)  | 13.36 | 12.25 | 11.53 | 10.93 | 12.01 |\n| +MemDec(345M)  | 12.08 | 11.59 | 10.92 | 10.43 | 11.26 |\n| +MemDec(774M)  | **11.67** | **11.23** | **10.83** | **10.28** | **11.00** |\n\nThese results reveal that our small Memory Decoder (117M) already achieves excellent performance across all model sizes. While larger decoders provide incremental improvements, the additional gains do not justify the increase in computational cost.\n\nThis finding validates our design choice of using compact Memory Decoders, demonstrating that effective domain adaptation can be achieved without scaling to larger decoder sizes. The efficiency of our approach makes it practical for real-world deployment.\n\n> Question 3: Setup for DAPT\n\nWe appreciate the reviewer's question about our baseline setup. For full-parameter finetuning, we follow the standard Domain Adaptive Pre-Training (DAPT) approach from [3]. Specifically, we continue full-parameter training on the domain-specific corpus for 1 epoch using a learning rate of 1e-5.\n\nWhile we acknowledge that various DAPT approaches exist with different training strategies like [4], we chose the standard DAPT setup as it remains the most prevalent and widely-adopted method in domain adaptation literature. This baseline provides a well-established benchmark that enables meaningful comparison and validates the effectiveness of Memory Decoder within the broader domain adaptation landscape. Our consistent improvements over this standard approach across multiple domains and model scales demonstrate the practical advantages of our plug-and-play memory paradigm.\n\n> Question 4: Additional references\n\nWe thank the reviewer for bringing these relevant works to our attention. We will include these citations in the camera-ready version.\n\nAfter reviewing these papers, we do note that these works address different problems at smaller scales. Eyuboglu et al. (2025) focuses on task-specific memory augmentation for reasoning, while Zhang et al. explores memory mechanisms mainly on a small synthetic dataset. In contrast, Memory Decoder tackles domain adaptation at scale, enhancing up to 72B parameters across entire domains (500M+ tokens). \n\nNevertheless, these works provide valuable insights, particularly their architectural innovations, that could inform future Memory Decoder designs. We appreciate the reviewer highlighting these connections.\n\n> Question 5: Capacity limitations\n\nWe appreciate the reviewer's thoughtful question about capacity limitations and performance on larger benchmarks.\n\n- Regarding capacity limitations: \n\nRecent theoretical work [5] provides insight into memorization capacity in language models. According to [5], models can store approximately 3.6 bits per parameter when completely eliminating generalization. Given that English text contains 1-2 bits of information per character [6] and following the general approximation of 1 token ≈ 4 characters, **we can estimate that complete memorization of the Wikitext-103 dataset would theoretically require at least 200M parameters**.\n\nRemarkably, our Memory Decoder achieves exceptional memorization and generalization with only 117M parameters on Wikitext-103, suggesting that our proposed paradigm exhibits novel scaling properties distinct from traditional language models. This efficiency indicates the existence of new scaling laws specific to our training objective, presenting exciting opportunities for future theoretical exploration.\n\n- Regarding performance on Large-Scale Benchmarks:\n\nApart from large-scale QA datasets that we showed in our response to Weakness 2, we also conducted extensive evaluations on **13 large-scale in-context learning benchmarks** across all three domains. Due to length constraints, we provide detailed results for all datasets in the last section of our response to Reviewer pgYC. \n\nIn summary, across 13 benchmarks spanning biomedical (ChemProt, MQP, PubmedQA, RCT, USMLE), financial (FiQA_SA, FPB, Headline, NER, ConvFinQA), and legal domains (SCOTUS, CaseHOLD, UNFAIR-ToS), Memory Decoder consistently improves performance while DAPT shows degradation in both zero-shot and few-shot settings, consistent with the findings from [4]. For instance, on finance tasks, DAPT drops average performance from 73.59 to 71.17, while Memory Decoder improves to 74.68. This demonstrates our approach successfully scales to large-sized in-context learning benchmarks.\n\n> References\n\n[1] Geng, Shangyi, Wenting Zhao, and Alexander M. Rush. \"Great Memory, Shallow Reasoning: Limits of $ k $ NN-LMs.\" arXiv preprint arXiv:2408.11815 (2024).\n\n[2] Shi, Weijia, et al. \"knn-prompt: Nearest neighbor zero-shot inference, 2022b.\" URL https://arxiv. org/abs/2205.13792 (2022).\n\n[3] Gururangan, Suchin, et al. \"Don't stop pretraining: Adapt language models to domains and tasks.\" arXiv preprint arXiv:2004.10964 (2020).\n\n[4] Cheng, Daixuan, Shaohan Huang, and Furu Wei. \"Adapting large language models via reading comprehension.\" The Twelfth International Conference on Learning Representations. 2023.\n\n[5] Morris, John X., et al. \"How much do language models memorize?.\" arXiv preprint arXiv:2505.24832 (2025).\n\n[6] Shannon, Claude E. \"Prediction and entropy of printed English.\" Bell system technical journal 30.1 (1951): 50-64."}, {"review_id": "tgWtGoZaNR", "reviewer": "Reviewer_mSKs", "summary": "This paper introduces Memory Decoder (MemDec), a domain adaption method that aims to combine the benefits of Domain Adaptive Pretraining (DAPT) and kNN-based RAG while avoiding their drawbacks. The core idea of the method is to train a transformer decoder. that learns to emulate the output distributions of a non-parametric k-NN retriever. Once trained on a specific domain, the MemDec can be intergrated with any pretrained LLM that shares the same tokenizer. The experiments show that a small MemDec of 0.5B can significantly reduce the perplexity across various models from the Qwen and Llama families.", "questions": "Q1: Can you include evaluations on real-world domain-specific tasks (e.g., biomedical QA or legal reasoning) rather than relying mainly on perplexity and general NLP benchmarks?\n\nQ2: Have you considered comparing against in-context learning baselines in domain settings?\n\nQ3: Could the interpolation weight be adaptive rather than fixed globally? And could you analyze how the interpolation weight affect the performance?", "limitations": "The author assumes full access to domain-specific corpora to build the k-NN datastore for MemDec training, but some domains may lack publicly available data or impose strict privacy constraints. The author may reflect on whether MemDec is feasible in low-resource or privacy-sensitive domains.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength\nS1: The MemDec method distills the benefits of non-parametric retrieval into a small parametric model, providing an efficient alternative to traditional domain adaption. it also reduces the inference latency compared to kNN RAG method while avoiding the cost of full-parameter fine-tuning.\n\nS2: The MemDec has great generalizability with the \"plug-and-play\" capability for models using the same tokenizer.\n\nS3: MemDec achieves significant perplexity reduction across all tested models and domains.\n\nWeakness\nS1: Lack of realistic evaluation. The experiment is centered around perplexity and some downstream NLP tasks. While perplexity is a useful proxy for language modeling, it may not capture effectiveness in domain-specific tasks such as question answering or reasoning.\n\nS2: While MemDec avoid retrieval overhead at inference time, it requires constructing large kNN datastore and performing kNN-based retrieval to generate supervision signal for the training of MemDec, which is very computationally expansive.\n\nS3: There is no comparison to in-context learning baselines, which is a common baseline in modern domain adaption.\n\nS4: The inference-time blending between the base LLM and Memory Decoder is controlled by a static interpolation weight α. However, different inputs, domains, or tasks may benefit from different levels of domain-specific influence, but the paper does not analyze the effects of this parameter or provide justification for the chose value (i.e., α = 0.6 for all tasks and models).", "was_revised": true, "final_justification": "The author has performed extensive additional experiments to address my concerns about lacking realistic evaluation and comparison with in context learning methods, and their results show impressive improvements compared with the baselines.", "rebuttal": "We sincerely thank the reviewer for the thoughtful and constructive feedback. We greatly appreciate the reviewer's recognition of the efficiency gains and plug-and-play capabilities of Memory Decoder.\n\n> Weakness 1 & Question 1: Lack of realistic evaluation\n\nWe appreciate this valuable suggestion and have conducted extensive evaluations on **13 real-world domain-specific tasks** across all three domains (biomedicine, finance, law) and large-scale memory-intensive reasoning tasks (Natural Questions and HotpotQA).\n\nDue to space constraints, we provide detailed results in **the last section of our response to Reviewer pgYC**. Our key findings demonstrate that Memory Decoder significantly outperforms baselines across diverse real-world applications:\n\n- *Domain-specific in-context learning*: Across 13 benchmarks spanning biomedical (ChemProt, MQP, PubmedQA, RCT, USMLE), financial (FiQA_SA, FPB, Headline, NER, ConvFinQA), and legal domains (SCOTUS, CaseHOLD, UNFAIR-ToS), Memory Decoder consistently improves performance while DAPT shows degradation, consistent with the findings from [2]. For instance, on finance tasks, DAPT drops average performance from 73.59 to 71.17, while Memory Decoder improves to 74.68.\n\n- *Knowledge-intensive reasoning*: On Natural Questions and HotpotQA, where traditional kNN-LM actually hurts performance [1], Memory Decoder achieves remarkable improvements of +4.37 and +2.58 points respectively. This demonstrates our method's unique ability to **combine the memorization benefits of non-parametric retrievers with the innate reasoning capabilities of LLMs**, addressing a critical limitation where retrieval-based methods typically degrade reasoning performance.\n\nThese comprehensive evaluations confirm that Memory Decoder's effectiveness extends well beyond perplexity metrics to real-world domain-specific applications where both factual knowledge and reasoning capabilities are essential.\n\n> Weakness 2: Training overhead\n\nWe appreciate the reviewer's attention to computational considerations. While Memory Decoder requires constructing a kNN datastore during training, this cost is significantly amortized compared to other methods:\n\n- *One-time vs. per-model*: Unlike DAPT requiring separate training for each model, Memory Decoder trains once for an entire model family sharing the same tokenizer. Training one 0.5B Memory Decoder is far more efficient than performing DAPT on multiple models (0.5B-72B).\n\n- *Model-agnostic datastore*: **Unlike kNN-LM which requires model-specific datastores** for each base model, our approach constructs the datastore only once. Traditional kNN-LM must rebuild datastores when switching models, while Memory Decoder eliminates this redundancy.\n\n- *Long-term efficiency*: Organizations train one Memory Decoder per domain and use it across their model ecosystem until major version updates (e.g., 9 months between Qwen2 and Qwen3). Section 5.4 demonstrates efficient cross-vocabulary adaptation with minimal training, further amortizing costs.\n\nThe one-time training cost, acknowledged in our limitations, yields significant long-term savings compared to repeated DAPT or maintaining model-specific datastores for each deployment.\n\n> Weakness 3 & Question 2: In-context learning baselines\n\nThank you for this important question about in-context learning comparisons. We'd like to clarify our method's positioning and provide additional experimental evidence:\n\n- *Memory Decoder's positioning*: Memory Decoder is designed to enhance learning from unlabeled domain corpora, which is why we primarily compare against DAPT throughout the paper. Our approach should be viewed as **orthogonal to methods like in-context learning and supervised fine-tuning**, rather than a replacement. The key advantage is that Memory Decoder can work synergistically with these methods.\n\n- *In-context learning experiments*: To address your specific concern, we conducted additional experiments comparing in-context learning performance. Our results from the domain-specific tasks already demonstrate that Memory Decoder outperforms both DAPT and base models in zero-shot and few-shot settings. We further conducted in-context learning ablations on the CaseHOLD legal reasoning task:\n\n| Method | CaseHOLD(micro) | CaseHOLD(macro) |\n|---|---|---|\n| Qwen2.5-7B (0-shot) | 35.92 | 35.93 |\n| +Memory 0.5B (0-shot) | 37.58 | 37.59 |\n| Qwen2.5-7B (4-shot) | 36.34 | 36.35 |\n| +Memory 0.5B (4-shot) | 37.95 | 37.95 |\n| Qwen2.5-7B (8-shot) | 36.34 | 36.35 |\n| +Memory 0.5B (8-shot) | **38.29** | **38.30** |\n| Qwen2.5-7B (16-shot) | 36.09 | 36.10 |\n| +Memory 0.5B (16-shot) | 37.53 | 37.53 |\n\nThese results reveal several key insights: \n\n-  Zero-shot Memory Decoder outperforms all few-shot base models, demonstrating the effectiveness of our domain knowledge encoding.\n-  Memory Decoder continues to benefit from in-context examples, with performance improving from 0-shot to 8-shot settings, which shows our method is orthogonal to in-context learning methods.\n-  Even when both methods experience slight degradation at 16-shot (likely due to noise introduced by increased examples), Memory Decoder maintains its advantage over the base model.\n\nThis demonstrates that Memory Decoder not only provides strong zero-shot domain adaptation but also **preserves and enhances the model's ability to leverage in-context examples**—a critical capability that DAPT often compromises.\n\n> Weakness 4 & Question 3: Adaptive interpolation weight\n\nWe apologize for the confusion in our description. To clarify: we use α=0.6 consistently for all domain adaptation experiments (Sections 5.3 and 5.4), while for Sections 5.1 and 5.2, we tune α on validation sets following established practices from prior work [3,4].\n\nTo address your concern about sensitivity, we conducted comprehensive experiments across all 11 Qwen models on the law domain with different α values:\n\n| α | 0.4 | 0.45 | 0.5 | 0.55 | 0.6 | 0.65 | 0.7 | 0.75 | 0.8 |\n|---|---|---|---|---|---|---|---|---|---|\n| Avg PPL | 4.21 | 4.17 | 4.14 | 4.119 | 4.110 | 4.112 | 4.126 | 4.155 | 4.201 |\n| Relative to α=0.6 | 1.0244 | 1.0146 | 1.0073 | 1.0022 | 1.0000 | 1.0005 | 1.0039 | 1.0109 | 1.0221 |\n\nThe results demonstrate remarkable robustness: performance varies by less than 2.5% across the entire α range (0.4-0.8), with optimal performance near α=0.6. This validates our choice and shows that Memory Decoder is not overly sensitive to this hyperparameter, making it practical for deployment without extensive tuning.\n\nFor completeness, the specific α values used in other experiments were:\n\n- Wikitext-103: (0.8, 0.6, 0.55, 0.55) for GPT2-(small, medium, large, xl) models respectively\n- Downstream tasks: α=0.3 for both Memory Decoder and kNN-LM following [3]\n\nThe trend of smaller α for larger GPT models aligns with intuition—stronger base models require less augmentation from the memory component. The general pattern still centers around α=0.6, confirming it as a robust default choice.\n\nRegarding adaptive interpolation: As proven in [5], non-parametric LMs have very high ceiling if the interpolation weight is optimal. We explored training-free adaptive methods like those mentioned in [6] but observed minimal improvements. For simplicity, we don't include it in our final model. However, we note that the parametric nature of our method makes it particularly suitable for adaptive weights, which we leave for future work.\n\n> Limitation:\n\nThank you for raising this important consideration about privacy-sensitive and low-resource domains. Memory Decoder actually offers significant advantages in these scenarios:\n\n- Privacy-preserving training: For privacy-sensitive domains, differential privacy techniques from [7] can be incorporated during kNN datastore construction. This adds calibrated noise to k-NN searches while maintaining utility, enabling Memory Decoder training on sensitive data without memorizing individual examples.\n\n- Enhanced deployment privacy: Crucially, Memory Decoder actually **offers superior privacy guarantees** compared to non-parametric methods. Unlike kNN-LM/RAG requiring corpus access at inference, Memory Decoder distills knowledge into parameters. Once trained, the original sensitive data can be completely removed from the deployment environment.\n\n- Low-resource advantages: Memory Decoder's efficiency is particularly valuable for low-resource domains. The one-time training cost is amortized across all models in a family, and the compact 0.5B model is far more practical to deploy than large datastores, making domain adaptation accessible in resource-constrained settings.\n\nThese characteristics position Memory Decoder as an attractive solution for organizations dealing with sensitive or limited domain data.\n\n> References:\n\n[1] Geng, Shangyi, Wenting Zhao, and Alexander M. Rush. \"Great Memory, Shallow Reasoning: Limits of $ k $ NN-LMs.\" arXiv preprint arXiv:2408.11815 (2024).\n\n[2] Cheng, Daixuan, Shaohan Huang, and Furu Wei. \"Adapting large language models via reading comprehension.\" The Twelfth International Conference on Learning Representations. 2023.\n\n[3] Shi, Weijia, et al. \"knn-prompt: Nearest neighbor zero-shot inference, 2022b.\" URL https://arxiv. org/abs/2205.13792 (2022).\n\n[4] Khandelwal, Urvashi, et al. \"Generalization through memorization: Nearest neighbor language models.\" arXiv preprint arXiv:1911.00172 (2019).\n\n[5] Xu, Frank F., Uri Alon, and Graham Neubig. \"Why do nearest neighbor language models work?.\" International Conference on Machine Learning. PMLR, 2023.\n\n[6] Li, Minghan, et al. \"Nearest neighbor speculative decoding for llm generation and attribution.\" Advances in Neural Information Processing Systems 37 (2024): 80987-81015.\n\n[7] Zhu, Yuqing, et al. \"Private-knn: Practical differential privacy for computer vision.\" proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020."}, {"review_id": "stVhiQo8fa", "reviewer": "Reviewer_pgYC", "summary": "The authors propose a plug-and-play memory module that enables domain adaptation for LLMs. The so-called memory decoder is designed to imitate retrieval-based methods like kNN-LM but without requiring non-parametric retrieval at inference time. The method is evaluated across various model families (e.g., Qwen, LLaMA) and tasks, demonstrating reduced perplexity and preserved downstream performance. The approach is positioned as a middle ground between DAPT and RAG.", "questions": "- How many runs were performed for the experiment in Table 2?\n- How sensitive is the performance to the alpha value?\n- Can you improve the formatting conventions used in the paper?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n+ The paper is clearly written and well-organized. The method and experimental results are presented in a coherent way. Strengths and weaknesses of the methods have been elaborated. \n \n+ The proposed method is novel and technically sound. The idea of mimicking non-parametric retrieval distributions through KL divergence during pretraining is innovative and addresses a real efficiency bottleneck in domain adaptation.\n\n+ The contribution sounds practically significant. The Memory Decoder module demonstrates strong plug-and-play capability across multiple models sharing the same tokenizer. This flexibility, combined with notable reductions in inference latency compared to retrieval-based methods like kNN-LM, makes the approach attractive for deployment in resource-constrained or latency-sensitive settings.\n\n**Weaknesses:**\n\n- A more thorough factor analysis and clearer interpretability would strengthen the findings. For example, the interpolation weight alpha, which balances the base model and Memory Decoder outputs during inference, is fixed to 0.6 throughout all experiments without enough justification. It will be better if the paper explores how performance varies with different alpha values and whether the gains are robust to this choice.\n\n- Some of the experimental results are not clearly presented. For example, Table 1 includes underlined values without explanation, leaving the readers puzzled about their significance. In Table 2, the best results are not consistently highlighted across all columns. The paper also makes strong claims, such as Memory Decoder outperforming full-parameter fine-tuning on GPT2-medium, without specifying whether these conclusions hold in general. Is this result statistically significant and generally true across all settings? The number of experimental runs and whether results are averaged or from single trials are not disclosed, which leaves the readers puzzled. \n\n- Minor suggestions on formatting. The manuscript uses \\cite{} instead of the appropriate \\citep{}, leading to inconsistent in-text citations (e.g., “Author et al. [2020]”). This indicates a lack of attention to submission formatting standards and should be corrected.  The spacing between tables and their captions is inconsistent (e.g., Tables 1 and 2 have unusually small gaps, while Tables 3 and 4 appear properly spaced). Such inconsistencies detract from the professional appearance of the submission.", "was_revised": true, "final_justification": "After carefully reading the authors’ response and other reviewers’ comments, I am raising my score to 5 (Accept). The authors have addressed my concerns, particularly around the interpolation factor and evaluation on additional datasets. The new results on benchmarks are convincing and strengthen the paper’s contributions. In general, I agree with Reviewer hqPH that the paper is well-written and presents a novel approach to domain adaptation.", "rebuttal": "Thank you for the thoughtful feedback. We greatly appreciate the reviewer's recognition of our novel approach, technical soundness, and practical significance.\n\n> Weakness 1 & Question 2: Alpha sensitivity analysis\n\nWe apologize for the confusion in our description. To clarify: we use α=0.6 consistently for all domain adaptation experiments (Sections 5.3 and 5.4), while for Sections 5.1 and 5.2, we tune α on validation sets following established practices from prior work [1,2].\n\nTo address your concern about sensitivity, we conducted comprehensive experiments across all 11 Qwen models on the law domain with different α values:\n| α | 0.4 | 0.45 | 0.5 | 0.55 | 0.6 | 0.65 | 0.7 | 0.75 | 0.8 |\n|---|---|---|---|---|---|---|---|---|---|\n| Avg PPL | 4.21 | 4.17 | 4.14 | 4.119 | 4.110 | 4.112 | 4.126 | 4.155 | 4.201 |\n| Relative to α=0.6 | 1.0244 | 1.0146 | 1.0073 | 1.0022 | 1.0000 | 1.0005 | 1.0039 | 1.0109 | 1.0221 |\n\nThe results demonstrate remarkable robustness: performance varies by less than 2.5% across the entire α range (0.4-0.8), with optimal performance near α=0.6. This validates our choice and shows that **Memory Decoder is not overly sensitive to this hyperparameter**, making it practical for deployment without extensive tuning.\n\nFor completeness, the specific α values used in other experiments were:\n\n- Wikitext-103: (0.8, 0.6, 0.55, 0.55) for (small, medium, large, xl) models respectively\n- Downstream tasks: α=0.3 for both Memory Decoder and kNN-LM following [2]\n\nThe trend of smaller α for larger models aligns with intuition—stronger base models require less augmentation from the memory component. The general pattern still centers around α=0.6, confirming it as a robust default choice.\n\n> Weakness 2 & Question 1: Experimental clarity and additional results\n\nWe apologize for the confusion in our presentation and thank the reviewer for helping us improve clarity.\n\nFirst, regarding Table 1, the underlined values represent the second-best results. We will clarify this in the revised manuscript to avoid confusion.\n\nTo further validate our claim about Memory Decoder outperforming full-parameter fine-tuning, we conducted additional experiments with different Memory Decoder sizes:\n\n| | small | med | large | xl | Avg |\n|---|---|---|---|---|---|\n| Base | 24.89 | 18.29 | 15.80 | 14.39 | 18.34 |\n| DAPT | 15.47 | 12.78 | 11.10 | 10.16 | 12.38 |\n| +MemDec(small) 117M | 13.36 | 12.25 | 11.53 | 10.93 | 12.01 |\n| +MemDec(med) 345M | 12.08 | 11.59 | 10.92 | 10.43 | 11.26 |\n| +MemDec(large) 774M | **11.67** | **11.23** | **10.83** | **10.28** | **11.00** |\n\nThese results confirm that our observation holds across different scales:\n\n- GPT2-medium + small MemDec (12.25) outperforms GPT2-medium DAPT (12.78)\n- GPT2-large + medium MemDec (10.92) outperforms GPT2-large DAPT (11.10)\n- GPT2-xl + large MemDec (10.28) slightly underperforms GPT2-xl DAPT (10.16) by only 1.2%\n\nThis pattern demonstrates that Memory Decoder consistently achieves competitive or superior performance compared to full-parameter fine-tuning while maintaining the crucial advantage of being plug-and-play across all model sizes.\n\nRegarding experimental runs: For perplexity experiments, we performed single runs as the evaluation is deterministic. For downstream tasks in Table 2, we follow the domain-conditional PMI scoring method from [2], the result is also deterministic since it does not involve sampling-based generation, therefore we also only perform single runs. We will clarify this methodological detail in our revised manuscript to ensure transparency.\n\n> Weakness 3 & Question 3: Formatting improvements\n\nThank you for your constructive feedback on formatting. We sincerely apologize for these oversights. The citation formatting issue arose from switching our bibliography style from IEEE to plainnat during preparation, which requires different citation commands.\n\nWe will correct this in the camera-ready version by:\n\n- Replacing \\cite{} with \\citep{} for all references that do not function as grammatical elements in the sentence.\n- We will standardize the spacing between all tables and their captions to maintain consistency throughout the manuscript.\n\nWe appreciate the reviewer bringing these formatting issues to our attention, as they help us improve the professional presentation of our work. All formatting corrections will be implemented in the final version.\n\n> **Evaluation on modern tasks**\n\nFollowing suggestions by other reviewers, we conducted extensive evaluations on **13 real-world domain-specific tasks** across all three domains and large-scale memory-intensive reasoning tasks.\n\n- Domain-Specific Downstream Tasks\n\nFollowing the evaluation framework from [3], we tested on large-scale in-context learning benchmarks across biomedicine, finance, and law domains. Our results demonstrate that Memory Decoder maintains strong performance in both zero-shot and few-shot settings, while DAPT shows significant performance degradation—consistent with findings in [3] that DAPT adversely affects prompting abilities.\n\n*Biomedical Domain:*\n\n| | ChemProt(13-shot) | MQP(4-shot) | PubmedQA(0-shot) | RCT(10-shot) | USMLE(0-shot) | Avg |\n|---|---|---|---|---|---|---|\n| Qwen2.5-7B | 24.4 | 83.44 | 63.7 | 70.1 | **36.92** | 55.71 |\n| +DAPT | 17.59 | 76.22 | 65.7 | 21.0 | 34.64 | 43.03 |\n| +Memory 0.5B | 24.4 | **84.09** | **64.4** | **74.06** | 36.84 | **56.76** |\n\n\n*Financial Domain:*\n\n| | FiQA_SA(5-shot) | FPB(5-shot) | Headline(5-shot) | NER(20-shot) | ConvFinQA(0-shot) | Avg |\n|---|---|---|---|---|---|---|\n| Qwen2.5-7B | 80.46 | 70.96 | 87.08 | 68.92 | 60.53 | 73.59 |\n| +DAPT | 75.59 | 66.39 | 86.03 | **69.32** | 58.52 | 71.17 |\n| +Memory 0.5B | **81.34** | **71.25** | **87.95** | 69.21 | **63.69** | **74.68** |\n\n*Legal Domain:*\n\n| | SCOTUS (micro)(0-shot) | SCOTUS (macro) | CaseHOLD (micro)(0-shot) | CaseHOLD (macro) | UNFAIR-ToS(4-shot) | Avg |\n|---|---|---|---|---|---|---|\n| Qwen2.5-7B | 26.66 | 17.90 | 35.92 | 35.93 | 87.05 | 40.69 |\n| +DAPT | 28.33 | 16.82 | 35.70 | 35.69 | **87.18** | 40.74 |\n| +Memory 0.5B | **31.66** | **21.05** | **37.58** | **37.59** | 87.05 | **42.99** |\n\nThis **preservation of in-context learning capabilities while enhancing domain-specific performance** represents a significant advantage over traditional domain adaptation methods.\n\n- Knowledge-Intensive Reasoning Tasks\n\nFollowing [4], we evaluated on Natural Questions and HotpotQA. Notably, [4] found that kNN-LM actually hurts performance on these reasoning-intensive tasks, despite the answer strings existing in the non-parametric datastore.\n\nTo address this challenge, we trained a 1B Memory Decoder on the same large heterogeneous unlabeled corpus (513M tokens) used in [4]:\n\n| Method | NQ | HotpotQA |\n|--------|------|----------|\n| Llama3-8B | 23.64 | 25.14 |\n| + kNN-LM | 24.00 (+0.36) | 24.48 (-0.66) |\n| + Memory Decoder (1B) | **28.01** (+4.37) | **27.72** (+2.58) |\n\nMemory Decoder provides substantial improvements on both benchmarks while maintaining strong reasoning capabilities, unlike kNN-LM which experiences performance degradation. This demonstrates our approach successfully **enhances memorization capabilities through learned retrieval patterns without compromising reasoning abilities**—validating our method's effectiveness on knowledge-intensive applications where both factual recall and reasoning are crucial.\n\nThese comprehensive results confirm that Memory Decoder successfully scales to modern benchmarks.\n\n> References\n\n[1] Khandelwal, Urvashi, et al. \"Generalization through memorization: Nearest neighbor language models.\" arXiv preprint arXiv:1911.00172 (2019).\n\n[2] Shi, Weijia, et al. \"kNN-Prompt: Nearest Neighbor Zero-Shot Inference.\" arXiv preprint arXiv:2205.13792 (2022).\n\n[3] Cheng, Daixuan, Shaohan Huang, and Furu Wei. \"Adapting large language models via reading comprehension.\" The Twelfth International Conference on Learning Representations. 2023.\n\n[4] Geng, Shangyi, Wenting Zhao, and Alexander M. Rush. \"Great Memory, Shallow Reasoning: Limits of $ k $ NN-LMs.\" arXiv preprint arXiv:2408.11815 (2024)."}, {"review_id": "zesc6MwxkL", "reviewer": "Reviewer_2Vi7", "summary": "This paper proposes a memory decoder, an LM that is trained to mimic a nonparametric LM (specifically, kNN-LM) and then augmented with a standard LM at inference time. The idea is that for domain adaptation, continued training (or fine-tuning) leads to catastrophic forgetting, whereas standard retrieval-augmented LMs or nonparametric LMs increase the latency. The memory decoder aims to address these problems and benefit from both worlds. Experiments are done in various evaluation settings including: a standard perplexity eval, downstream task fine-tuning, and domain adaptation. Results show that the proposed method outperforms baselines including retrieval-augmented LMs, kNN-LM, full fine-tuning and LoRA tuning, while having very little computational/latency overhead.", "questions": "Written in the previous field", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n- The problem is well motivated – the limitations of standard training and retrieval augmentation are well explained, and the memory decoder is shown to be effective in achieving the best of both worlds.\n- Experiments are comprehensive - although it is on a small scale and evaluated on relatively easy datasets by today’s standard, it is quite comprehensive and includes a thorough set of baselines.\n- All the execution steps are all reasonable and sound.\n- The paper is well written and is easy to follow.\n\nWeaknesses/questions\n- A fundamental question that arises to me after reading this paper is – a nonparametric LM and finetuning of LM have different properties, rather than one being universally better than the other, e.g., a nonparametric model may be better at long-tail or uncommon patterns like factoid knowledge, while finetuning of an LM may make the model generalize better. Is the proposed model closer to a nonparametric LM, or to finetuning of LM, in terms of model properties? I think this is a fundamental question that should be discussed and analyzed.\n- The paper claims that  the memory decoder prevents catastrophic forgetting that standard fine-tuning would have, which is likely because it is still preserving the original LM. One important baseline would be to train a small model fine-tuned on the target data and augment that with the original model, e.g., through logit ensembling, or weight merging (if the model sizes are the same). like how memory decoder is being incorporated currently. I think this is likely a competitive baseline for catastrophic forgetting (related work: [1], which is using a model soup). However, there is a chance memory decoder still has interesting properties that this baseline won’t have (related to the first bullet point).\n- Question: One benefit of nonparametric LMs is that the data can be replaced without training, although it still requires the datastore & the nearest neighbor index to be updated. In the case of memory decoder, it requires building a nearest neighbor index AND training a new LM. How much additional overhead is training? I assume it depends on the training steps. If that cost is negligible, I think it is good to highlight; if not, a trade-off analysis would be helpful.\n- Question: In Section 5.2, is each model fine-tuned individually on the training set of each benchmark and evaluated on its corresponding test set? I believe the very poor performance of full-param fine-tuning on some datasets like HYP, AGN, and Yahoo is quite suspicious – especially for AGN and Yahoo, as they are balanced and stable, unlikt HYP which is a tiny dataset with imbalanced labels. Understanding this result would be helpful.\n- As mentioned earlier briefly, experiments are done on a small scale and evaluated on easy datasets by today’s standard, e.g., there are many published papers evaluated on these benchmarks with GPT-2 models in 2021-2022, but not recently. Also, many experiments are using PPL as a metric, e.g., domain adaptation serves as a primary motivation for the whole work, but is only doing PPL evals, which are less convincing.\n\nOverall, I don’t see major flaws in the paper – it is a solid and well-executed paper, and includes experiments that the paper is supposed to have. However, I think the paper is not *exciting* enough, largely because it doesn’t fully explore some of the deeper questions it could have touched on, e.g., where the memory decoder sits in the spectrum of nonparametric LMs vs. fine-tuning (related to the first bullet point). Moreover, expanding experiments is likely to signify its impact, e.g., adding stronger models with more challenging benchmarks, adding downstream evals to domain adaptation experiments, and including the baseline mentioned in thes second bullet point.\n\n[1] “Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time” https://proceedings.mlr.press/v162/wortsman22a.html", "was_revised": true, "final_justification": "Revised after rebuttal", "rebuttal": "We sincerely thank the reviewer for the thoughtful and constructive feedback, and we are pleased that the reviewer found our work well-motivated, comprehensively evaluated, and clearly written.\n\n> Weakness 1: Positioning of Memory Decoder\n\nWe appreciate the reviewer's insightful question about where Memory Decoder sits in the spectrum between non-parametric LMs and fine-tuning. This is indeed an important question that we are happy to address.\n\nMemory Decoder is **fundamentally a parametric method** that learns to mimic the behavior of non-parametric retrievers. Our key motivation is to compress the knowledge stored in large non-parametric datastores into a compact parametric model, thereby **combining the memorization capabilities of non-parametric methods with the efficiency and generalization abilities of parametric approaches**.\n\nWe hypothesize that Memory Decoder can effectively learn long-tail knowledge (similar to non-parametric methods) while maintaining the coherence and reasoning capabilities of parametric models. Unlike non-parametric methods that rely on simple similarity matching and suffer from reasoning tasks (as noted in [1]), our parametric approach should preserve the model's ability to perform reasoning and maintain semantic coherence.\n\nTo validate our hypothesis, we conducted case studies on WikiText-103 examining two key aspects (We compare label probability of different methods on the highlighted token, using tokens before as context):\n\n- Long-tail Knowledge Learning:\n\n| Sentence | Memory Decoder | kNN-LM | Base LM |\n|----------|---------------|---------|----------|\n| he starred alongside actors Mark Strong and Derek ***Jacobi***. | **68.94%** | 9.39% | 0.12% |\n| The launch of HMS Dreadnought in ***1906*** by the Royal Navy raised the stakes | **98.65%** | 40.62% | 1.57% |\n\nMemory Decoder assigns significantly higher probabilities to factual long-tail knowledge compared to the base LM, successfully capturing the memorization benefits of non-parametric methods.\n\n- Semantic Coherence and Basic Reasoning:\n\n| Sentence | Memory Decoder | kNN-LM | Base LM |\n|----------|---------------|---------|----------|\n| In 2000 Boulter had a guest-starring role ***on*** the television series The Bill. | 40.11% | 8.07% | **45.51%** |\n| the Royal Armoured Corps in Britain created three tank squadrons for special overseas operations, known as 'A', 'B' and '***C***' Special Service Squadrons. | 50.10% | 10.76% | **63.04%** |\n\nMemory Decoder maintains probabilities closer to the base LM rather than blindly following kNN-LM's lower probabilities, demonstrating its ability to preserve coherent language modeling capabilities.\n\nThese results suggest that Memory Decoder successfully occupies a unique position: **it enhances memorization of domain-specific and long-tail knowledge like non-parametric methods, while maintaining the generalization and reasoning capabilities inherent to parametric models**. We further validate this hypothesis through extensive downstream evaluations of ***13 modern datasets*** (see response to Question 3 and the last section of response to reviewer pgYC for detail), where Memory Decoder consistently outperforms both non-parametric methods and standard fine-tuning approaches.\n\n> Weakness 2: Ablations on fine-tuned model\n\nWe thank the reviewer for this important suggestion. We have indeed compared with this baseline in our ablation study (Table 5), where \"CE Only\" represents a small model fully fine-tuned on the domain corpus. Our results show that augmenting the base LLM with this fine-tuned model cannot achieve the same performance as Memory Decoder (4.69 vs 3.59), and we will highlight this observation more prominently in our revised manuscript.\n\nTo further validate our approach, we conducted additional ablation studies on GPT-2 models, demonstrating that Memory Decoder provides substantial improvements over logit ensembling with a fine-tuned model:\n\n| Model | Base | +Finetune(small) | +MemDec(small) |\n|-------|------|------------------|----------------|\n| small | 24.89 | 15.95 | 13.36 |\n| medium | 18.29 | 14.26 | 12.25 |\n| large | 15.80 | 13.13 | 11.53 |\n| xl | 14.39 | 12.30 | 10.93 |\n| Avg | 18.34 | 13.91 | 12.01 |\n\nMemory Decoder **achieves an additional 1.9 perplexity improvement on average compared to the fine-tuned model baseline**, demonstrating that our novel training objective (combining distribution alignment with language modeling) enables superior memorization of domain knowledge beyond standard fine-tuning.\n\nRegarding Model Soup, we appreciate the suggestion but note that its main contribution focuses on \"averaging weights of multiple models fine-tuned with different hyperparameter configurations.\" Our method has a fundamentally different objective—learning to mimic non-parametric retrieval behavior—and achieves much better generalizability across different model sizes and architectures. We leave the exploration of weight merging approaches as an interesting direction for future work.\n\n> Question 1: Trade-off analysis\n\nWe thank the reviewer for this important question about training overhead. The reviewer is correct that non-parametric LMs offer data updates without retraining, though they still require rebuilding the datastore and kNN index.\n\nFor Memory Decoder, our training budget is comparable to standard DAPT approaches—approximately 1 epoch of training a 7B model, which we consider a reasonable one-time cost for domain adaptation.\n\nWe acknowledge this trade-off, but emphasize significant practical advantages:\n\n- Model-independent deployment: While **kNN-LM requires building model-specific datastores for each model**, a single Memory Decoder serves multiple models sharing the same tokenizer.\n\n- Inference efficiency: Memory Decoder incurs only 28% inference overhead vs kNN-LM's 117% (Figure 4). Additionally, kNN-LM requires 500GB memory for 100M tokens, while Memory Decoder compresses this into a 0.5B model.\n\nThe one-time training cost is offset by dramatic improvements in inference efficiency, storage, and cross-model reusability—highly favorable for production scenarios with stable domain data. We will make this trade-off analysis more explicit in our revised manuscript.\n\n> Question 2: Understanding significant drop on AGN and Yahoo\n\nWe appreciate the reviewer's question and would like to clarify that the downstream tasks in our paper follow zero-shot evaluation protocols from [2] and do not involve fine-tuning on individual datasets. We apologize for the confusion caused by the term \"full-param fine-tuning\"—this should be understood as Domain Adaptive Pre-training (DAPT) on the Wikitext-103 corpus. We will revise this terminology for clarity.\n\nPrevious work has shown that DAPT can adversely affect a model's prompting ability [3]. We believe this effect is amplified in our evaluation due to our use of domain-conditional PMI (DCPMI) scoring. To investigate this, we compared standard language modeling (LM) scores with DCPMI scores:\n\n| Model | Yahoo(LM) | Yahoo(DCPMI) | HYP(LM) | HYP(DCPMI) | Avg |\n|-------|-----------|--------------|---------|------------|-----|\n| gpt2-small | 0.466 | 0.495 | 0.639 | 0.638 | 0.559 |\n| +DAPT | 0.429 | 0.244(-0.251) | 0.608 | 0.361(-0.277) | 0.410 |\n| gpt2-xl | 0.520 | 0.499 | 0.628 | 0.609 | 0.564 |\n| +DAPT | 0.490 | 0.491 | 0.624 | 0.618 | 0.556 |\n\nWhile direct LM evaluation shows modest drops, DCPMI scores drop significantly for GPT-2 small. We hypothesize that since we use domain-conditional scoring with fuzzy verbalizers (following [2]), and since words in Yahoo and AGN's label space (e.g., \"politics,\" \"technology\") appear frequently in Wikitext-103, DAPT increases the domain probability for these terms, causing DCPMI scores to drop dramatically. The GPT-2 XL results show that larger models are less susceptible to this dataset bias, but the general trend holds. We will clarify this evaluation setup in our revised manuscript.\n\n> Question 3: Evaluation on modern datasets\n\nWe conducted extensive evaluations on 13 real-world domain-specific tasks across all domains and large-scale reasoning tasks (NQ and HotpotQA). Due to limited space, detailed results are shown in the last section of our response to Reviewer pgYC.\n\nKey findings:\n\n- Domain-Specific Tasks: Memory Decoder consistently outperforms DAPT across all domains: biomedical (56.76 vs 43.03), financial (74.68 vs 71.17), and legal (42.99 vs 40.74). While DAPT hurts prompting ability [3], Memory Decoder preserves and enhances zero-shot/few-shot performance.\n\n- Memory-Intensive Reasoning: Unlike kNN-LM which hurts reasoning [1], Memory Decoder achieves substantial improvements: NQ (28.01 vs 24.00 for kNN-LM) and HotpotQA (27.72 vs 24.48 for kNN-LM).\n\nThese results validate that Memory Decoder successfully combines memorization benefits with reasoning capabilities, addressing the limitations of both approaches on real-world tasks.\n\n> Conclusion\n\nWe sincerely thank the reviewer for the constructive feedback. We have addressed the deeper questions about Memory Decoder's position in the spectrum between non-parametric LMs and fine-tuning through detailed case studies (Response to Weakness 1), expanded the suggested baseline comparisons (Response to Weakness 2), and evaluated on 13 modern datasets and challenging reasoning benchmarks (Response to Question 3). We hope these additions have addressed your concerns about the paper's depth and impact, and we will incorporate all these analyses into our revised manuscript.\n\n> References\n\n[1] Geng, Shangyi, Wenting Zhao, and Alexander M. Rush. \"Great Memory, Shallow Reasoning: Limits of $ k $ NN-LMs.\" arXiv preprint arXiv:2408.11815 (2024).\n\n[2] Shi, Weijia, et al. \"knn-prompt: Nearest neighbor zero-shot inference, 2022b.\" URL https://arxiv. org/abs/2205.13792 (2022).\n\n[3] Cheng, Daixuan, Shaohan Huang, and Furu Wei. \"Adapting large language models via reading comprehension.\" The Twelfth International Conference on Learning Representations. 2023."}]

Author rebuttal:

Meta-review: The work proposes memory decoder, a small pretrained decoder that can adapt LLMs to specialized domains without domain specific pre-training or RAG. This is achieved by training a memory module to imitate the output the distribution of a k-NN retriever. This is a well motivated work solving an important problem. The proposed approach is novel and shown to be effective based on empirical studies. The "plug-and-play" nature of the proposal makes the work quite practical. 

The authors' rebuttal addressed most of the concerns, and all reviewers recommended acceptance after the rebuttal and discussion period. Overall, the AC finds the work to be a good contribution to the community. Please incorporate the new experiments and discussions in the final version. 

One question from 2Vi7 worths further discussions in the final version:

"a nonparametric LM and finetuning of LM have different properties, rather than one being universally better than the other, e.g., a nonparametric model may be better at long-tail or uncommon patterns like factoid knowledge, while finetuning of an LM may make the model generalize better. Is the proposed model closer to a nonparametric LM, or to finetuning of LM, in terms of model properties?"

The authors' reply and additional experiments in the rebuttal show that "Memory Decoder can effectively learn long-tail knowledge (similar to non-parametric methods) while maintaining the coherence and reasoning capabilities of parametric models. "

However it is easy to construct an external knowledge set such that the "long-tail" knowledge cannot be fully compressed/memorized by a small LM. For example, in the extreme case, one can construct a table (as external knowledge) mapping a set of keys to a set of (random) values, and the size of the table is larger than the size of the small memory decoder. In such a case, it is not possible to compress the knowledge into the memory decoder. 

The authors should tone down the statement a bit and discuss the limitations.
Decision: Accept (poster)