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

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.

Weakness 1: Positioning of Memory Decoder

We 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.

Memory 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.

We 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.

To 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):

• Long-tail Knowledge Learning:

Memory Decoder assigns significantly higher probabilities to factual long-tail knowledge compared to the base LM, successfully capturing the memorization benefits of non-parametric methods.

• Semantic Coherence and Basic Reasoning:

Memory 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.

These 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.

Weakness 2: Ablations on fine-tuned model

We 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.

To 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:

Memory 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.

Regarding 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.

Question 1: Trade-off analysis

We 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.

For 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.

We acknowledge this trade-off, but emphasize significant practical advantages:

• 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.

• 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.

The 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.

Question 2: Understanding significant drop on AGN and Yahoo

We 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.

Previous 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:

While 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.

Question 3: Evaluation on modern datasets

We 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.

Key findings:

• 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.

• 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).

These results validate that Memory Decoder successfully combines memorization benefits with reasoning capabilities, addressing the limitations of both approaches on real-world tasks.

Conclusion

We 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.

References

[1] Geng, Shangyi, Wenting Zhao, and Alexander M. Rush. "Great Memory, Shallow Reasoning: Limits of $k$ NN-LMs." arXiv preprint arXiv:2408.11815 (2024).

[2] Shi, Weijia, et al. "knn-prompt: Nearest neighbor zero-shot inference, 2022b." URL https://arxiv. org/abs/2205.13792 (2022).

[3] Cheng, Daixuan, Shaohan Huang, and Furu Wei. "Adapting large language models via reading comprehension." The Twelfth International Conference on Learning Representations. 2023.

Thank you for the thoughtful feedback. We greatly appreciate the reviewer's recognition of our novel approach, technical soundness, and practical significance.

Weakness 1 & Question 2: Alpha sensitivity analysis

We 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].

To address your concern about sensitivity, we conducted comprehensive experiments across all 11 Qwen models on the law domain with different α values:

The 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.

For completeness, the specific α values used in other experiments were:

• Wikitext-103: (0.8, 0.6, 0.55, 0.55) for (small, medium, large, xl) models respectively
• Downstream tasks: α=0.3 for both Memory Decoder and kNN-LM following [2]

The 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.

Weakness 2 & Question 1: Experimental clarity and additional results

We apologize for the confusion in our presentation and thank the reviewer for helping us improve clarity.

First, regarding Table 1, the underlined values represent the second-best results. We will clarify this in the revised manuscript to avoid confusion.

To further validate our claim about Memory Decoder outperforming full-parameter fine-tuning, we conducted additional experiments with different Memory Decoder sizes:

These results confirm that our observation holds across different scales:

• GPT2-medium + small MemDec (12.25) outperforms GPT2-medium DAPT (12.78)
• GPT2-large + medium MemDec (10.92) outperforms GPT2-large DAPT (11.10)
• GPT2-xl + large MemDec (10.28) slightly underperforms GPT2-xl DAPT (10.16) by only 1.2%

This 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.

Regarding 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.

Weakness 3 & Question 3: Formatting improvements

Thank 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.

We will correct this in the camera-ready version by:

• Replacing \cite{} with \citep{} for all references that do not function as grammatical elements in the sentence.
• We will standardize the spacing between all tables and their captions to maintain consistency throughout the manuscript.

We 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.

Evaluation on modern tasks

Following 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.

• Domain-Specific Downstream Tasks

Following 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.

Biomedical Domain:

Financial Domain:

Legal Domain:

This preservation of in-context learning capabilities while enhancing domain-specific performance represents a significant advantage over traditional domain adaptation methods.

• Knowledge-Intensive Reasoning Tasks

Following [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.

To address this challenge, we trained a 1B Memory Decoder on the same large heterogeneous unlabeled corpus (513M tokens) used in [4]:

Memory 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.

These comprehensive results confirm that Memory Decoder successfully scales to modern benchmarks.

References

[1] Khandelwal, Urvashi, et al. "Generalization through memorization: Nearest neighbor language models." arXiv preprint arXiv:1911.00172 (2019).

[2] Shi, Weijia, et al. "kNN-Prompt: Nearest Neighbor Zero-Shot Inference." arXiv preprint arXiv:2205.13792 (2022).

[3] Cheng, Daixuan, Shaohan Huang, and Furu Wei. "Adapting large language models via reading comprehension." The Twelfth International Conference on Learning Representations. 2023.

[4] Geng, Shangyi, Wenting Zhao, and Alexander M. Rush. "Great Memory, Shallow Reasoning: Limits of $k$ NN-LMs." arXiv preprint arXiv:2408.11815 (2024).

Dear authors,

Sorry, we didn't realize that this year's NeurIPS setting is withholding the Final Justification at this Discussion phase. Below, please see the paragraph we posted two days ago:

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.

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.

Weakness 1 & Question 1: Lack of realistic evaluation

We 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).

Due 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:

• 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.

• 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.

These 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.

Weakness 2: Training overhead

We 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:

• 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).

• 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.

• 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.

The 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.

Weakness 3 & Question 2: In-context learning baselines

Thank you for this important question about in-context learning comparisons. We'd like to clarify our method's positioning and provide additional experimental evidence:

• 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.

• 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:

These results reveal several key insights:

• Zero-shot Memory Decoder outperforms all few-shot base models, demonstrating the effectiveness of our domain knowledge encoding.
• 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.
• 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.

This 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.

Weakness 4 & Question 3: Adaptive interpolation weight

We 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].

To address your concern about sensitivity, we conducted comprehensive experiments across all 11 Qwen models on the law domain with different α values:

The 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.

For completeness, the specific α values used in other experiments were:

• Wikitext-103: (0.8, 0.6, 0.55, 0.55) for GPT2-(small, medium, large, xl) models respectively
• Downstream tasks: α=0.3 for both Memory Decoder and kNN-LM following [3]

The 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.

Regarding 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.

Limitation:

Thank you for raising this important consideration about privacy-sensitive and low-resource domains. Memory Decoder actually offers significant advantages in these scenarios:

• 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.

• 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.

• 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.

These characteristics position Memory Decoder as an attractive solution for organizations dealing with sensitive or limited domain data.

References:

[1] Geng, Shangyi, Wenting Zhao, and Alexander M. Rush. "Great Memory, Shallow Reasoning: Limits of $k$ NN-LMs." arXiv preprint arXiv:2408.11815 (2024).

[2] Cheng, Daixuan, Shaohan Huang, and Furu Wei. "Adapting large language models via reading comprehension." The Twelfth International Conference on Learning Representations. 2023.

[3] Shi, Weijia, et al. "knn-prompt: Nearest neighbor zero-shot inference, 2022b." URL https://arxiv. org/abs/2205.13792 (2022).

[4] Khandelwal, Urvashi, et al. "Generalization through memorization: Nearest neighbor language models." arXiv preprint arXiv:1911.00172 (2019).

[5] Xu, Frank F., Uri Alon, and Graham Neubig. "Why do nearest neighbor language models work?." International Conference on Machine Learning. PMLR, 2023.

[6] Li, Minghan, et al. "Nearest neighbor speculative decoding for llm generation and attribution." Advances in Neural Information Processing Systems 37 (2024): 80987-81015.

[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.

Thank you for the thoughtful feedback. We greatly appreciate the reviewer's recognition of our novel approach and strong empirical results.

Weakness 1: Architecture choice

We 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.

Our 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.

Weakness 2: Performance on Large-scale QA benchmarks

We 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.

Interestingly, [1] found that kNN-LM actually hurts performance on these reasoning-intensive tasks, despite the answer strings existing in the non-parametric datastore.

To 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:

We 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.

Question 1: Computational Cost and Dataset size

We 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.

For 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.

While 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.

Question 2: Ablation on the size of Memory Decoder

We 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.

To address the reviewer's suggestion, we conducted ablations on Memory Decoder size:

These 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.

This 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.

Question 3: Setup for DAPT

We 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.

While 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.

Question 4: Additional references

We thank the reviewer for bringing these relevant works to our attention. We will include these citations in the camera-ready version.

After 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).

Nevertheless, these works provide valuable insights, particularly their architectural innovations, that could inform future Memory Decoder designs. We appreciate the reviewer highlighting these connections.

Question 5: Capacity limitations

We appreciate the reviewer's thoughtful question about capacity limitations and performance on larger benchmarks.

• Regarding capacity limitations:

Recent 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.

Remarkably, 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.

• Regarding performance on Large-Scale Benchmarks:

Apart 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.

In 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.

References

[1] Geng, Shangyi, Wenting Zhao, and Alexander M. Rush. "Great Memory, Shallow Reasoning: Limits of $k$ NN-LMs." arXiv preprint arXiv:2408.11815 (2024).

[2] Shi, Weijia, et al. "knn-prompt: Nearest neighbor zero-shot inference, 2022b." URL https://arxiv. org/abs/2205.13792 (2022).

[3] Gururangan, Suchin, et al. "Don't stop pretraining: Adapt language models to domains and tasks." arXiv preprint arXiv:2004.10964 (2020).

[4] Cheng, Daixuan, Shaohan Huang, and Furu Wei. "Adapting large language models via reading comprehension." The Twelfth International Conference on Learning Representations. 2023.

[5] Morris, John X., et al. "How much do language models memorize?." arXiv preprint arXiv:2505.24832 (2025).

[6] Shannon, Claude E. "Prediction and entropy of printed English." Bell system technical journal 30.1 (1951): 50-64.