PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Plug-n-Play Knowledge Modules with Deep Context Distillation

Lucas Caccia, Alan Ansell, Edoardo Ponti, Ivan Vulić, Alessandro Sordoni

arXiv:2503.08727v4cs.LGcs.AI

TL;DR

Integrating new or rapidly evolving documents after pretraining remains difficult, especially with limited data, and in-context learning and RAG face limits in global information or inference cost. The paper trains plug-and-play Knowledge Modules with Deep Context Distillation, which outperforms other methods across evaluated datasets and settings and shows synergy with RAG.

  • Problem

    Post-pretraining integration of new, private, or specialized documents is challenging, while in-context learning and RAG face limits with long documents, global information, and inference cost.

    Source: §1 · p. 1

  • Method

    The paper encodes document knowledge in plug-and-play LoRA Knowledge Modules trained to match a document-conditioned teacher’s hidden states and output probabilities, using synthetic data.

    Source: §1 · p. 2

  • Results

    DCD-based KMs outperform other methods across QuALITY and NarrativeQA, two base models, and open- and closed-book settings, with reported synergies with RAG.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    Document-level KMs support fine-grained control over knowledge integration, and the reported results highlight their modularity and synergy with RAG.

    Source: §1 · p. 2, §5 · p. 10

    The study assumes the document relevant to each question is known; document-agnostic use across multiple KMs is left for future work.

    Source: §5 · p. 10

Abstract

from arXiv · show

Dynamically integrating new or rapidly evolving information after (Large) Language Model pre-training remains challenging, particularly in low-data scenarios or when dealing with private and specialized documents. In-context learning and retrieval-augmented generation (RAG) face limitations, including their high inference costs and their inability to capture global document information. In this paper, we propose a way of modularizing knowledge by training document-level Knowledge Modules (KMs). KMs are lightweight components implemented as parameter-efficient LoRA modules, which are trained to store information about new documents and can be easily plugged into models on demand. We show that next-token prediction performs poorly as the training objective for KMs. We instead propose Deep Context Distillation: we learn KMs parameters such as to simulate hidden states and logits of a teacher that takes the document in context. Our method outperforms standard next-token prediction and pre-instruction training techniques, across two datasets. Finally, we highlight synergies between KMs and RAG.

1 Introduction

Integrating new, private, or specialized documents after pretraining is difficult in low-data settings, while in-context learning and RAG have limits in handling very long documents, global information, or inference cost. The paper proposes document-level Knowledge Modules trained with Deep Context Distillation to encode document knowledge in plug-and-play adapters.

1 Introduction

  • In-context learning is strained by extremely long documents, while RAG supplies local passages at higher inference cost and may omit global document information.Continual pretraining on new documents is another alternative, but next-token prediction may be ineffective with limited data.Source: §1 · pp. 1–2
  • Document-level KMs are parameter-efficient LoRA modules that can be loaded on demand and managed individually, including for access control or removal without retraining.The paper proposes one module per document or document group.Source: §1 · p. 2
  • DCD trains each KM to reproduce the document-conditioned model’s hidden states and output probabilities, with synthetic data providing additional training signals.The paper explores summaries, question-answer pairs, and knowledge-graph data as inputs for eliciting the teacher’s behavior.Source: §1 · p. 2, Figure 1 · p. 2
  • The authors report that DCD pairs well with synthetic-data generation, with performance increasing as generation compute scales.Source: §1 · p. 2
  • Across QuALITY and NarrativeQA, Phi-3 3B and Llama-3.1 8B, and open- and closed-book settings, DCD-based KMs outperform other methods and show synergies with RAG.The evaluation spans two long-context question-answering datasets, two base models, and two evaluation settings.Source: §1 · p. 2

2 Knowledge Modules

Knowledge Modules encode document knowledge in plug-in parameters, while Deep Context Distillation trains them to reproduce a document-conditioned teacher. Knowledge Extractors add task-specific adaptation, with an extension that incorporates retrieved context.

2 Knowledge Modules

  • KMs store information about a document in a small parameter set that can be plugged into a base model on demand.They are parameter-efficient LoRA adapters, which add the outer product of two learnable low-rank matrices to each linear layer.Source: §2 · p. 3, §2.1 · p. 3

2.2 Next-Token Prediction Loss

  • Next-token prediction trains a document KM to predict subsequent document tokens, but can overfit perplexity without extracting all relevant knowledge.The paper motivates distillation as a richer learning signal for this low-data setting.Source: §2.2 · p. 3

2.3 Deep Context Distillation

  • DCD trains a student KM without document access to match a teacher that sees the document, using output probabilities and hidden states.The full teacher distribution supplies a rich signal, while hidden-state matching gives more direct credit assignment to LoRA layers.Source: §2.3 · p. 3
  • DDCD distills on contiguous chunks sampled from the document, using one chunk as context for the next, without augmented data.SDCD instead generates synthetic data from a sampled document chunk; explored inputs include question-answer pairs, summaries, and entity relationships.Source: §2.3 · p. 4

2.4 Task Adaptation with Knowledge Extractors

  • A task-trained Knowledge Extractor combines with document KMs to optimize task performance, and can generalize to unseen documents and their KMs.The modules are combined using learned weights for the KM and KE at each layer where LoRA is applied.Source: §2.4 · p. 4
  • The task-adaptation loss can also include retrieved passages conditioned on each query, extending the KM-and-KE setup with RAG context.The passage describes this as a straightforward extension of the loss.Source: §2.4 · p. 4

2.5 Training Considerations

  • KMs can be trained independently in parallel, whereas training a KE requires the set of KMs for joint multitask training.The setup assumes some query-time computation can move to a more powerful training-time index; KMs require gradient descent, and SDCD requires generated summaries.Source: §2.5 · p. 5

3 Experiments

Experiments compare Knowledge Module training objectives and combinations with retrieval across two question-answering datasets. They report gains from Deep Context Distillation, synthetic-data choices, and combining KMs with RAG.

Setup

  • Evaluation uses QuALITY and NarrativeQA with Phi-3 and Llama-3.1 in closed- and open-book settings; closed-book comparisons include LM, DCD, SDCD, and PIT.QuALITY documents average about 5,000 tokens, while NarrativeQA documents average about 60,000 tokens.Source: §3 · p. 5

Results

  • Across both datasets and base models, SDCD performs best among closed-book LM, PIT, and DDCD variants; training on Summary and QA together performs best.The comparison covers QuALITY and NarrativeQA with Phi-3 and Llama-3.1.Source: Table 1 · p. 6, §3 · p. 6
  • Compared with RAG + KE, RAG + KM + KE yields 4.2 and 4.1 Rouge-L gains on NarrativeQA and 2.4% and 4.1% Accuracy improvements on QuALITY, for Phi-3 and Llama 3.1, respectively.The authors attribute slight forgetting when comparing RAG with RAG + KM to KMs not being trained with long contexts, and report that longer concatenated training samples address this gap.Source: §3 · p. 7

Ablations

  • QuALITY ablations find gains from diverse synthetic inputs, stronger generators, more datapoints per chunk, and combined KL and L1 losses when no KE is used.A task-specific KE can mitigate the performance gap from adding hidden-state matching.Source: Table 2 · p. 7, §3 · pp. 7–8

Transfer of Knowledge Extractors

  • On LongHealth, a QuALITY-trained Knowledge Extractor transfers zero-shot to patient-document KMs and improves accuracy by 1.3% over ICL.Zero-shot KM+ICL alone does not add gains.Source: §3 · p. 7, Table 2 · p. 7

Inference cost

  • KM+KE beats RAG+KE at k=1 using 40 vs 200 prompt tokens; RAG+KM+KE reaches 42.4 Rouge-L at k=8 versus 42.5 at k=16, saving 50%.Figure 2 plots NarrativeQA Rouge-L against context-token cost.Source: §3 · p. 7, Figure 2 · p. 8

Ablations

  • A random KM reduces Llama 3.1 (8B) MMLU accuracy by 1%, from 63.9% to 62.9%.The paper characterizes this change as negligible.Source: §3 · p. 8

4 Related Work

The related work situates Knowledge Modules among retrieval, context distillation, and modular knowledge-injection approaches. It distinguishes this paper's per-document gradient-trained modules from prior prompt compression and multi-task module-generation approaches.

Knowledge injection and retrieval

  • RAG and in-context learning can incur storage costs or miss global document context; GraphRAG also uses costly offline graph construction to capture global context.The paper suggests GraphRAG-generated data could potentially serve as KM distillation data.Source: §4 · p. 8

Knowledge distillation

  • Unlike context-distillation work focused on prompt internalization, this paper uses distillation for document knowledge acquisition and reports transfer across unseen KMs.It also reports gains from synthetic data generated by a stronger, off-policy model.Source: §4 · p. 9

Knowledge injection with modules

  • Unlike multi-task hypernetwork approaches, KMs are trained by gradient descent on individual documents, increasing per-document cost while reducing domain-shift risk.The paper also distinguishes its fixed inference compute from document-length-dependent soft-prompt approaches.Source: §4 · p. 9

KV cache compression

  • Knowledge Modules are framed as query-agnostic context compression, alongside learned compressors using cross-entropy or reconstruction objectives and distillation-based prompt compression.The comparison table covers knowledge-injection methods relative to KMs.Source: §4 · p. 9, Table 3 · p. 9

5 Conclusion

The conclusion presents Deep Context Distillation (DCD) as a way to encode document-specific knowledge in lightweight, plug-and-play Knowledge Modules, with gains across question-answering benchmarks and synergy with RAG. The investigation assumes the question’s relevant document is known; document-agnostic deployment is left for future work.

5 Conclusion

  • DCD-based Knowledge Modules outperform traditional next-token prediction and pre-instruction tuning across several question-answering benchmarks.The conclusion also highlights consistent gains and synergy with retrieval-augmented generation (RAG).Source: §5 · p. 10
  • Lightweight, plug-and-play Knowledge Modules encode document-specific knowledge using Deep Context Distillation.The modularity enables fine-grained control over knowledge integration.Source: §5 · p. 10
  • The investigation assumes the document relevant to each question is known; document-agnostic deployment with multiple modules is future work.The authors identify zero-shot routing as a potential route to that setting.Source: §5 · p. 10

A Synthetic Data Generation

The appendix describes how experimental documents are chunked for generated-data experiments.

A Synthetic Data Generation

  • Documents are split into non-overlapping chunks of 2048 tokens for the generated-data experiments.The passage specifies this as the document-splitting procedure.Source: A Synthetic Data Generation · p. 13
  • The experiments use RecursiveCharacterTextSplitter from the LangChain text splitters package.This is the named tool used to create the chunks.Source: A Synthetic Data Generation · p. 13
  • The appendix indicates that the prompts used for generation are provided below the description.The passage does not specify their contents here.Source: A Synthetic Data Generation · p. 13

A.1 Summary Generations

The appendix describes three synthetic-data approaches: chunk-level summaries, questions and answers, and Entigraph-generated relationship data.

A.1 Summary Generations

  • The summary-generation setup allows up to 16 summaries per chunk and uses 8 by default in the main experiments.The prompt asks for a grammatical summary that retains important details across the text’s sections.Source: §A.1 · p. 13

A.2 Q/A Generation

  • The Q/A prompt requests challenging questions answerable from the text, encouraging aggregation across it and coverage of different sections.The prompt provides a question-and-answer format with numbered placeholders.Source: §A.2 · p. 14

A.3 Entigraph Generation

  • Entigraph data comes from a public QuALITY dataset and is split into relationship-focused chunks averaging about 700 tokens.Because the original location of each fact is unknown, training uses the full document in context.Source: §A.3 · p. 14

B.1 Impact of different losses

The appendix identifies two ablation scopes and reports that the NQA ablation uses a 140-document subset.

B.1 Impact of different losses

  • The NQA ablation uses 140 documents: 100 for training, 20 for development, and 20 for testing.This is the subset used for the following ablation.Source: §B.1 · p. 14
  • The distillation-component ablation uses Phi-3 as its backbone.The caption identifies the table’s scope but provides no component-level outcomes.Source: Table 4 · p. 14
  • A separate ablation examines the data used for DCD training.The supplied caption does not state which data choice performs best.Source: Table 5 · p. 14

B.2 Impact of Document DCD over standard next-token prediction

Progressively removing distillation components makes performance converge toward the language-model baseline, supporting the authors’ conclusion that the overall distillation objective is central to the gains. The ablation varies hidden-state matching and teacher-probability temperature.

B.2 Impact of Document DCD over standard next-token prediction

  • The authors hypothesize that hidden-state matching and KL distillation provide a richer learning signal than standard next-token prediction.They relate this hypothesis to findings that teacher-model matching can be more data-efficient than data matching (He et al., 2022; Minixhofer et al., 2025).Source: §B.2 · p. 15
  • The ablation removes hidden-state l1 loss and lowers teacher-probability temperature until the target becomes one-hot, mimicking next-token prediction.Source: §B.2 · p. 15
  • As distillation components are stripped away, performance converges toward the LM baseline, indicating the overall objective is central to the gains.The experiment uses NarrativeQA with phi-3 as the base model; Table 6 reports Rouge-L scores for DDCD variants.Source: §B.2 · p. 15, Table 6 · p. 15

B.3 Forgetting of General Knowledge after KM training

On a 400-example MMLU subset, a base model equipped with a randomly sampled Knowledge Module retains overall general knowledge, dropping only 1 point. The paper frames modules as removable without modifying the base model.

B.3 Forgetting of General Knowledge after KM training

  • The MMLU evaluation uses 400 examples evenly split across subjects, with a randomly sampled KM used for each question.Source: §B.3 · p. 15
  • Because KMs do not modify the base model, the authors say a module can be downweighted or removed if it alters other information unintentionally.Source: §B.3 · p. 15
  • The base model with a Knowledge Module retains overall general knowledge, dropping only 1 point on MMLU.Table 7 compares accuracy for the base model and its KM variant.Source: §B.3 · p. 15, Table 7 · p. 15
Loading 2503.08727v4…