PaperlayerView exact revision on arXiv

Source-linked AI summary

From Lookups to Lessons: Turning Retrieved Knowledge into Parametric Memory

Changjiang Han, Bowei He, Ye Yuan, Weixu Zhang, Haolun Wu, Xue Liu, Jikun Kang

arXiv:2611.00450v1cs.CLcs.AIcs.LG

TL;DR

RAG provides evidence beyond an LLM’s prompt but adds storage and inference costs, while retriever-trained memories may not improve predictions when combined with the base model. ResMem distills retrieved evidence into a residual memory, improving five-dataset knowledge-QA results and reducing retained construction storage; its memory reads one hidden state and writes only to output logits.

  • Problem

    RAG requires corpus and index storage plus retrieval and evidence-reading latency, while retriever-trained memories do not train against the base prediction they later modify.

    Source: §1 · p. 1

  • Method

    ResMem uses a frozen base that reads retrieved evidence as teacher to train a parametric memory’s additive logit residual through the joint base-plus-memory prediction.

    Source: §1 · p. 2

  • Results

    Across five knowledge-QA datasets, ResMem improves EM and token F1 over the base on every dataset, raising mean F1 from 32.83 to 34.48 and breaking 21% fewer correct base answers than MLP Memory.

    Source: §1 · p. 2

  • Takeaways & Limitations

    On a Wikipedia subset, evidence-reading teacher targets cut retained construction storage by 94.7% (18.8×), and the reported extra repairs come with that teacher.

    Source: §1 · p. 2, §6 · p. 9

    The memory reads one hidden state and writes only to output logits, so the teacher’s gain is not integrated through the base model’s intermediate layers.

    Source: §6 · p. 9

Abstract

from arXiv · show

Retrieval-augmented generation (RAG) helps large language models (LLMs) answer factual questions beyond their training data, but it adds storage and inference costs. Distilling previously retrieved information into parametric memories enables reuse without online retrieval. However, training the memory to mimic a retriever rewards accurate independent predictions, which do not necessarily improve the final output quality when combined with the LLM's existing predictions. To address this, we introduce ResMem, in which the LLM reads retrieved text offline to generate teacher targets. This design avoids storing massive token representations, unlike existing parametric memory approaches such as MLP Memory, which relies on k-nearest neighbor (kNN) datastores. Furthermore, by using the LLM as an offline teacher that reads privileged retrieved evidence alongside the input, we train a parametric memory to learn an additive logit residual. This residual aligns the predictions of a student sharing the same backbone with the teacher distribution, while self-distillation updates remain confined to the parametric memory. On a Wikipedia subset, ResMem reduces the storage retained for teacher construction by 94.7% compared with the evaluated kNN pipeline. Across five question answering datasets, ResMem raises mean token F1 from 32.83% for the frozen LLM to 34.48%. It also breaks 21% fewer of the LLM's correct answers than a kNN-trained memory. These findings validate our motivation of treating retrieval as a learning resource, with parametric memory shaped by how retrieved knowledge is used to inform the model's future predictions. Code and model weights are available at https://github.com/SSRand/From-Lookups-to-Lessons and https://huggingface.co/RandHan/ResMem-Mistral-7B-v0.3.

1 INTRODUCTION

ResMem trains a parametric memory through the base prediction it will modify, using an evidence-reading teacher rather than retriever-only targets. It improves knowledge-QA results while reducing retained teacher-construction storage.

1 INTRODUCTION

  • Retriever-pretrained memories match the retriever without seeing the base prediction during training, so an accurate retriever imitation can still override a correct base answer.Their base model and memory interact only at inference.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • ResMem trains a logit residual through the joint base-plus-memory prediction, using a frozen base that reads retrieved evidence as teacher and a student that sees only the query.At inference, only the frozen base and trained memory are needed; γ controls the residual’s strength.Source: §1 · p. 2
  • 94.7% (18.8×) less retained construction storage is needed with ResMem’s BM25 retrieval and cached top-64 teacher targets than with the evaluated kNN pipeline.On a subset of English Wikipedia, retained storage falls from 602.0 to 32.0 GiB.Source: §1 · p. 2
  • Mean F1 rises from 32.83 to 34.48 across five knowledge-QA datasets, with ResMem improving EM and token F1 over the base on every dataset.It corrects 11% more base errors and breaks 21% fewer correct base answers than MLP Memory.Source: §1 · p. 2

2 RELATED WORK

Prior work stores retrievable evidence or trains parametric memories to imitate retriever distributions. ResMem instead trains an offline memory contribution through the combined base-plus-memory prediction.

2 RELATED WORK

  • Retrieval methods such as RAG, RETRO, and Atlas condition generation on retrieved text, while parametric memories such as MLP Memory and Memory Decoder imitate retriever distributions for retrieval-free inference.ResMem trains the memory contribution through the joint prediction instead.Source: §2 · p. 3
  • ResMem uses privileged retrieved evidence for self-distillation, with a shared frozen backbone and a student that lacks the evidence; TransMem instead reads long context and outputs gated hidden-state shifts.ResMem learns a logit residual through a loss on the combined prediction and uses evidence only offline.Source: §2 · p. 3
  • Unlike decoding-time methods that combine or contrast predictions, ResMem learns its residual through a combined-prediction loss and trains on fixed teacher-forced continuations.On-policy distillation methods instead use student-generated trajectories.Source: §2 · p. 3

3 METHODOLOGY

ResMem adds a trainable memory to a frozen language model and trains it through the combined base-plus-memory prediction. Retrieved evidence informs the teacher offline, while inference uses the learned memory without retrieval.

3.1 MEMORY INPUT AND ARCHITECTURE

  • The memory maps the frozen base model’s contextual representation to vocabulary logits using feed-forward blocks; only memory parameters are trained.Because the base representation carries preceding context, the memory needs no token-level attention.Source: §3.1 · p. 3

3.2 RESIDUAL SELF-DISTILLATION

  • Retrieved passages give the frozen teacher privileged input, while the student sees only the query and adds a memory residual to the base logits.Training compares their predictions and updates only the memory; the residual learns the log-ratio between the target and base prediction.Source: §3.2 · p. 4
  • Teacher targets cache renormalized top-64 token probabilities, and the loss combines teacher-to-student KL with next-token cross-entropy.The cached target has zero mass outside its top-K tokens, so the method uses forward KL.Source: §3.2 · p. 4

3.3 INFERENCE AND MATCHED CONTROL

  • At inference, γ scales the memory logits before they are added to the base logits; the additive readout follows a log-linear path, unlike the control’s linear probability mixture.In the stated override comparison, the additive rule requires a correspondingly confident target to override a confident base token; the mixture can override any base margin once λ > 1/2.Source: §3.3 · p. 5

3.4 CONSTRUCTION STORAGE

  • With d = 4096, P = 64, and K = 64, the kNN route retains 8652 bytes per position, while the RAG route retains 384 + b bytes.The kNN route retains datastore and index entries as well as cached targets; the RAG route retains retrieved text and a BM25 index instead.Source: §3.4 · pp. 5–6

4 EXPERIMENTAL SETUP

The experiments use Mistral-7B-v0.3 as the base model and compare ResMem with memory and retrieval baselines. Training and evaluation span Wikipedia-based memory construction and several QA and HaluEval datasets.

Models

  • Mistral-7B-v0.3 is the base model, and ResMem and MLP Memory each attach an eight-layer memory with approximately 1.54B parameters.Additional baselines include BM25-RAG, LoRA, and Memory Decoder.Source: §4 · p. 6

Training

  • On the Wikipedia subset, two memories use a kNN teacher and train for three epochs; knowledge-QA and other evaluations train the memories for one epoch on full English Wikipedia.The shared-teacher ablation trains eight paired seeds per objective and for Memory Decoder with equal exposure.Source: §4 · p. 6

Datasets

  • Knowledge QA comprises 38,627 questions from five datasets, and HaluEval evaluation uses its Dialogue, QA, and Summarization splits.Memory training uses the December 2021 English Wikipedia dump; the storage subset contains 74 of 4,716 source shards.Source: §4 · p. 6

5 RESULTS

Across knowledge-QA task families, ResMem leads on EM, F1, and PPL while reducing retained construction assets and improving the correction–harm balance over MLP Memory. Shared-teacher ablations show a preservation–repair trade-off, while additional repairs in the full recipe accompany the evidence-reading teacher.

5.1 KNOWLEDGE QA TASK PERFORMANCE

  • Across three task families, ResMem has the highest EM and F1 and lowest PPL; over Base, mean EM rises from 25.91 to 26.74 and mean F1 from 32.83 to 34.48.Source: §5.1 · p. 6, Table 1 · p. 7

5.2 CONSTRUCTION STORAGE

  • 602.047 vs 32.033 GiB: the RAG route reduces retained construction assets by 94.7% (18.8×) versus kNN on the Wikipedia subset.Across nested shards up to 652M positions, kNN files scale as (2d + 4)n bytes, while BM25 index and passages grow at 3.5–3.7 bytes per position.Source: §5.2 · pp. 6–7

5.3 CORRECTION AND HARM ANALYSIS

  • 704 corrections and 381 harms: ResMem nets 0.84 EM points versus MLP Memory’s 0.39, with most of the difference coming from 101 avoided harms.Corrections and harms are paired against Base across 38,627 questions; each method’s global coefficient is selected on disjoint development questions.Source: §5.3 · p. 7, Figure 3 · p. 8
  • At 87% of cached positions, kNN puts at least 95% of its mass on one token, versus 24% for the evidence-reading teacher and 22% for Base.The paper identifies target shape as one source of the harm gap.Source: §5.3 · p. 7, Figure 4 · p. 8
  • With the shared kNN teacher, standalone MLP Memory’s 95%-mass support shrinks from 21,051 tokens to 310, and its mixture component can override confident Base answers.MLP Memory breaks 482 correct answers, compared with 381 for ResMem.Source: §5.3 · p. 8, Figure 4 · p. 8

5.4 ABLATIONS

  • Under the shared BM25-RAG teacher, the residual configuration leads the standalone configuration by 0.117 F1 points, with a 95% interval of [−0.020, 0.258].Source: Table 2 · p. 9, §5.4 · p. 9
  • Memory Decoder reaches F1 32.87, 0.26 and 0.37 points below the standalone and residual configurations, respectively, under the tested training budget and model sizes.Development selection sets its coefficient to zero for six of eight seeds.Source: Table 2 · p. 9, §5.4 · p. 9
  • At native operating points, the residual configuration preserves 96.11% of Base-correct answers versus 94.05%, but repairs 1.51% of Base errors versus 2.25%.With the teacher fixed, the residual configuration thus trades some repair for preservation.Source: §5.4 · p. 9

6 CONCLUSION

ResMem distills a retrieval-augmented teacher into an additive logit residual on a frozen base model, reducing construction storage and improving knowledge-QA results. The conclusion identifies the memory’s single-hidden-state input and output-logit-only write as a limitation; the paper also reports its AI use and provides code and a checkpoint.

6 CONCLUSION

  • On five knowledge-QA datasets, ResMem raises mean F1 from 32.83 to 34.48, while correcting more base errors and harming fewer correct answers than MLP Memory.With a shared teacher, the residual objective mainly preserves correct answers; extra repairs come with the evidence-reading teacher.Source: §6 · p. 9
  • The memory reads one hidden state and writes only to output logits, so the teacher’s gain is not integrated through the base model’s intermediate layers.The authors identify memories that read from and write to several layers as a natural next step.Source: §6 · p. 9

AI USE STATEMENT

  • The authors report using generative AI for mathematical formalization, experiment feedback, code implementation, paper structure, and readability editing, while retaining research decisions and verification.They state they did not use it to generate synthetic data, propose research ideas or hypotheses, or clean datasets.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The ResMem training and main evaluation pipelines are available as code, alongside a knowledge-QA checkpoint.The passage provides the repository and checkpoint locations.Source: REPRODUCIBILITY STATEMENT · p. 10

A IMPLEMENTATION DETAILS

The implementation details specify where the Mistral memory reads, how the subset kNN teacher and Memory Decoder are configured, and the evaluation and hardware settings.

Memory input

  • In Mistral, the memory input h_t is the state after the attention residual connection and following RMS normalization; base logits use the remaining feed-forward computation, final normalization, and frozen vocabulary head.Source: A IMPLEMENTATION DETAILS · p. 14

Subset kNN teacher and Memory Decoder

  • The Wikipedia-subset kNN teacher reproduces the original position count and 26.967 GiB target cache, while Memory Decoder uses a Qwen3-0.6B-Base body with the Mistral vocabulary.The Memory Decoder has 0.47B parameters.Source: A IMPLEMENTATION DETAILS · p. 14

Knowledge-QA configuration

  • Knowledge-QA coefficients are selected by validation EM and reused for EM, F1, and PPL; BM25-RAG retrieves five Wikipedia passages, unlike the other methods’ shorter-input, shorter-output settings.BM25-RAG uses a 4096-token input and up to 64 output tokens; other methods use 384 input tokens and at most 12 output tokens.Source: A IMPLEMENTATION DETAILS · p. 14

Hardware

  • Memory training uses A100 GPUs, while deployment timing uses one GPU; QA inference uses BF16 base computation and greedy decoding.The passage also specifies batch size 16 and eager attention for QA inference.Source: A IMPLEMENTATION DETAILS · p. 14

B PER-DATASET KNOWLEDGE-QA RESULTS

Per-dataset results show ResMem improves EM and F1 over Base on every dataset, while MLP Memory edges it on WebQuestions.

B PER-DATASET KNOWLEDGE-QA RESULTS

  • ResMem improves EM and F1 over Base on every dataset and has the lowest PPL on each; MLP Memory is slightly higher on WebQuestions EM and F1.WebQuestions is identified as the smallest dataset.Source: B PER-DATASET KNOWLEDGE-QA RESULTS · p. 14

C HALLUCINATION DETECTION

Using the QA checkpoints without task-specific tuning, ResMem improves mean hallucination-detection accuracy over Base and MLP Memory.

C HALLUCINATION DETECTION

  • ResMem raises mean detection accuracy from 54.83 for Base to 63.09, compared with 60.34 for MLP Memory.The passage also reports that ResMem’s margin over MLP Memory is 7.00 points on QA and 1.48 on summarization, while MLP Memory is 0.22 points higher on dialogue.Source: C HALLUCINATION DETECTION · p. 14

D PRECISION AND DEPLOYMENT COST

BF16 loading produces small mean-F1 changes while reducing ResMem’s latency and additional peak memory; deployment measurements use fixed coefficients and specified hardware and decoding settings.

Precision

  • BF16 changes mean F1 by −0.0159 points for ResMem, +0.0227 for MLP Memory, and +0.0129 for Memory Decoder across 38,627 questions.The respective 95% paired question-cluster intervals are [−0.0539, 0.0195], [−0.0195, 0.0684], and [−0.0063, 0.0338].Source: D PRECISION AND DEPLOYMENT COST · p. 14

Deployment cost

  • At batch 16, BF16 ResMem reduces fixedlength total latency from 1,254 to 1,098 ms and raises decode throughput from 595 to 727 tokens/s.Its additional peak allocated memory over Base falls from 5.750 to 2.875 GiB.Source: D PRECISION AND DEPLOYMENT COST · p. 14
  • BF16 ResMem adds 8.6% latency over Base, or 16.3% at batch one, while BF16 Memory Decoder adds 63.5% latency over Base on the same host.Memory Decoder adds 1.01 GiB of peak memory in BF16 and has 0.47B parameters; its timing used a second identical host.Source: D PRECISION AND DEPLOYMENT COST · pp. 14, 16

Evaluation setup

  • Timing uses batch 16, 384 input tokens, and exactly 32 output tokens; quality covers 38,627 questions, with fixed coefficients and a BF16 Base.Times and peak allocated memory are medians over ten steady batches; Memory Decoder rows use a second identical host.Source: Table 5 · p. 16

E TOKENWISE GATE ABLATION

The tokenwise gate selectively applies ResMem’s residual when its score meets a threshold, but its measured QA changes are small and Table 1 reports the ungated model.

Gate design

  • The gate aims to retain residual changes that repair base errors while blocking changes that override correct base tokens.It acts only at decoding steps where the residual changes the base’s top token.Source: E TOKENWISE GATE ABLATION · p. 16
  • A logistic classifier uses 22 decoding-time features and applies fused logits when its score is at least 0.5, without labels or future tokens.The features include distribution statistics, top-token comparisons, step and prefix information, and whether either top token ends the answer.Source: E TOKENWISE GATE ABLATION · p. 16

Gate evaluation

  • On 38,627 questions, the gate changes mean EM by +0.07 points and mean F1 by −0.02 points; every dataset interval includes zero, so Table 1 omits it.It alters 190 answers, 113 for the better and 77 for the worse.Source: E TOKENWISE GATE ABLATION · p. 16, Table 6 · p. 16
  • The gate is fitted on 264 NQ/TriviaQA training questions and selected on a disjoint 4,091-question development split, adding QA supervision beyond memory training.The fused top token differs from the base at 19,254 of 463,520 decoding steps; the gate accepts 10,614 changes and blocks 8,640.Source: E TOKENWISE GATE ABLATION · p. 16

Related papers

Loading 2611.00450v1…