MLP Memory: A Retriever-Pretrained Memory for Large Language Models
Authors: Rubin Wei
Jiaqi Cao
Jiarui Wang
Jushi Kai
Qipeng Guo
Bowen Zhou
Zhouhan Lin
Venue: ICLR 2026 Poster (2026)
Keywords: external memory
parametric memory
Primary area: foundation or frontier models, including LLMs

Abstract: Modern approaches to enhancing Large Language Models' factual accuracy and knowledge utilization face a fundamental trade-off: non-parametric retrieval-augmented generation (RAG) provides flexible access to external knowledge but suffers from high inference latency and shallow integration, while parametric fine-tuning methods like LoRA risk catastrophic forgetting and degraded general capabilities. In this work, we propose MLP Memory, a lightweight parametric module that learns to internalize retrieval patterns without explicit document access. By pretraining an MLP to imitate a $k$NN retriever's behavior on the entire pretraining dataset, we create a differentiable memory component that captures the benefits of retrieval-based knowledge access in a fully parametric form. Our architecture integrates this pretrained MLP Memory with Transformer decoders through simple probability interpolation, achieving 12.3\% relative improvement on five question-answering benchmarks and 5.2 points absolute gain across nine general NLP tasks, while reducing hallucinations by up to 10 points on HaluEval. Moreover, MLP Memory delivers 2.5$\times$ faster inference than RAG with superior accuracy. Our findings show that learning retrieval patterns parametrically bridges the gap between efficient inference and effective knowledge access, offering a practical alternative to both RAG and fine-tuning approaches.

Published as a conference paper at ICLR 2026
MLP MEMORY: A RETRIEVER-PRETRANEE MEMORY FOR LARGE LANGUAGE MODELS
Rubin Wei Jiaqi Caol.4 Jiarui Wang'
Jushi Kai*
Qipeng Guo?
Bowen Zhou Zhouhan Lin
LLUMIA Lab, Shanghai Jiao Tong University
2 Shanghai Artificial Intelligence Laboratory
Flectronic Engineering, Tsinghua University
*STU Paris Elite Institute of Technology
we' rubruunngeee nn@gmail 11. com, 1 wmw..hhhhngggmmee
Code & model available at https://github... httpss//gitthbb..o///ll/// http://ww....oouuuaaaaaarrrrraaa
ABSTRACT
Modern approaches to enhancing Large Language Models' factual accuracy and knowledge utilization face fundamental trade-off: non-parametric retrieval- augmented generation (RAG) provides flexible access to external knowledge but suffers from high inference latency and shallow integration, while parametrie fine- tuning methods like LoRA risk catastrophic forgetting and degraded general capa- bilities. In this work, propose MLP Memory, lightweight parametric module that learns to internalize retrieval patterns without explicit document access. By pretraining MLP to imitate ANN retriever's behavior the entire pretraining dataset, we create a differentiable memory component that captures the benefits of retrieval-based knowledge access fully parametric form. Our architec- ture integrates this pretrained MIP Memory with Transformer decoders through simple probability interpolation, achieving 12.3% relative improvement on five question-answerigg benchmarks and 5.2 points absolute gain across nine general NLP tasks, while reducing hallucinations by up to points on HaluEval. More- over, MLP Memory delivers 2.5x faster inference than RAG with superior accu- racy. Our findings show that learning retrieval patterns parametrically bridges the gap between efficient inference and effective knowledge access, offering practi- cal alternative to both RAG and fine-tuning approaches.
?
NQ
WebQA
TriviaQA
Time First token()) TTFT(ms)
Tokens second())
TPS(tokens/s)
28
40
84
48
23
37
140
36
120
8x
20
6.5%
67ms
17
28
60
55
40
3.0x
11
Base
+CPT
-MPP Memory
LM(7B)
+RAG
+NNN-LM
-ML Memory(1B)
Figure 1: Performance and efficiency comparison. Left: accuracy across three QA benchmarks.
MLP Memory consistently outperforms the base model, surpassing both parametric methods (CPT, LoRA) and nor-parametric retrieval (RAG). Right: inference efficiency, measured by time to first token (TTFT, + lower is better) and tokens per second (TPS, + higher is better). RAG results shown for top-5 retrieval. kNN-LM is accelerated via dimension reduction
and both RAG and &NN-LM use the Wikipedia-200. -021 retrieval corpus. MLP Memory uses IB parameters.
"Equal contribution.
Zhouhan Lin is the corresponding author.

---

Published as a conference paper at ICLR 2026
Non-Parametrie
Parametric
Query
RAG
Documents
Corpus
Query
Continue Pre-train
Retrieve
MLP Memory
RAG-Inspired Training
Efficient Inference w/o RAG
ML.P Memory Imitate
RAG
Query
MLP Memory
Interpolate
Figure 2: Approaches to enhance factual accuracy and knowledge utilization. Top left: Non- parametric RAG provides flexible knowledge access but suffers from high latency. Top right: Para- metric fine-tuning risks catastrophic forgetting. Bottom: MLP Memory learns retrieval patterns during training (left) and enables efficient inference without explicit retrieval (right).
INTRODUCTION
Decoder-only architectures such GPT [Bowwn et al. 2020), LLaMA (Grattafiori et al. 2024),
Qwen (Qwen et al., 2025), and DeepSeek (Liu et al. 2024) have achieved remarkable success in various tasks, including open-ended text generation (OpenAl et al. 2024), code completion (Chen
ett al. 202T). image synthesis (Chen et al.] [2020). and multimodal reasoning (Liu et al., 2023) However, despite their impressive capabilities, these models often struggle with effective knowl- edge utilization, producing responses that may be fluent but fail to accurately leverage the factual information encoded in their parameters.
Current approaches to enhance knowledge utilization in LLMs face significant trade-offs. Retrieval- augmented generation (RAG) methods (Lewis et al.) 2021 Peng et al., [023] Gao al., [202]] Izacard et al. 2022) dynamically fetch relevant documents ground model outputs, providing flex- ible access to external knowledge sources. However, these non-parametric approaches introduce substantial inference latency through expensive nearest-neighhor searches and longer context from retrieved documents. They also suffer from shallow integration with the base model, as the retrieval component remains isolated from the LLM's computational graph. Conversely, parametric adap- tation methods such as continued pre-training (CPT) and LoRA (Hu et al., 2022) directly modify model weights to incorporate comain-speciffe knowledge. While computationally efficient infer- ence time, these approaches risk catastrophic forgetting of previously learned capabilities and often degrade performance on general tasks, requiring careful task-specific tuning that limits their broader applicability. Figure iillsstrtes how our approach differs fundamentally from both non -on-maracctiic retrieval methods and parametric adaptation approaches.
In contrast to decoder-only LLMs, neuroscience research reveals a lateralized human brain where language processing is dominated by the left hemisphere while memory formation occurs in the hippocampus (Gazzanigal 2005b a: Douglas 1967). This insight has inspired gemory-augmennt models machine learning. Early approaches like Memory Networks (Weston al., 2015) enabled read/write operations external memory, while Sparse Access Memory introduced differentiable memory access schemes. However, these were task-specific with limited general applicability. In the LLM era, methods such Memory Transformers (Burtsev al.) (021) incorporate trainable memory tokens for global context, while AutoCompressors (Chevalier et al. 2023) compress long contexts into summary vectors. Nevertheless, these memory tokens primarlly function work- ing memory supplements for context extension rather than long-term memory capable of retaining information from the entire training corpus.
In this work, we propose an external memory for LLM that pretrained to mimic a retriever on the entire pretraining dataset. Specifically, following the RAG setting in &-N-LMM (Khandelwal al.
2

---

Published as a conference paper at ICLR 2026
RAG
Refer to the background document and answer the question:
Background: SS Security Service, known the SD-Amt, became official security organization the Nazi Party 1934. Consisting at first paid agents and a few hundred unpaid informants scattered across Germany, the SD was quickly professionalized under Heydrich, who commissioned National Socialist academics and lawyers ensure that the and its Security Service in particular, operated \"withnn the framework of National Socialist ideology."" Heydrich was given the power to select men for the SS Security Service from among any SS subdivisions since Himmler considered the organization of the SD as important, September 1939, the SD was divided into two departments. the interior department (Inland-SD) and the foreign department (Ausland-SD), and placed under the authority of the Reich Security Main Office (RSHA).
Question: In WWII, who was the head of the Nazi party's security service? The answer is:
Output: Heinrich Himm
Base LM
MLP Memory
Answer the question:
Answer the question:
Question: In WWII, who was the head of the Nazi party's security service? The answer is:
Question: In WWII, who the head of the Nazi party's security service? The answer is:
Output: Question: In WW
Output: Reinhard Heydrich.
Comparison of model outputs on factual question. Despite retrieving relevant documents with correct information (highlighted in green), RAG is misled by contextual distractors and pro- duces an incorrect answer. MLP Memory generates the correct answer without explicit retrieval.
Figure 3:
2020), this memory learns to map the LLM hidden state at certain step a vocabulary distribution matching the output of the &NN retriever. During inference, the I.MM's 'ss native output is interpolated with the retriever- preetraiedd output from the external memory. Our resulting architecture, illustrated in Figure 4 consists of transformer decoder and external MLP memory, each pretrained sepa- rately with different pretraining tasks. For our pretrained external memory, we aim to achieve the following features simultaneously:
1) End-to-end differentiabiliity Unlike the non-parametric nature of retrievers, MLP mem- ory is fully parameterized and allows gradieat flow during training. This enables end-to-end joint optimization of the entire model architecture.
2) Highly compressible memory. The MLP memory compresses large datastores (c.g.. 40TB for 5B tokens in &NN-LM) into a compact parametric form (c.g.. 4GB for IB parameters storing 5B tokens), facilitating efficient deployment without performance degradation.
3) Low inference-time latency. MLP memory eliminates costly retrieval operations, achieving 2.5x faster inference than RAG methods and 5.6x faster inference than kNN-LM when us- ing a 5B-token retrieval corpus. Crucially, unlike retrieval-based approaches, our method's inference speed remains constant regardless of the retrieval corpus size.
4) Long-term memory, covering the whole pretraining corpus. While existing memory tokens serve primarily working memory by storing local context for immediate use, our MLP memory functions as a long-term repository of generalizable knowledge acquired during the pretraining phase.
Experimental results demonstrate that MLP Memory significantly outperforms existing approaches across multiple dimensions. It achieves average relative improvements of 12.3% (Mistral-7B) and 7.8% (Llama2-7B) on five QA benchmarks, with WebQA showing exceptional gains (37.45% vs. 29.28% baseline). On nine general NLP tasks, it delivers a 5.2 points absolute improvement. MLP Memory also substantially reduces hallucinations on HaluEval, with accuracy improvements of 9.68, 10.08, and 2.14 points on dialogue, QA, and summarization tasks respectively. Most notably, it achieves 2.5x faster ing constant inference speed regardless of corpus size, unlike retrieval methods whose latency scales with data size. iillustrats MLP Memory's 's performance gains and inference efficiency over baselines and Figure demonstrates case where ML. Memory correctly answers factual ques-
than RAG and 5.6x faster than ANN-LM, while maintain-
3

---

Published as a conference paper at ICLR 2026
tions while RAG fails despite retrieving correct information. These results confirm that parametric compression of retrieval patterns offers a more efficient and effective alternative to explicit retrieval.
PRELIMINARY, K-NEAREST NEIGHBORS LANGUAGE MODEL
The ANN-LM Khandelwal et al. 2020) augments a pre-trained LM by interpolating its parametric distribution with non- -on-aammet distribution from nearest neighbor retrieval. Given context e2 = (Wi, W/-1), 1), and w, denotes the next token. The next-token probability is:
p(w, c.) = X PENN (wi ()) + (1 - X) PLad (we ct),
(1)
where À E 10, is the interpolation parameter, PLM is the LM's distribution, and PENN is retrieval- based distribution.
Datastore Constructed via forward pass corpus, the datastore consists of key-value pairs (ka, ve) where k = f(ce) encodes context es using LM representations, and oy is the next token wy: (XC, V2 ((f(c)) w.) (ca, we) D)
(2)
Inference The LM encodes context c into query f(e) and retrieves k-nearest neighbors / from (K,, V) using distance metric d(:.-) (typically squared L2).. The non-parametric distribution is:
PENN (y | c) x Σ Ly= exp(-d(ki, (c))).
While &NN-LM improves predictions through explicit memory, it suffers from substantial storage requirements and high-latency retrieval. For instance, the Wikitext- 103 datastore requires nearly 500 GB of storage even for the GPT2-small model (He et al.) 202T). These limitations motivate our MLP Memory, a compact parametric model pretrained to approximate the retrieval function: given query embedding, it directly outputs kNN-like next token distribution, thereby eliminating both the substantial storage requirements and high-latency retrieval.
MLP MEMORY
In this section, we present MLP Memory, lightweight parametric module that learns to internalize retrieval patterns without explicit document access. Our approach consists three key components: stack of MLPs that processes hidden representations without token-mixing operations (Section 3.I. a specialized pre-training procedure that enables the MLP to mimic non-parametric retrieval distributions (Section (.2),, and an efficient inference mechanism for deployment (Section 3.3). As illustrated in Figure 4 MLP Memory first learns to mimic non-parametric retrieval distributions dur- ing pre-training (Figure ())) then seamlessly integrates with the language model during inference (Figure [a)), eliminating both the storage requirements of large datastores and the computational cost of nearest neighbor search.
3.1
ARCHITECTURE
Our MLP Memory learns to mimic non- paraparmeetiic retrieval by mapping query embeddings to KNN distributions. Given query q f(c) from context c, the MI.P directly predicts PANNN (y]c) without neighbor search, transforming discrete retrieval into a differentiable mapping M Rd RIV), where d is the embedding dimension and V is the vocabulary size.
In designing the memory module, we observe from Section 2 that the retriever imitation task processes single-vector representation without requiring token-mixing operations. Recent stud- ies (Geva et al., 2020) have identified that FFN layers function as key-value memories, suggesting that MLPs play specialized role in knowledge memorization within LLMs. Based these in- sights, we propose pretraining an all-MLP memory that effectively functions as a non-parametric retriever, as illustrated in Figure 4
The MLP Memory takes hidden representations f(c) from the pretrained LM as input and is trained to predict the corresponding &NN distribution PENN c) as its target. Once trained, the MLP's is out- put distribution is interpolated with the LM's parametric distribution during inference, following the same interpolation scheme as &-NN-MM but without requiring datastore access neighbor search.
4

---

Published as a conference paper at ICLR 2026
(a) Inference
(b) Training
P(welo)) (P + (1 (W:le))
== alyı + (1 (1-a)AcE
Pors
L
Beiies
London
Ace
PeN(Weloe)
(Welee)
groundtruth
Pawcy (Welce)
Bejiz
Parss
Palt
Lendon
Lecakn
Reijing
LLM
Memory
&NN Retriever
Softmax
Soflmax
Training Contexts
Targets
Bejiing
Layer
MLP
The Fiffel Tower stands
Add & Norm
Brisiin 's expital is
London
M
hg
h
R-11
h
Feed Forward
Layer K
Layer
Training Time
You'll find Bifiel Tower
Figure 4: Overview of MLP Memory architecture. (a) Inference: MLP Memory processes context representations from specific L.MM layer, generating token probabilities that are interpolated with LLM outputs for final predictions. (b) Training: MLP Memory learns to imitate retriever behavior using LLM representations as input and distributions generated by kNN retrievers as targets, opti- mized through hybrid objective.
3.2 TRAINING
The training procedure for MLP Memory consists of two primary stages; constructing supervision signals from non-parame con-paraeetrie retrieval distributions, and optimizing the MLP to mimic these distri- butions through a carefully designed loss function.
Data Construction To generate supervision for training MLP Memory, we leverage the datastore construction process described in Section 2 We build the datastore (AC, V) through a forward pass the training corpus, storing context representations and their corresponding
tokens. For each training example (Cz, wt) E D. we compute the non-parae parametric distribution P&NN (ylc.) by retrieving k-nearest neighbors from the datastore. To prevent trivial self-retrieval that would con- taminate the learning signal, we exclude the query itself from the neighbor set when constructing the target distribution. These beteeeddinndistribuuuu offline and cached for efficient training.
pairs ((f(c)) PANN (-lex))) are precomputed
Loss Function Unlike traditional language modeling with single-label targets, ANN distributions capture the diversity of plausible continuations by encoding multiple valid next tokens weighted by their contextual similarity. Our ablation studies Section 5.4 demonstrate that hybrid objective combining two complementary losses yields optimal performance. Our approach centers min- imizing the Kullback- Leibler divergence (Van Erven & Harremos 2014] output distribution and the cached iNN distributions:
between MLP Memory's
(xL(c)) (L(peNN(-le)) || PMLP(-1ee)
(4)
This encourages the memory module match the full probability distribution rather than merely predicting the most likely token. To prevent excessive deviation from the underlying corpus distri- bution, we integrate complementary Cross-Entrope loss (Zhang & Sabuncu 2018):
Lce(c) log (ogpm)((w))))
(5)
5

---

Published as a conference paper at ICLR 2026
The final training objective balances these two components through hyperparameter at
[(c,) a. CkL(c)) + (1-a)
-L)))) LCx(c))
(6)
The KL term encourages learning distributional patterns while the CE term ensures accurate ground- truth prediction, preventing the overfitting that occurs with cross-entropy alone.
3.3 INFERENCE
Once trained, MLP Memory integrates with the base language model through simple probability interpolation. During inference, MLP Memory processes hidden representations from the language model MLM and produces a distribution that is interpolated with the LM's output:
Pfinat (w)lc) A PMMP + (1- - A) PLAs (welez)
where À E [0, controls the influence of retrieval-based knowledge.
approaches that require nearest neighbor search and extended context processing, MLP Memory requires only a single forward pass through lightweight all-MLP ar- chitecture. demonstrated Figure [1] our method achieves 2.5x faster tim-tto-frss--okee than RAG (top-5) and 5.6x faster than kNN-LM, despite &NN-LM employing dimension reduction from 4096 to 256 for acceleration. For tokens per second, MLP Memory delivers 1.5x higher through- put than RAG and 6x higher than ANN-LM, while introducing only 1.2x overhead relative to the base model. Crucially, this performance remains constant regardless of retrieval corpus size, unlike retrieval-baseee methods whose latency scales with datastore size.
Unlike
4 EXPERIMENTAL SETUP
Overview We conduct comprehensive experiments to evaluate MLP Memory across four critical dimensions. First, we assess performance on five question-answering benchmarks (5.1) demon- strate that our approach represents novel form of parametric memory that surpasses both traditional parametric methods (continued pretraining, LoRA) and por-mmametriic approaches (RAG). Second, we evaluate on fundamental NI.P tasks (5.2) to verify that integrating MLP Memory preserves the base model's general capabilities. Third, we examine hallucination reduction (5.3) on HaluEval to validate our method's 's effectiveness in improving factual accuracy. Finally, we present ablation study (5.4 to analyze design choices such as loss weighting and layer selection.
Implementation Details We conduct our experiments 32xA800 80GB GPUs. To demonstrate the generalizability our approach, we employ two distinct backbone models: Llama-2-7B (Tou- vron et al., 2023) and Mistral-7B-vO.. (Jiang al., 2023). For question-answerie
benchmarks, we build key-value datastores and non non-parametric distributions using both models on preprocessed Wikipedia-2021 (Izacard al. 2022). and train separate BB-parameter MILP Memory modules with learning rate 4e-4. The MLP Memory uses layers by default. See Appendix [ for details. For general NLP tasks, build datastores using Mistral-7B-v0.3 on a heterogencous corpus follow- ing Geng et al. (2024). and train the MLP Memory with learning rate 4c-4. For hallucination evaluation, directly apply the MLP Memory trained from question-answering experiments. All experiments use a training budget equivalent to the computational cost of training a 7B parameter model for 1 epoch. The training hyperparameter a is set to 0.4 across all tasks. The interpolation hyperparameter À is tuned on the validation split of each task following Khandelwal et al. 2020)). see more details in Appendix D
Baselines We compare MIP Memory against established methods for improving factual accuracy and knowledge utilization: RAG, which employs BGE (Chen et al. (2024) as the retrieval model and retrieves top-5 documents to ensure comprehensive context coverage. ANN-LM (Khandelwal al. 2020), configured with interpolation parameter = 0.1 and temperature 7 = 10.0 following (Geng et al. 2024). LoRA (Hu et al., 2022), applied to query, key, value, and MLP layers, with rank adjusted to match the parameter count our MLP Memory modules. Continued Pretraining (CPT), which involves further training of all model parameters on the corresponding corpus.

---

Published as a conference paper at ICLR 2026
Table Question answering performance across five benchmarks. Positive gains shown in green and negative changes in red. Percentage in parentheses denotes the relative improvement over the base model. All methods use the same Wikipedia corpus for training or retrieval.
Open-Domain QA
Long-fore form QA
Multhhop QA
Methods
Average
NQ
WebQA
TriviaQA
TruthfulQA
HoipoiQA
Llama2-7B
23.18
32,09
56.91
29.16
22.72
32.81
Non-parametric methods
+RAG
14.60-8.58
36.71+4.62
62.20+5.299
31.59+2.48
19.60-3.J
$2.94(=0.49%
+ANN-LM
23.16-0.02
33.4644137
47.31
29.22+0:06
22.66-006
33.16(-1.1%)
Paranetric methods
+CPT
12,90 10.28
31.55 0.54
58.81 +.90
29.56
15.49 7223
29.66( 9.6%%
+LoRA
17.88-5.30
35.19+3.10
58. 14+1
28.33-089
17.18-5.54
31.34( 4.5%%
+MLP Mem
27.04
36.61+4152
57.50+0.59
30.04+0088
25.69+297
35,38( 7.8%%
Mistral-7B-vii
20.63
29.28
57.65
32,09
20.96
32.12
Non-parametric methods
+RAG
1.95
24.90
54.21
35.47++3.88
29.77+8.81
3
+ÀNN-LM
21.05 +-0.2
30.51
57.77 40.12
32.33
21.20
32.57(+
Parametric methods
+CPT
12.16 8.47
34.0614.78
61.21 13.56
29.18 2.91
16.04 4.92
30.53( (1.53(-5.0%%
+LoRA
18.17 18.47
4..50
61.60 43.9
30.91 -1.18
16,23 4.73
32.28( -0.5%%
+MLP Mem
25.20++455
37.45+8.17
60.99 +3.34
32.54+0:45
24.14+3.18
36.06(+ (1.33%)
Performance nine general NLP tasks spanning sentiment classification, textual entail- ment, and topic classification. indicate improvement over the Misvral-7B-v0.3 baseline, while indicate decreased performance.
Table 2:
Sentiment Classification
Textual.
Topic.
Methods
Average
SST2
MR
CR
RT
HYP
CB
RTE
AGN
Yahoo
Mistral-7B-v0.3
81.21
75.35
62.30
74.95
55.42
69.64
59.57
75.95
56.36
67.86
Non-purumetric methods
+RAG
87.201
83.701
71.551
82.361
$4.65
57.14_
66.431
75.64.]
58.431
70.791
+ANN-LM
82.151
76.851
61.70]
74.95
56.781
71.42"
60.287
76.131
56,26]
68.501
Parametric methods
+CPT
87.091
82.851
82.601
77.487
60.65t
14_
52.71
83.101
51.56]
70.581
+LoRA
86.541
83.201
75.101
79.831
55.42
51.78.
56.31
65.46]
57.301
67.881
+MLP Mem
83.191
79.901
75.951
75.421
64.151
76.79*
64.621
80.281
57.331
73.07
EXPERIMENTAL RESULTS
5.1
QUESTION ANSWERING PERFORMANCE
We evaluate MLP Memory five diverse QA benchmarks: Natural Questions (NQ) (Kwiatkowski
et al. 2019), WebQA (Berant et al. 2013), TriviaQA (Joshi et al. 2017). TruthfulQA (Lin et al.] 2022 and HotpotQA (Yang al. 2018), comparing against CPT. LoRA, and RAG. As shown in Table Mistral-7B--0.3 with MLP Memory achieves an average relative improvement of 12.3% over the baseline across five benchmarks. with particularly striking improvements on NQ (25,20% vs. baseline (0.63%) and WebQA (37,45% vs. baseline 29.28%). While CPT and LoRA suffer sig- nificant degradation across all tasks tasks-likely due to catastrophic forgetting during cemaai--ppecffc training-Mi. Memory maintains improves performance by learning emulate retrieval behav- ior without modifying the base model's parameters. Notably, our approach outperforms both RAG and KNN-LM, even though they leverage the same Wikipedia-2022 corpus for retrieval inference time, suggesting that our parametric compression of retrieval patterns captures richer contextual relationships than explicit document retrieval. The consistent gains across both factoid QA (NQ, TriviaQA) and multi-hop reasoning (HotpotQA) demonstrate that MLP Memory effectively bridges the gap between parametric and par-marametric memory systems.

---

Published as a conference paper at ICLR 2026
Table 3: Performance on HaluEval benchmark across question answering, dialogue, and summa- rization tasks. Results show accuracy (%). RAG is not evaluated on summarization as this task requires only the source document.
Dialogue
QA
Summarization
Mistral-7B-v0.3
57.18
53.99
50.27
+CPT
31.68-5.50
46,49-7.55
47.39-
+LoRA
55.51-1.66
50.02-3.0
50,38+0.11
+RAG
59.06+1.88
65.09+11.D
+MLP Mem
66.86+9.68
64.07+10.00
52.41+2..1
5.2 GENERAL NLP TASK PERFORMANCE
To ensure MLP Memory doesn't 't compromise fundamental language understanding, we evaluate on nine standard NLP tasks spanning sentiment classification (SST-2 (Socher et al. 2013), MR (Pang)
& Lee, 2005a)) CR (Hu & Liu 2004). RT (Pang & Leel 2005b), HYP (Kiesel et al. 2019 tex- tual entailment (CB (De Marneffe al. 2019), RTE (Dagan al. 2010 and topic classification (AGNews (Zhang et al. (2015a), Yahoo (Zhang et al. 2015b)). Table 2 reveals that MLP Memory achieves comprehensive improvements across all general tasks, achieving the highest average score compared to all baselines. The improvements are particularly pronounced on reasoning-intenc tasks like RTE (64.62% vs. baseline 59.57%) and CB (76.79% vs. baseline 69.66 69.64%) suggesting that the retrieval patterns learned by MIP Memory provide useful inductive biases even for tasks that don't explicitly require factual knowledge retrieval. contrast, CPT and LoRA show mixed results with improvements some tasks but degradation others. The robust performance across this diverse task suite demonstrates that MLP Memory's 's external parametric memory complements rather than interferes with the base model's learned representations.
5.3 HALLUCINATION REDUCTION
We assess MIP Memory's ability to reduce hallucinations using HaluEval (Li et al. 2023) across three generation tasks: dialogue, question answering, and summarization, where models must iden- tify factual inconsistencies in generated content. As shown in Table parametric methods (CPT and LoRA) severely degrade performance, confirming the risks of weight modification. MLP Memory consistently improves hallucination detection across all three domains, with gains of 9.68, 10.08, and 14 points respectively. These substantial improvements indicate that the retrieval patterns encoded in MLP Memory significantly help the model better distinguish factual from hallucinated content. The effectiveness across diverse generation contexts-from free-form dialogue to constrained sum- marization- suggests that MLP Memory's 's learned retrieval behavior provides a general mechanism for grounding language generation in factual knowledge. This hallucination reduction, combined with strong QA performance and enhanced general capabilities, validates our core hypothesis that decoupling memorization from reasoning through retriever-pretrained external memory can enhance factual accuracy without the typical trade-off's of parametric or parametric approaches.
5.4 ABLATION STUDY
Ablation Setup We conduct ablation experiments across three GPT2 (Radford et al. 2019) scales: small (12 layers), medium (24 layers), and large (36 layers), paired with corresponding MLP Mem- ory modules of 117M, 345M, and 774M parameters respectively. All experiments are evaluated on WikiText- 103 (Merity et al. 2016) to investigate loss weighting and optimal layer selection.
Impact of Loss Weighting We examine how balancing KL and CE losses affects retriever im- itation by varying a from 0.0 to 1.0. As Figure 5aa) shows, extreme values produce suboptimal res values prevent the MLP memory from learning from the ANN distribution, while high values cause overfitting to the language modeling objective. The optimal balance occurs at a = 0.4, indicating both objectives necessary. KL divergence leverages the information-rice kNN distri- bution, enabling more effective generalization, while CE loss provides essential token-level predic- tion accuracy. This balanced approach prevents overfitting while maintaining predictive power.
8

---

Published as a conference paper at ICLR 2026
(8)
(b)
100
GPr2-Large GPT2-Med
100
GPTe-Lage (3L))
GPT2-Med (24L)
GPT2-Small (121.)
Nomm(PL)) (%6)
Norn(PPL) (*))
60
60
40
20
0.0
0.1
0.2
0.3
0.4
0.5
0.5
0.7
0.8
0.9
1.0
20
40
60
80
100
Loss Weight
Layer Percentage (%)
Figure 5: (a) Impact of KL loss weight retriever imitation. Lower PPL (min-max normalized for clarity) indicates better performance, with optimal balance at a = 0.4. (b) Impact of input layer depth MIP Memory performance decoder stack (e.g., 70% corresponds to layer 25 in GPT2-large).
model sizes. Layer percentage denotes depth in the
Which Layer Provides the Best Representation for MLP Memory? While ANN-LM performs best using the input to the final feedforward layer as the retrieval key, our MLP Memory consistently achieves optimal performance at around 70% of network depth, regardless of model scale. Our finding aligns with Memorizing Transformers (Burtsev et al. (021). which also selected around 75% depth for optimal retrieval performance. We evaluate GPT2-small (12 layers), GPT2-medium (24 layers), and GPT2-large (36 layers), attaching the MLP Memory to various transformer blocks. As shown in Figure [5]b), the x-axis indicates relative depth (20%-100%), and the y-axis shows min-max normalized perplexity (0% best, 100% worst). This consistent pattern across all model sizes contrasts with the &NN-LM convention of using final-layer representations.
RELATED WORK
Retrieval-Augmeente Generation RAG (Lewis al. 2021 Peng et al. 2023 Gao et 2022) mitigates hallucinations by grounding generation in external knowledge. Despite improving factual accuracy, RAG faces limitations: retrieval latency, coarse granularity, and limited LLM integra- tion (Zhang et al., 2024). Recent work (Su et al. 2025) explores enhanced retrieval with LLM priors. Our approach proposes parametric memory mimicking non- parametrii retrieval, eliminat- ing explicit document retrieval while preserving knowledge augmentation.
Various architectures explored memory augmentation, from Memory Networks (Weston et al., 2015) with explicit read-write components to Memory Transformers (Burtsev et al. 2021) with extended attention. Long.Mmm (Wang et al. 2023) and MemoRAG (Qian et al. 2025) introduced decoupled architectures for long-term history storage. While these focus on context extension, our MLP memory expands across the entire pre-training corpus, enabling long- term generalizable knowledge storage.
Memory-Augmented Language Models
MLP Architectures All-MLP architectures emerged as transformer alternatives, with gMLP (Lu al. 2021) matching transformer performance and sparse MLPs (Yu al. 2022) showing superior training efficiency. Studies (Geva et af.) 2020) identified FFN layers as key-value memories in LLMs. Inspired by this, propose pretraining all-MLP memory non-parametric retriever, leveraging MLPs' memorization capabilities for a compact, differentiable knowledge store.
7 CONCLUSION
In this paper, we introduced MLP Memory, a novel approach for enhancing language models by learning to internalize retrieval patterns. By pretraining lightweight MLP module to imitate kNN retriever behavior on the entire pretraining corpus, MLP Memory captures the benefits of retrieval- augmented generation in fully parametric form, without requiring explicit document access.

---

Published as a conference paper at ICLR 2026
The key advantage of MLP Memory lies in its efficiency and effectiveness. Our approach achieves 12.3% relative improvement on question-anae answering benchmarks, 5.2 points gain on general NLP tasks, and up to 10 points reduction in hallucinations-wwwe while delivering 2.5x faster inference than RAG and maintaining constant speed regardless of corpus size. Unlike parametric fine-tuning that risks catastrophic forgetting or non-parametric RAG that suffers from high latency, MLP Memory enhances model capabilities without these typical trade-off's.
MLP Memory introduces a new paradigm for language model enhancement that fundamentally reimagines how models access and utilize knowledge. By parametrically encoding retrieval behav- ior through a pretrained memory component, our approach creates more efficient, accurate, and scalable framework that bridges the gap between parametric and non-parametric methods.
ACKNOWLEDGMENTS
This work is sponsored by the National Natural Science Foundation of China (NSFC) grant (No. 62576211) and the National Key Research and Development Program of China (No. 20232D0121402). It is also the result of collaborative project on novel language model archi- tectures between Shanghai Jiao Tong University (SJTU) and the Shanghai Artificial Intelligence Laboratory. The computational resources required for pretraining the models were provided by the Shanghai AI Lab. This work is also supported by the Specialized Program on Fundamental Research from Science and Technology Commission of Shanghai Municipality (No. 2025SH2DDZX0256009,
REFERENCES
Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. Semantic parsing on freebase from question-answer pairs. In Proceedings of the 2013 conference on empirical methods in natural language processing, pp. 1533-1544, 2013.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Berbert-Vosss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https://aree //tp:v//xxvvvv. org/ abs 14165
Mikhail S. Burtsev, Yuri Kuratov, Anton Peganov, and Grigory V. Sapunov. Memory transformer, 2021.
11527
Jianly Chen, Shitao Xiao, Peitian Zhang, Kun Luo, Defu Lian, and Zheng Liu. Bge m3-embedding: Multi-lingual, fultiiffuctionnalit,, quati-granullritt text embeddings through self-knowledge dis- tillation. arXiv preprint arXiv:2402.032/1 2024.
Mark Chen, Alec Radford, Rewon Child, Jeff Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In Proceedings of the 37th International Conference on Ma- chine Learning, ICML'20 20. JMLR.org, 2020.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 2021.
Xin Cheng, Xun Wang, Xingxing Zhang, Tao Ge, Si-Qing Chen, Furu Wei, Huishuai Zhang, and Dongyan Zhao. xrag: Extreme context compression for retrieval-augmmented generation with one token. Advances in Neural Information Processing Systems, 37:109487-109516, 2024.
Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts, 2023. T /ttp//aaiiii wtt.aaaivvvvvo... abaggg/ab///3300 2305. 14788
Ido Dagan, Bill Dolan, Bernardo Magnini, and Dan Roth. Recognizing textual entailment: Rational, evaluation and peproocches-erratunn Natural Language Engineering, 16(1):105-105, 2010.
10

---

Published as a conference paper at ICLR 2026
Marie-Catherin De Marneffe, Mandy Simons, and Judith Tonhauser. The commitmentbank: In- vestigating projection in naturally occurring discourse. In proceedings of Sinn und Bedeutung, volume 23, pp. 107-124, 2019.
Robert Douglas. The hippocampus and behavior. Psychological bulletin, 67(6):416, 1967.
Luyu Gao, Xueguang Ma, Jimmy Lin, and Jamie Callan. Precise zero-shot dense retrieval without relevance labels, 2022. URthttts:: //appss:////vvvi////a http//aawwv..r..aaaaaeeeaeeoeeee 2221 10496
Michael S Gazzaniga. The ethical brain. Dana press, 2005a.
Michael S Gazzaniga. Forty-five years of split-brain research and still going strong. Nature Reviews Neuroscience, (18,,,53-659 2005b.
Shangyi Geng, Wenting Zhao, and Alexander M Rush. Great memory, shallow reasoning: Limits of k nn-lms. arXiv preprint arXiv:2408. 11815, 2024.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers key-value memories. arXiv preprint arXiv:2012.14911 14911,,,. 2020.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21773, 2024.
Felix Hamborg, Norman Meuschke, Corinna Breitinger, and Bela Gipp. news-please: A generic news crawler and extractor, March 2017. URL https: // /oi. /ttp:////o..rrr/15211111111 zzeodo.
4120316
Junxian He, Graham Neubig, and Taylor Berg- Kirkparik.....ck. Efficient nearest neighbor language models, 2021. htthstt:::// arxiv. wwwpxivv.org/aaoo http:/211...2211
Ruining He and Julian McAuley, Ups and downs: Modeling the visual evolution of fashion trends with one-class collaborative filtering. In proceedings of the 25th international conference on world wide web, pp. 507-517, 2016.
Ari Holtzman, Peter West, Vered Shwartz, Yejin Choi, and Luke Zettlemoyer. Surface form compe- tition: Why the highest probability answer isn't always right. arXiv preprint arXiv:2104.08315. 2021.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.
Minqing Hu and Bing Liu. Mining and summarizing customer reviews. In Proceedings of the Tenth ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD '04, pp. 167, New York, NY, USA, 2004. Association for Computing Machinery. doi: 10.1145/1014055.10000 1014033. URL http://///cc..o11111111515
ISBN 1581 138881.
101 4052. 1014073
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: Few-shot learning with retrieval augmented language models, 2022. URL httttttt//aaxtttti wtwpaaawwww..raa//////a000 03299
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023. URL https:
//axxiv.org/abs//311
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv: 1705.v:70500355/ 2017.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.

---

Published as a conference paper at ICLR 2026
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models, 2020. URL https://arxix:
org/abs/ 1911. 00172
Johannes Kiesel, Maria Mestre, Rishabh Shukla, Emmanuel Vincent, Payam Adinch, David Cor- ney, Benno Stein, and Martin Potthast. SemEval-2019 task 4: Hyperpartisan news detection. In Jonathan May, Ekaterina Shutova, Aurelie Herbelot, Xiaodan Zhu, Marianna Apidianaki, and Saif M. Mohammad (eds.), Proceedings of the 13th International Workshop on Semantic Evalua- tion, pp. 829-839, Minneapolis, Minnesota, USA, June 2019, Association for Computational Lin- guistics. doi: 1/S56///1//1122455 ottp:///ss////////aaaaaaoaooooo/
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453 2019.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. Remmened generation for geowledgeintensiing nlp tasks, 2021. URL https: 11401
Junyi Li, Xiaoxue Cheng, Wayne Xin Zhao, Jian- Yun Nie, and Ji-Rong Wen. Halueval: A large-
scale hallucination evaluation benchmark for large language models, 2023. URL https://
arxiv. wtt.:orgg.b.......r.r
Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human wwwp/rgg//////eeeee..eeeeeeeeeee 07958
falsehoods, 2022. http::ttt///att
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024.
Hanxiao Liu, Zihang Dai, David R. So, and Quoc Le. Pay attention to mlps, 2021. URL https: 7/arxiv. httpg//bs/2100 08050
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. Advances in neural information processing systems, 06:34892-34916 2023.
Andrew I.. Maas, Raymond F. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Dekang Lin, Yuji Matsumoto, and Rada Mihalcea (eds.), Proceedings of the 49th Annual Meeting of the Association for Compu- tational Linguistics: Human Language Technologies, pp. 142-150, Portland, Oregon, USA, June org/
2011. Association for Computational Linguistics. URL https::/aclanthhaaet
P11- 211-1
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016. URL http:p/://axxxv..oo/a/bbso///// 0atp///11...08833333 07843
OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Floren- cia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, and Shyamal Anadkat. Gpt-4 technical report, 2024. URL httpttttp:/aaait http:/awwww......//a/.a////////e http://bbb/23330877.2
Bo Pang and Lillian Lee. Seeing stars: exploiting class relationships for sentiment categorization with respect to rating scales. In Proceedings of the 43rd Annual Meeting Association for Com- putational Linguistics, ACL '05 pp. 115-124, USA, 2005a. Association for Computational Lin- guisties. doi: 19555
1219855 URL https: / http:////.....ro/...11111//////.
Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categorization with respect rating scales. In Proceedings of the ACI., 2005b.
12

---

Published as a conference paper at ICLR 2026
Baolin Peng, Michel Galley, Pengcheng He, Hao Cheng, Yujia Xie, Yu Hu, Qiuyuan Huang, Lars Liden, Zhou Yu, Weizhu Chen, and Jianfeng Gao. Check your facts and try again: Improving large language models with external knowledge and automated feedback, 2023. URL https:
12813
Hongjin Qian. Zheng Liu, Peitian Zhang, Kelong Mao, Defu Lian, Zhicheng Dou, and Tiejun Huang. Memorag: Boosting long context processing with global memory-enhanced retrieval augmentation. In Proceedings of the ACM on Web Conference 2025, pp. 1366-2377, 2025.
Qwen, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang. Jiaxi Yang. Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tianyi Tang, Tingyu Xia. Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yu Wan, Yugiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu. Qwen2.5 technical report, 2025. URL https: arxiv. org/ abs /es.s/412... 15115
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.
Weijia Shi, Julian Michael, Suchin Gururangan, and Luke Zettlemoyer. Nearest neighbor zero- shot inference. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. Abu Dhabi, United Arab Emirates, December 2022a. Association for Computational Linguis- URLhttps: / aac antho logy. org 2022. haa.aaacaaaaooocc
tics. doi: 10.
emn p-ma ma in. 2147
Weijia Shi, Julian Michael, Suchin Gururangan, and Luke Zettlemoyer. Nearest neighbor zero-shot inference. In Proceedings the 2022 Conference Empirical Methods in Natural Language Processing, pp. 3254-3265, 2022b.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng. and Christopher Potts. Recursive deep models for semantic compositionality over sentiment tree- bank. In David Yarowsky, Timothy Baldwin, Anna Korhonen, Karen Livescu, and Steven Bethard (eds.), Proceedings of the 2013 Conference Empirical Methods in Natural Language Process- ing. pp. 1631-1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics. URL http:ttt///aaaattll.l.
Weihang Su, Yichen Tang, Qingyao Ai, Junxi Yan, Changyue Wang, Hongning Wang. Ziyi Ye, Yujia Zhou, and Yiqun Liu. Parametric retrieval augmented generation, 2025. URL https:
//axxiv.rr//abll aaaaaab///000.0 1
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2330.009288 2023.
Tim Van Erven and Peter Harremos. Rényi divergence and kullback-leibler divergence. JEEE Transactions on Information Theory, 60(7):3797-3820, 2014.
Weizhi Wang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Aug- menting language models with long-term memory. Advances in Neural Information Processing Systems, 16:74530-74544 2023.
Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks, 2015. URL https://
erxiv.org/abs//147 3916
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv preprint 1rXiv: (80900900000,.....0 2018.
Ping Yu, Mikel Artetxe, Myle Ott, Sam Shleifer, Hongyu Gong, Ves Stoyanov, and Xian Li. Ef- ficient language modeling with sparse all-mlp. 2022. URL ptttps://arxx org/ abs/ 2203. 06850
13

---

Published as a conference paper at ICLR 2026
Lingxi Zhang, Yue Yu, Kuan Wang, and Chao Zhang. Arl2: Aligning retrievers for black-box large language models via self-guided adaptive relevance labeling. arXiv preprint 2XXi::2402.1354,,, 2024.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015a.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015b.
Zhilu Zhang and Mert Sabuncu. Generalized cross entropy loss for training deep neural networks with noisy labels. Advances in neural information processing systems, 31, 2018
14

---

Published as a conference paper at ICLR 2026
A THE USE OF LARGE LANGUAGE MODELS
We utilized ChatGPT as an editorial assistant to enhance the manuscript's language quality, correct grammatical errors, and ensure academic writing standards. All scientific contributions, including the research design, experimental procedures, and analytical interpretations, were originated and performed independently by the authors.
B REPRODUCIBILITY STATEMENTS
We provide comprehensive implementation details of method throughout the paper. Our method is thoroughly described in Section 3 We present detailed settings for our experiments and analyses
in Section 4 5.4 and Appendix
C
Code and checkpoints will be released upon acceptance.
C IMPLEMENTATION DETAILS
we utilize heterogeneous corpus constructed by aggregating several publicly available sources that cover diverse domains relevant to common NLP tasks. Following the methodology from (Geng et al.) 2024), this corpus comprises WikiTText- 103 MMrrty et al. (2016) for encyclopedic content, Amazon Reviews He & McAuley (2016) for user-generated product feedback, CC-NEWS Hamborg et al. (2017) for journalistic content, and IMDB Maas et al. (2011) for movie reviews and discussions.
Datasets For the general NLP tasks in Section 5.2
This diverse mixture captures both formal and informal language patterns, spans multiple domains from news articles to consumer opinions, and provides comprchensive coverage of linguistic phe- nomena encounnteed in real-world NLP applications. The complete dataset publicly available at: -ataattrtssssss
https://huggingfaeee wtt.atasetaaaaaaaee /ttp///r...ettt///ee.eggnaa///aa
Evaluation Metrics For question answering benchmarks, following Cheng et al. (2024), we eval- uate three Open Domain Question Answering datasets and HotpotQA using the Exact Match (EM) metric. For long-form QA evaluation, we employ three complementary metries: MC1, which mea- sures whether the model assigns the highest likelihood the most accurate answer; MC2, which sums the normalized probabilities over all correct answers; and MC3, which evaluates whether the model assigns higher average likelihood to true answers than false ones. We report the aver- age of these three metrics as the final performance measure for long-form QA tasks. For general NLP tasks, following the methodology from (Shi et al. (2022b), we report results using the domain- conditional PMI scoring rule Holtzman et al. (2021). For hallucination reduction evaluation, accuracy as the primary metric to assess the model's 's ability to generate factually correct responses.
Hyperparameters In Table we list the hyperparameters used for training the IB MLP Memory model (excluding embedding parameters).
Table 4: Hyperparameters for training the IB MLP Memory model.
Hyperparameter
Assignment
optimizer learning rate Ir scheduler type warmup ratio weight decay epochs
AdamW
4e-4
linear
0.03
0.0
5
flash attention batch size
False
4
gradient accumulation steps
4
num GPUs max train samples
32
33,000,000
15

---

Published as a conference paper at ICLR 2026
D SENSITIVITY TO INTERPOLATION WEIGHT X
We conducted comprehensive analysis of method's sensitivity to the interpolation weight on the HaluEval benchmark using Mistral-7B--0.33 T S presents the results across three tasks: Dialogue, QA, and Summarization, with À values ranging from 0.1 to 0.9.
Our findings demonstrate that the method exhibits robust performance across a wide range of À val- ues, with optimal performance generally observed in the 0.35-0.55 range. Specifically, the Dialogue task achieves its best performance at X 0.45 (64.07%) QA peaks at À = 0.55 (66. 86%), and Summarization reaches its maximum at 0.35 (5).41%). Notably, all three tasks show consistent improvements over the baseline Mistral-7B-vdl model across the optimal range, with QA showing the most substantial gains (up to 10.08 points improvement).
The performance remains relatively stable within the 0.3-0.6 range. with only gradual degradation outside this interval. When approaches extreme values (e.g. 0.9), performance deteriorates signif- icantly, particularly for Dialogue and Summarization tasks, though still maintaining improvements over the baseline in the QA task.
These results confirm that our method is not overly sensitive to the specific choice of À within reasonable range, making it practical for deployment without extensive hyperparameter tuning. The consistent improvements across different values and tasks validate the robustness of our approach.
Table 5: Performance sensitivity analysis of interpolation weight À on HaluEval benchmark us- ing MB-0ral-7B-v0.3, Results are reported accuracy (%) across three tasks: Dialogue, QA, and Summarization. The first row shows baseline Mistral-7B-v0.3 performance without memory aug- mentation. Bold values indicate the best performance for each task.
Dialogue
QA
Summarization
Mistral-7B-vO.
57.18
53.99
50.27
0.10
56.80
59.86
50.92
0.20
59.43
62.01
51.87
0.30
61.99
64.11
52.17
0.35
63.01
64.96
52.41
0.40
63.88
66.03
52.11
0.45
64.07
66.55
51.55
0.50
63.57
66.57
51.39
0.55
63.24
66.86
50.73
0.60
62.38
66.30
49.79
0.70
59.65
64.77
47.72
0.80
56,42
62.78
46.53
0.90
49.71
60.36
46.67
E INFERENCE EFFICIENCY ANALYSIS
Table 6 presents the computational cost breakdown for both Transformer and MLP architectures in terms of FLOPs per token. As demonstrated, the primary difference in computational efficiency stems from the absence of attention mechanisms in pure MLP models.
By comparing these computational requirements, we derive the theoretical speed ratio between the Transformer (denoted as FLOPss) and the MLP models (denoted as FLOPsm):
FLOPst
AMtayrrdodde dmodel (2dattn + drs) + Rea
Bete
FLOPsm
Groderddoddelyy
12dmode!
(8)
with the slandard dutta d dmodel-
This relationship in Equation [ reveals that MLPs maintain a consistent computational advantage across all context lengths, with the efficiency gap widening context length increases.
16

---

Published as a conference paper at ICLR 2026
Table 6: Flops per Token inference time. Following Kaplan et al. (2020), we analyze com- putational requirements for Transformer and MLP architectures where Mayer (number of layers), dmodel (dimension of the residual stream), a adimeennso of the intermediate feed-forward layer), dattn (dimension of the attention output) Rheads (number of attention heads per layer), nett (the length of input context), Trroes size). Forword denotes computational cost per token inference step.
Openration
FLOPs per Token(Transformen)
FLOPs per Token(MLP)
Embed
Admodel
Attention: QKV
Mayger 3dattta
Attention: Mask
Mayer Decterreetalle Deta
Attention: Project Feedforward
21Maaeer dattn moddet
Mnyyer 2omodedd
3 2dmoddddy
De-embed
2dmoder Wooons Boooab
2mmodded Pl Bryonh
Total(Non-Emmeedding)
C forward
Aayaer mooda (2daten drs) + M Rety Outtn
Corward Gralaeer moddad) modeld
F SCALING LAW
Setup We conduct scaling law experiments using standard decoder-only models and our overall model architecture. As basclines, we use four GPT-2 Radford et al. (2019) variants with increasing parameter counts: GPT2-small (124M), GPT2-medium (345M), GPT2-large (774M). and GPT2- xl (1.5B). For MLP Memory, we define three configurations: small (I (124M), medium (335M), and large (774M) that align with the scaling trend of standard architectures. The MLP Memory module is externally integrated with matching-sized GPT-2 variant, resulting in total parameter counts of approximately 248M, 7 710M, and 1.5B for the small, medium, and large configurations. respectively. All models are trained two datasets: WikiiTex-1103 MMrrty et al. (2016) (around 100M tokens) and mixed Web dataset (around 600M tokens). Following Shi et al. 2022a). our Web dataset com- bines diverse knowledge sources relevant to common NLP tasks, including WikiText- 103, Amazon
Reviews He & McAuley 2016), CC-NEWS Hamborg et al. (017). and IMDB Maas et al. (001)).
GPT2 Coeffakk
GPT2 Coefaai
GPTT- GITT2 MML
GPTT MLL
MLPMin
7.8
11 HPE ecl be
7.5
Pappocty
Parpessty
Poplauty
10
11-
1.2
1.0
0.1
0.5
2.0
1.0
10
20
30
40
Model Paramweers (Billion)
Model Partmeters (Billion)
Triniig CompateFarpp1100
(a) Scaling law on 1ik
(b) Scaling law on Web
(c) Training compute Scaling Law
Figure 6: Power-law scaling behavior with model size N and training compute C. (a) Scaling re- sults compare the continued training of GPT2 (GPT2-ConTrain) with our overall model architecture (GPT2+MLP Mem) under fixed compute. Our fitted curve shows 17.5% exponent improvement on WikiText-103 103. (b) On the larger Web dataset, our architecture exhibits stronger scaling gains from increased data size, with an exponent improvement of 24.1% (c) At the GPT2-xl scale, our architecture continues to benefit from additional training on the Web dataset without overfitting.
Scaling law with model parameters N Following Kaplan et al. 2020). we model perplexity scaling PPL = (8-N). Under fixed compute, we compare our architecture continued GPT-2 training on WikiText- 103 and Web datasets in terms of test perplexity scaling with model size N. Results in sig our architecture demonstrates steeper scaling curve than the decoder-only model, indicating improved scaling efficiency. The power-law scaling laws on WikiText-103 can be expressed as:
PPL (9.3 10 (0-8N)-00 -0.143
-0.143
and PPL (1.1 10 10011110000 -N00600 0.168
17

---

Published as a conference paper at ICLR 2026
where PPLa and PPLm denote the test perplexity of the decoder model and our architecture, respectively. The corresponding power-law scaling laws on the Web dataset are as follows:
PPL (7.7 10- and PPLm (5.0 10 1N) -0.26
(10)
These results highlight the superior scaling efficiency of our overall model architecture compared to the standard decoder-only baseline, on both WikiText-103 and the Web dataset.
Scaling law with training compute C We further examine how model performance scales with training compute C' while keeping model size fixed. At the GPT2-xl scale, we conduct experiments on the Web dataset, measuring test perplexity after varying amounts of training flops. As illustrated in Figure our overall model architecture achieves significantly lower perplexity with increasing training compute, with no signs of overfitting. This suggests that the retriever imitation pretraining task is more challenging and continues to benefit from additional compute.
G COMPARING OUTPUT DISTRIBUTION CHARACTERISTICS OF LM, KNN, AND MLP MEMORY
As two samples illustrated in Figure distributions produced by LM, ANN search, and MLP Mem- ory exhibit distinct characteristics. LM typically yields smooth and dense probability distributions, as it is trained generalize across large corpora and capture broad contextual patterns.
In contrast, the kNN approach produces sparse and spiky distributions, concentrating most of the probability mass only few retrieved neighbors. For instance, when using a GPT-2 model (vo- cabulary size 50,257), even after retrieving the top-k neighbors (e.g., k = 1024), only a small subset of these neighbors meaningfully influences the output distribution, while the majority receive near- zero probability.
The MLP Memory lies between LMs and kNN in terms of distribution characteristics. As a neural model, it is trained using a combination of KL loss and CE loss to approximate the sparse and spiky distributions produced by the kNN approach. While its outputs remain somewhat smoother due to the training objective, the resulting distributions are sharper than those of standard LMs, yet not as extreme as the highly concentrated outputs of kNN.
0.5
Probability Comparison: LM vs KNN vs MLP
0.5
Probability Comparison: LM vs kNN vs MLP
LM
LM
KNN
0.4
KNN
MLP
MLP
Probability
0.3
Probability
0.3
0.2
0.2
0.1
0.1
0.0
2
4
10
11
12
13
14
15
5
6
7
8
10
11
12
13
14
15
16
Sample
Rank
(Top
Top
16)
Sample 2: Rank (Top
Top
16)
Figure 7: Comparison of output probability distributions. Two samples show the top- probabilities from the LM and kNN distributions using GPT2-large, along with the distribution generated by the MLP Memory based the same large model size.
Table compares the output sparsity of LM, ANN, and MLP Memory by reporting the number of tokens assigned non-zero probabilities at various thresholds. The LM assigns non-zero mass to all 50,257 tokens, reflecting its dense distribution. However, this number drops sharply at higher thresholds, with only 2 tokens receiving probabilities above 0.1, indicating a rapid decay despite its broad support.
18

---

Published as a conference paper at ICLR 2026
Table Number of tokens with non-zero probability mass different thresholds. This table reports the number of tokens assigned non-zero probabilities by the LM, ANN, and MLP Memory, across range of probability thresholds. All values are averaged over 20,000 test samples.
Types
>
0.0
10-6
>
10-6
10
> 10-3
>
10
2
10-1
LM
50257
1760
562
148
34
7
2
KNN
251
217
197
146
43
2
MLP
50257
1151
388
115
30
2
In contrast, the KNN output is highly sparse, with only 251 tokens assigned any non-zero probability. Even at low thresholds (e.g., 10 10-"), the number remains limited, confirming its concentrated nature shaped by a small set of retrieved neighbors.
MLP Memory exhibits intermediate hehavior. Although it outputs the full vocabulary like the LMM, the number of tokens exceeding higher thresholds aligns more closely with &NN. This suggests that MLP Memory learns to approximate the spiky distributions of kNN while maintaining some smoothness from its parametric formulation.
Table 8: Cumulative token count required to reach probability mass thresholds. This table indicates the number of top-ranked tokens needed accumulate total probability mass exceeding thresholds such as 0.8, 0.9, etc. All values are averaged over 20,000 test samples.
Types
Top Prob Count(sum 0.8)
sum 0.9
sum
0.95
sum > 0.99
LM
23
63
142
617
ANN
22
43
68
126
MLP
13
33
72
308
Table further examines distribution sharpness by reporting the number of top-ranked tokens needed to accumulate a specified proportion of total probability mass. Here, we observe that the &NN distribution reaches 99% cumulative probability with only 126 tokens, while the language model (LM) requires 617 tokens to achieve the same threshold. This suggests that the LM's probability mass is more broadly spread across the vocabulary, in contrast the highly concentrated outputs of KNN.
Interestingly, the MLP Memory achieves 99% cumulative probability with 308 tokens, placing it between L.M and ANN. Notably, the MLP reaches 80% total probability with only 13 tokens tokens-fewer fewer than both LM and kNN-indicati support the observation that MLP Memory produces sharper distributions than LMs, yet avoids the extreme sparsity of &NN.
that captures prominent signals more efficiently. These results
H EFFECT OF DIFFERENT KNN TARGET DISTRIBUTIONS
Figure presents the test perplexity of our overall model architecture evaluated at various training steps. In all settings, both the base language model and the MLP Memory are of small size (GPT2- small), with the MLP Memory trained to mimic ANN target distributions constructed from different base models: GPT2-small, GPT2-medium, GPT2-large, and GPT2-xl. As training progresses, test perplexity steadily declines across all variants, indicating stable optimization and effective learning. Among them, the model trained on the kNN-XX XL distribution achieves the lowest final test perplexity (12.84), closely followed by the trained on ANN-large (12.85). In contrast, the models trained on --N-mmm and smal converge to higher perplexities approximately 12.87 and 12.91, respectively.
19

---

Published as a conference paper at ICLR 2026
These results demonstrate that ANN target distributions derived from larger base models lead to improved performance when used to train the MLP Memory. The richer and more informative su- pervision encoded in these distributions appears to enhance the parametric memory's generalization ability.
13.8
kNN-small
keN-mm
KNN-large
kNN-xl
Perplexity
20k
Training Steps
Figure 8: Test perplexity of overall model architecture, where both the base language model and the MLP Memory are of small size (GPT2-small). The MLP Memory is trained to mimic different KNN target distributions constructed from various base models: &NN-small (GPT2-small), ANN- med (GPT2-medium), kNN-large (GPT2-large), and ANN-XL (GPT2-xI).
SENSITIVITY TO IN TARGET DISTRIBUTION GENERATION
We used & = 1024 for generating all target distributions. Table shows the sensitivity analysis using GPT2-large-CPT, namely GPrg-large with continued pre-training on WikiText-103. While smaller & values degrade performance, values beyond 1024 yield minimal gains while significantly increasing computational costs, making & 1024 optimal for practical deployment.
Table 9: Test perplexity sensitivity to different values of k in target distribution generation using GPT2-large-CPT on Wik
Models
Perplexity k
GPT2-large-CPP
10.43
10.30
10.11
9.95
9,83
9,71
+&NN-LM
32
9,63
9.57
128
9.52
256
9,48
512
9.46
1024
9,43
2048
9.42
20

---

Published as a conference paper at ICLR 2026
RAG
Refer to the background document and answer the question:
SD-Amt, became official security organization the Nazi Party 1934. Consisting at first paid agents and a few hundred unpaid informants scattered across Germany, the SD was quickly professionalized under Heydrich, who commissioned National Socialist academics and lawyers ensure that framework of National Socialist ideology."" Heydrich was given the power select men for the SS Security Service from among any SS subdivisions since Himmler considered the organization of the SD as important. September 1939, the SD was divided into two departments. the interior department (Inland-SD) and the foreign department (Ausland-SD), and placed under the authority of the Reich Security Main Office (RSHA).
Background: SS Security Service, known
and its Security Service in particular, operated \"within the
Question: In WWII, who was the head of the Nazi party's security service? The answer is:
Output: Heinrich Himm
Base LM
MLP Memory
Answer the question:
Answer the question:
Question: In WWII, who was the head of the Nazi party's security service? The answer is:
Question: WIII, who the head of the Nazi party's security service? The answer is:
Output: Question: In WW
Output: Reinhard Heydrich.
Figure 9: Comparison of model outputs on factual question. Despite retrieving relevant documents with correct information (highlighted in green), RAG is misled by contextual distractors and pro- duces an incorrect answer. MLP Memory generates the correct answer without explicit retrieval.
RAG
Refer the background document and the question:
Background: 1. The Theory of Flight (1998); Mystics (2003): Harry Potter and the Order of the Phoenix
(2007): Harry Potter and the Half-Blood Prince (2009): Harry Potter and the Deathly Hallows Part 1 (2010);
Harry Potter and the Deathly Hallows Part 2 (2011); The Company You Keep (2012): About Time (2013);
Ex Machina (2014): The Legend of Tarzan (2016) Fantastic Beasts and Where to Find Them (2016):
Fantastic Beasts: The Crimes of Grindelwald (2018): Fantastic Beasts: The Secrets of Dumbledore (2022)
Question:
what are the names of harry potter movies arder? The answer is:
Output: 1. Harry Potter and the Philosopher's Stone
Base LM
MLP Memory
Answer the question:
Answer the question:
Question: what are the names of harry potter movies in order?
Question: what the names of harry potter movies in order? The answer is:
The answer is:
Output:
Harry Potter and the Philosopher's Stone, Harry Potter and the Chamber of Secrets, Harry Potter and the Prisoner of Azkaban, Harry Potter and the Goblet of Fire, Harry Potter and the Order of the Phoenix, Harry Potter and the Half-Blood Prince, Harry Potter and the Deathly Hallows Part 1, Harry Potter and the Deathly Hallows Part 2.
Output: 1. Harry Potter and
Figure 10: RAG retrieves irrelevant documents that introduce interference, while MLP Memory demonstrates perfect performance.
J CASE STUDY ON DOWNSTREAM TASKS
As show in Figure we observe that RAG often fails even with relevant retrieved documents due to contextual noise interference. When documents contain related but distracting information, RAG's shallow integration cannot effectively filter these distractors, leading to incorrect answers. In con- trast, MLP Memory learns intelligent corpus compression during training, capturing richer contex- tual relationships that enable robust disambiguation without explicit retrieval.
21

---

Published as a conference paper at ICLR 2026
Refer the background document and answer the question:
Background: 1. The Nordic countries are a geographical and cultural region in Northern Europe and the North Atlantic, which includes the countries of Denmark, Finland, Iceland, Norway and Sweden, and the autonomous territories of the Faroe Islands and Greenland. Anthropology has diverse history the Nordic countries, tracing all the way back to the early nineteenth century with the establishmentof ethnographic museums.
Question: where are the nordic countries? The answer is:
Output: Nordic
X
Base LM
Memory
Answer the question:
Answer the question:
Question: where are the nordic countries? The answer
Question: where nordic countries? The answer is:
Output: Question:
X
Output 
 The Nordic countries are located in Northern
Europe. The Nordic countries are: Denmark,
Figure 11: RAG fails to extract answer despite retrieving relevant content, while MLP Memory provides accurate response
RAG
Refer to the background document and answer the question:
Background: 1. London
Question: what london city? answer
Output: 1. The London
Base LM
MLP Memory
Answer the question:
Answer the question:
Question:
what to visit in london city? The answer is:
Question: what to visit in london city? The answer
Output Big Ben 2. Buckingham Palace 3. Tower Bridge 4. Tower
London 5. London Eye 6. Westminster Abbey 7. St Paul's Cathedral Tate Modern National Gallery 10. British Museum 11. Natural
Output: London is city that
History Museum Science Museum 13. London Zoo Madame
Tussauds 15. London Dungeon 16
Figure 12:
RAG system distracted by retrieved content fails to answer the question, while MLP Memory provides comprehensive and accurate response
22

---

Published as a conference paper at ICLR 2026
K CASE STUDY ON THE DISTRIBUTION OF BASE LM, KNN-LM, AND MLP MEMORY
To further understand the mechanisms underlying the effectiveness of MLP Memory, analyze the token-level probability distributions produced by the base LM, KNN-LM, and MLP Memory. Our hypothesis is that kNN-based distributions are particularly effective at capturing long-tail knowl- edge, whereas LM distributions exhibit strong coherence on frequent function words. MLP Mem- ory appears to internalize both characteristics, acquiring long-tail information manner similar to non-parametric retrieval while preserving the stability of parametric models.
To examine this hypothesis, we perform case studies on examples drawn from WikiText- and report the probability assigned by each method to the target token (highlighted in bold) given its preceding context.
Table 10:
Token-level probability assignments for different methods on long-tail entities
( top block) ) and coherent function words bottom block
Long-tail Knowledge
Context (target token in bold)
Base LM
kNN-LM 0.74
MLP Memory 0.75
Southward. the Yongsan area, Keiser placed Brigadier General Joseph Bradley. Assistant Division Commander, in charge the 9th Infantry Regiment.
The song reached number ten in Mexico and number one on both the Billboard Latin Songs and Latin Pop Songs chart. Coherence
0.01
0.07
0.45
Context (target token in bold)
Base LM
KNN-LM 0.01
MLP Memory
As the threat of invasion was clearly felt in late 1941, an idea for a series of secret observation posts (first in Gibraltar and later in other places like Malta and Aden)...
0.65
0.44
Here the invasion force encountered the first French defences, con- sisting of camouflaged trenches and pillhoxes dug in along ridge.
0.45
0.06
0.53
As shown in Table 10, the case studies provide clear empirical support for our hypothesis. For rare entities such Bradley and Mexico, MLP Memory assigns probabilities comparable to exceeding those of KNN-LM, demonstrating effective acquisition of long-tail knowledge. In contrast, for func- tion words such in and and, MLP Memory maintains probability mass close to that of the base LM, whereas kNN-LM shows substantial degradation. These observations suggest that MLP Mem- ory successfully combines the advantages of non- non-ponnne parameerrii retrieval with the coherence properties of parametric language models.
L MLP ARCHITECTURE DETAILS
The MLP Memory used in LLaMA2- and Mistral-based experiments is initialized from the corre- sponding MLP modules of their original architectures maintain structural consistency and stable training behavior. Both settings employ 8 stacked MLP layers with their respective native hidden and intermediate dimensions. The total size of the MLP Memory about 1B parameters, excluding embedding parameters. The detailed architectural configurations are reported in Table I1
Table 11: MLP Memory architecture for LLaMA2 and Mistral experiments.
Layers
Hidden dim
Intermediate dim
8
4096
1 11008
8
4096
14336
In the study of the effect of MLP Memory size presented in Appendix M the number of layers is varied while the hidden and intermediate dimensions are fixed at 1280 and 5 5120, respectively. The corresponding configurations summarized in Table
23

---

Published as a conference paper at ICLR 2026
Table 12: MLP Memory size ablation configurations.
Layers
Hidden dim
Intermediate dim
8
1280
5120
15
1280
5120
36
1280
5120
Based on the MLP Memory size study, adopt the 8-layer configuration the default setting, as it provides favorable balance between performance and efficiency.
M EFFECT OF DIFFERENT MLP MEMORY SIZES ON PERFORMANCE
This section examines how different MLP Memory sizes affect language modeling performance on WikiText-103. The memory capacity is controlled by varying the number of MLP layers while keeping other architectural settings unchanged.
Table 13: Performance of different MLP Memory sizes on WikiText- 103. The base model is GPT2- large.
Model
WikiText-103 PPL
GPT2-large
15.81
+ MLP Memory (8 layers, 221M)
11.41
MLP Memory (15 layers. 359M)
11.35
MLP Memory (36 layers, 772M)
11.25
As shown in Table even the smallest MLP Memory achieves significant improvements over the base model. While scaling provides additional gains, smaller models offer the best performance- efficiency trade-off.
24


Reviews:
[{"review_id": "FIwHZMIegO", "reviewer": "Reviewer_xCCs", "summary": "This paper proposes MLP-memory that replaces the KV store in kNN-LM by an MLP. The goal is to save memory space requirement, which is high in kNN-LM. The problem is interesting. A more efficient implementation of kNN-LM idea can make the approach more feasible in practice.\nThe method is tested on several tasks of QA and general NLP, showing better performance compared to the baselines. Analyses are provided on the impact of different hyperparameters and the training objective function.", "questions": "see weakness.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The main idea is to make kNN-LM more storage efficient. It trains an MLP to replace the costly K-V storage. The idea is interesting.\n\n2. The paper describes well the approach. It contains good analyses about different hyperparameters and settings.\n\n3. The approach is compared to the relevant baselines (except the missing comparison mentioned in weakness).", "weaknesses": "1. The contribution of the paper is limited. It does not change the basic idea of kNN-LM, but tries to provide a better implementation of it.\n\n2. As the main contribution of MLP-memory is to reduce the storage cost of kNN-LM, the main comparison should be done with kNN-LM. However, in the main results on QA, this comparison is missing. The comparison is only done on general NLP tasks. This is insufficient.\n\n3. The goal of MLP-memory is to reproduce kNN-LM at lost storage cost. The key research question is whether MLP can approximate KV store. This question requires extensive comparison between them. In addition to the end tasks, it is also useful to compare the internal probability distribution of models (e.g. PPL). The paper contains some analysis about the distribution of probabilities of LM, kNN-LM and MLP-memory, but perplexity is only evaluated on MLP-memory. A comparison on perplexity may better reveal whether MLP-memory can approach kNN-LM.\n\n4. The analysis (in appendix) on the different distributions of probabilities in LM, kNN-LM and MLP-memory is interesting, showing that MLP-memory is between LM and kNN-LM. It is however unclear what the implication would be. So the authors suggest that a distribution between LM and kNN-LM sould be better? and why?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "erQEWDxvtA", "reviewer": "Reviewer_hdzS", "summary": "This paper introduces MLP Memory, a parametric memory module trained to imitate the behavior of retrieval-augmented models such as kNN-LM. The approach constructs a datastore from a training corpus and trains a small MLP to predict the kNN distribution given hidden representations from a pretrained language model. During inference, the predicted distribution is interpolated with the base model’s output probabilities, removing the need for explicit retrieval or datastore lookup. Experiments on question answering and factual recall tasks show that MLP Memory achieves comparable or better accuracy than kNN-LM and RAG while being more efficient in inference, as measured by time-to-first-token (TTFT) and tokens-per-second (TPS). Ablations explore the layer position for attaching the memory and the balance between cross-entropy and KL objectives.", "questions": "- Is MLP Memory pretrained separately for each task or domain, or can a single memory generalize across corpora?\n- Are the base language model weights frozen during memory pretraining, or do they also receive updates?\n- How large is the total training cost (in GPU-hours or FLOPs) relative to standard fine-tuning or RAG adaptation?\n- How sensitive are results to the choice of $\\lambda$, and how is it tuned for each task?\n- What is the scaling trend with different memory sizes? Does a smaller MLP maintain performance efficiency?\n- Could you discuss missing important related literature on (efficient) parametric memory architectures such as Product Key Memory [1, 2] and Entities as Experts [3, 4].\n\n[1] Lample et al., Large Memory Layers with Product Keys\n\n[2] Kim and Jung, Large Product Key Memory for Pretrained Language Models\n\n[3] Fevry et al., Entities as Experts: Sparse Memory Access with Entity Supervision\n\n[4] Verga et al., Facts as Experts: Adaptable and Interpretable Neural Memory over Symbolic Knowledge", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper addresses an important problem in retrieval-augmented modeling: maintaining factual recall benefits without external retrieval or large storage. MLP Memory provides a clean parametric alternative that approximates kNN-LM behavior through distribution imitation. The method is conceptually clear and empirically well evaluated, with strong results on factual QA and hallucination benchmarks. The use of TTFT and TPS as efficiency metrics is appropriate and provides a runtime-based comparison to retrieval baselines. Ablation studies on interpolation weighting, attachment layer depth, and loss balancing help clarify key design choices. Overall, the work presents a technically sound and well-executed exploration of parametric memory.", "weaknesses": "Despite strong results, several aspects limit the clarity and generality of the approach. The training cost is not clearly quantified: even if inference is faster, pretraining the MLP Memory requires generating kNN distributions and performing supervised training on them, which likely increases per-step computation and total cost. It is also unclear whether the backbone model is frozen or jointly updated, which affects the practical difficulty of integration. Moreover, the MLP Memory seems to require task-specific pretraining, since the datastore and supervision are derived from each corpus; this reduces its plug-and-play appeal compared to non-parametric retrieval systems like RAG. The method also depends on tuning the interpolation weight λ for each task, which introduces additional hyperparameter sensitivity. Finally, while the ablation identifies an optimal attachment depth (~70%), it remains uncertain how this generalizes across different architectures or scales. These issues make the approach less universally efficient than implied.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M5mk2k4Py7", "reviewer": "Reviewer_qhzi", "summary": "This paper introduces a new method called MLP Memory to help large language models be more factually accurate without slowing them down.  The proposed solution involves pre-training a lightweight, all-MLP (Multi-Layer Perceptron) module to imitate the output distribution of a kNN-LM retriever over the entire pre-training corpus. This MLP Memory module learns to map a hidden state from the LLM to a vocabulary distribution that approximates the retriever's output, effectively internalizing retrieval patterns. During inference, the MLP Memory's distribution is interpolated with the base LLM's output, providing a knowledge-augmented prediction. The results show that this method makes models better at answering questions, reduces hallucinations, and is much faster than using a traditional search (RAG).", "questions": "(1) Cost problem: What is the total computational budget (e.g., in GPU hours) required to pre-train the 1B-parameter MLP Memory for a 7B model? How does this cost scale with the size of the base model and the pre-training corpus?\n\n(2) Dynamic knowledge integration: Could the MLP Memory architecture be adapted for incremental learning? For instance, is it feasible to fine-tune the MLP Memory on a small corpus of new documents to update its knowledge without catastrophic interference?\n\n(3) Generalization to specialized domains: The experiments use general corpora like Wikipedia. How would the method perform when applied to highly specialized domains (e.g., legal or biomedical text) where the retrieval might be more complex? For example, in the legal domain, retrieving similar legal cases for generative legal judgment prediction is much difficult than the general RAG systems.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "(1) The motivation of the paper is clear. The paper identifies a clear gap in the literature and proposes an elegant hybrid approach that bridges parametric and non-parametric methods. \n\n(2) Comprehensive evaluation: The experimental design is thorough, evaluating the method across multiple critical dimensions: factual question answering, general NLP capability preservation, hallucination reduction, and inference efficiency. \n\n(3) Compelling performance: The results show that MLP Memory not only surpasses the base models but also outperforms RAG on accuracy metrics. This suggests the method successfully captures richer contextual relationships. Results also show that MLP Memory gets better efficiency than existing methods.", "weaknesses": "(1)  The method requires a pre-training phase for the MLP Memory module, which involves constructing a large datastore and optimizing a large model. The computational cost and time of this initial step could be a barrier for some. \n\n(2)  Unlike RAG, the knowledge stored in the MLP Memory is fixed after pre-training (just like the generative retrieval models). This limits the model's ability to access updated or real-time information without a costly retraining cycle, making it less suitable for domains requiring frequent knowledge updates.\n\n(3) While the paper demonstrates that the method works, a deeper analysis of what specific knowledge patterns the MLP Memory learns and how it interfaces with the transformer's inherent reasoning process would strengthen the theoretical contribution.\n\n(4) Some existing works are trying to solve similar problems. For example, RetroLLM also tries to leverage generative retrieval models to implement end-to-end RAG systems. The authors should discuss and compare with these existing methods. \n\n[1] RetroLLM: Empowering Large Language Models to Retrieve Fine-grained Evidence within Generation, Li et al..", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YmUX2TJFZn", "reviewer": "Reviewer_t9hD", "summary": "The paper introduces MemoryMLP, a parametric module to mimic the behavior of explicit, external knowledge. In comparison to classical RAG approaches, MemoryMLP offers faster inference due to simple probability interpolation using the trained MLP module. The authors conduct extensive experiments showing that (1) MemoryMLP improves performance over base models and selected baselines (RAG / LoRA), (2) improves efficiency and throughput over these approaches, (3) reduces hallucination rates as measured by HaluEval, (4) ablations on the loss-function trade-off.", "questions": "- The dataset construction in the main part is quite short - can you please specify how you create the embedding-distribution pairs? Specifically, how do ensure (i) that you generate useful examples such as in Figure 4 that actually depict a fact and (ii) that you have pair following identical format, also shown in Figure 4?\n- Do you have any inights on performance differences using different MLP sizes? Why have you initially chosen 1B? How many layers have you chosen for the MLP?\n- Do you think the training budget is better spend on optimizing a smaller LM on the target distribution? Why, why not?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The proposed idea of MemoryMLP is novel, simple and intuitive. It addresses a key limitation of many current RAG framework, overcoming search time and inference overhead. The presentation of the method including the respective image (Figure 4) is precise and easy to follow.\n- The experiements show strong improvement over RAG and LoRA approaches, especially for Natural Questions and HotpotQA where the respective baseline do not improve on the base model performance. The choice of experiments is sound and convincing.\n- The ablations and experiments on downstream tasks other than QA show the transferability of the proposed approach.", "weaknesses": "- The presentation may be contain flaws: While the overall presentation is good, I have stuggles on understanding several details of the method, either because the information is not where I would expect it to be or because it is simply not mentioned. For instance, some hyperparameters appear to be arbitrarily chosen such as the parameter count of the MLP or why the Wikipedia-2021 dump. Further, the MLP is not further specified, e.g., how many layers. See also the questions below.\n- The results may be considered weak: While we observe relative improvements over the selected baselines, e.g., for NQ or HotpotQA are low and there are known to be existing approaches that are better. While state-of-the-art methods are not necessarily comparable, I would like to understand how MemoryMLP performs against other approaches such as question decomposition or fine-tuning a smaller LM. Especially considering the training budget of 5B token. See also question below.\n- The evaluation may be considered incomplete: Do you also fine-tune the RAG baseline or do you employ them in a zero-shot setup? As proposed by Lewis et al. (2022), their RAG pipeline is also end-to-end differentiable and it would be interesting to see a comparison between MemoryMLP and end-to-end RAG.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)