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

This paper proposes a novel method for integrating new information into Large Language Models (LLMs) called Knowledge Modules (KMs) trained with Deep Context Distillation (DCD). The paper presents a new training objective, Deep Context Distillation (DCD), which aims to train lightweight, parameter-efficient KMs (implemented as LoRA modules) to encode document-level knowledge. The authors demonstrate that DCD, especially when combined with synthetic data generation, outperforms standard next-token prediction and pre-instruction tuning techniques in closed-book settings on two question-answering datasets, QUALITY and NarrativeQA, using Phi-3 3B and Llama-3.1 8B base models.

This paper proposes a knowledge module (KM), which is a LoRA adapter that is trained to mimic the logics of a teacher model that has access to new documents, without directly accessing the new documents. In other words, KM is designed to learn the new information by distilling from a model that has access to these new information, so that a model equiped with KMs no longer needs to access the new documents at inference-time, which saves compute and storage. The authors proposed several training variants of KMs, namely DDCD (which treats new documents as chunks and try to recover chunks given context), SDCD, which distills generated question and answer pairs based on the new documents instead of raw texts, and KMs can be optionally equiped with knowledge extractors (KE) that can further optimize based on task-specific supervision data.

Experiments show that the proposed KM methods outperform continual pre-training baselines, as well as in-context learning and RAG baselines in certain settings.

• This study introduces a new approach for integrating new knowledge into language models through a modular, plug-and-play framework based on knowledge distillation with next-token prediction. Experiments across standard benchmarks demonstrate the effectiveness of the proposed method.

In this paper, the authors introduce knowledge module, a lightweight adapter that encode document-level knowledge into continuous parameters that can be plugged into pretrained LLMs for improved knowledge-intensive tasks. They proposed deep context distillation (DCD) is a novel training paradigm that distills both hidden representations and output distributions from a teacher model with access to the document. To improve training, the authors generate synthetic data in the form of summaries, QA pairs etc., demonstrating that DCD can outperforms baselines on long-context QA tasks across close and open-book settings. The authors also introduce task-specific knowledge extractors to combine KMs with retrieval-augmented generation, with additional gains and highlighting the effectiveness of the proposed method.

We thank the reviewers for their overall positive outlook on our work and for their constructive feedback. Two reviewers raised questions regarding the training cost of Knowledge Modules, which we address below.

$` `$

On the cost of training Knowledge Modules

We generally find that the cost of training KMs is proportional to the length of said documents; longer documents require more training steps. That said, longer documents are also costlier to run in-context; Indeed, for our experiments on NarrativeQA, we find that the full cost of training KMs in FLOPS is equivalent to 173 ICL inference queries. Therefore, the training cost is amortized over time as the number of inference queries rises. Importantly, KM training shifts the cost incurred from the user to the developer, making the user experience better as it reduces latency, a property which we believe to be very important in many settings.

We will add a section on the cost-related matters into the camera-ready version, following the reviewers’ feedback.

$` `$

Cost breakdown

We use the approximation in [1], where a decoder-only Transformer with N parameter uses 2N FLOPS per token at inference and 6N FLOPS per token for training. For training of NarrativeQA with (at most) 1500 training steps, where the teacher sequence length is 4096 and the mean student sequence length is 942 tokens, the total cost is

n_training_steps x ( student_cost + teacher_cost) = 1500 x (6N * 942 + 2N * 4096) = 20.8M x N FLOPs

On the other hand, a single ICL query with the full document (assuming an average sequence length of 60000 tokens) gives:

2N * 60000 = 0.12M x N FLOPs

Thus, after 20.8 / 0.12 = 173 inference calls, the training cost is fully amortized.

We also compute this equivalence for RAG. Without counting the cost to embed the document chunks, and the cost to perform the retrieval, when performing inference calls with RAG, we typically use the full 4K context window. This gives us a cost of:

2N * 4096 = 8.192K x N FLOPs

Thus, the training cost is amortized after 20.8M / 8.192K = 2539 inference calls.

[1] Jared Kaplan, Sam McCandlish, T. J. Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeff Wu, and Dario Amodei. Scaling laws for neural language models.

The paper proposes a creative approach to injective new knowledge into LLMs using LoRA adapters (called Knowledge Modules) trained via distillation from a teacher LLM. The reviewers broadly find this paper exciting, particularly highlighting the (1) plug-and-play nature of the approach, (2) potential use cases in private domains.

The authors should address the following points to strengthen their next version of the paper:

1. (Reviewer LZaV) Conduct addtitional analysis on why distillation is superior to simple comtinued pre-trianing on the new data. Additionally, include results from prior continued pre-training baselines such as “Synthetic continued pretraining” (https://arxiv.org/abs/2409.07431) that report superior results on Quality.
2. Include discussion of the current approach’s limitations in scaling up the approach to large datasets.

[Automatically added comment] At least one review was discounted during the decision process due to quality]