PaperlayerView exact revision on arXiv

Source-linked AI summary

DKL: Decoupled Knowledge Learning for Instruction-Tuned Language Models

Kushagra Bhushan, Meghanadh Pulivarthi, Sai Krishna Reddy Sathi, Gaurav Pandey, Sonam Gupta, Vineet Kumar, Jaydeep Sen, Yatin Nandwani, Sachindra Joshi, Dinesh Raghu

arXiv:2609.02685v1cs.CLcs.AI

TL;DR

Existing approaches either depend on retrieval, large synthetic-QA corpora, or costly instruction fine-tuning, while EPT on instruction-tuned models can damage instruction-following. DKL trains a knowledge adapter through EPT on the base LLM and merges it with the Instruct LLM. The paper reports that DKL consistently outperforms SFT-based methods while requiring substantially less synthetic data.

  • Problem

    RAG is vulnerable to retrieval failures, while EPT can compromise instruction-following and synthetic-QA methods require large generated corpora or costly training resources.

  • Method

    DKL trains a knowledge adapter through extended pre-training on the base LLM, then transfers and merges it with the instruction-tuned model.

  • Results

    DKL consistently outperforms state-of-the-art SFT-based knowledge-infusion methods while requiring substantially less synthetic data.

  • Takeaways & Limitations

    DKL is presented as a practical and scalable alternative for rapidly incorporating domain-specific knowledge into LLMs without costly IFT.

  • Takeaways & Limitations

    DKL requires an available base model that has not undergone instruction fine-tuning, limiting applicability because many open-weight models are released instruction-tuned.

Abstract

from arXiv · show

RAG has become the de facto method for incorporating new, corpus-specific knowledge into an instruction following LLM (Instruct LLM). Although RAG-based prompting improves factual grounding, it fails when retrieval is incorrect or incomplete, leading to hallucinations. Finetuning methods such as RAFT and PA-RAG enhance RAG by injecting new knowledge into the model's parameters, but require generating a massive amount of synthetic QA that covers the entire corpus. Extended Pre-Training (EPT) on the text corpus avoids the need for comprehensive synthetic data generation but compromises an Instruct LLM's instruction-following capabilities, necessitating instruction fine-tuning (IFT) after pre-training. However, IFT is costly and may be infeasible due to the unavailability of an instruction-tuning corpus. In this work, we propose DKL-Decoupled Knowledge Learning for Instruction-Tuned Language Models. Instead of doing EPT on the Instruct LLM, DKL performs EPT on its corresponding base LLM to infuse new knowledge. These knowledge infused weights are then merged with the Instruct LLM, imparting new knowledge without affecting their instruction-following capabilities. DKL is a lightweight method that avoids expensive instruction fine-tuning and relies on model merging to infuse the new knowledge into the Instruct LLM without destroying its instruction following capabilities. Empirical results show that DKL improves RAG accuracy from 54.17 to 79.26 on retrieval failure cases, while outperforming prior approaches with substantially less training data.

1 Introduction

Specialized documents expose weaknesses in retrieval and knowledge-injection approaches: retrieval failures can cause hallucinations, while extended pre-training or synthetic-data methods compromise instruction-following or require costly resources. DKL addresses this by training a knowledge update on the base LLM and merging it with the Instruct LLM.

  • Motivation: RAG can produce hallucinated or incomplete answers when retrieval fails on scarce, proprietary, or specialized documents.Injecting document knowledge into model parameters may provide parametric fallback knowledge.
  • Motivation: EPT ingests new knowledge effectively, but applying it to Instruct LLMs can cause catastrophic forgetting of instruction-following skills.Prior approaches therefore retrain the base LLM and typically repeat IFT afterward.
  • Motivation: Synthetic-QA finetuning avoids costly post-ingestion IFT, but requires massive generated datasets that may be unavailable or prohibitively expensive.These datasets must generally cover the new documents to support knowledge infusion.
  • DKL method: DKL computes an instruction-following vector from instruct–base weight differences and a knowledge vector by EPT on new documents using the base LLM.The knowledge update uses unsupervised next-token prediction on the new corpus.
  • DKL method: Adding both vectors to the base LLM produces a model combining document knowledge with the Instruct LLM’s instruction-following skills.DKL can optionally use a small amount of synthetic QA supervision to improve query-time recall.
  • Contribution: DKL is presented as a lightweight method that avoids costly IFT and uses model merging to impart new knowledge without destroying instruction-following capabilities.The method also addresses transferability by using the instruct model’s token embeddings during knowledge-vector training.

2 Related Work

Prior work uses retrieval, parameter injection, synthetic-QA finetuning, or task-vector merging to add domain knowledge, but each leaves limitations that DKL targets. DKL specifically addresses vocabulary mismatch and optimizes the interpolation of knowledge and instruction-following vectors.

  • Retrieval-Augmented Generation: RAG grounds responses in retrieved passages but remains vulnerable to retrieval failures that can lead to hallucinations.Its effectiveness depends on retrieval quality at inference time.
  • Static knowledge injection: Static knowledge injection uses finetuning or EPT to place domain knowledge in model parameters for closed-book inference.EPT has been applied across medical, materials, and finance domains, but can regress instruction-tuned models on general capabilities.
  • Synthetic-QA finetuning: RAFT and PA-RAG combine parameterized knowledge with retrieved passages, yet rely heavily on large synthetic QA corpora.Their SFT-based approach avoids repeated IFT after knowledge ingestion but depends on extensive generated supervision.
  • Model merging via task vectors: Task vectors define weight differences between trained models for model merging, whereas DKL additionally trains on the base model and addresses vocabulary-distribution mismatch.The mismatch can involve instruction-specific tokens such as Mistral’s [INST].
  • Model merging via task vectors: Compared with Chat-Vector, DKL uses instruct-model embeddings during knowledge-vector training and optimizes the interpolation ratio when combining vectors.The paper reports both design choices as critical for effective knowledge injection.

3 Methodology

DKL separates instruction-following and knowledge acquisition by training a knowledge adapter on the base LLM, then combining it with the instruct model. This model-merging approach aims to retain instruction-following while adding document knowledge, with optional synthetic QA improving recall.

  • Decoupled knowledge and instruction vectors: DKL represents instruction-following and new knowledge as separate weight transformations derived from the instruct and base LLMs.The instruction-following vector is the instruct-minus-base weight difference, while the knowledge vector is learned from new documents through unsupervised next-token prediction.
  • Knowledge adapter training: Training the knowledge adapter on the instruct LLM can deteriorate instruction following, so DKL instead trains it on the base LLM.The base model uses the same unsupervised next-token-prediction objective as extended pre-training and is therefore treated as a more suitable starting point.
  • Model merging: The knowledge-infused base parameters are combined with the instruct model’s capabilities through task-vector-style model merging.The resulting combined model is intended to capture both the new corpus knowledge and instruction-following skills.
  • Optional synthetic QA: Synthetic question-answer pairs can be added as a small recall-enhancement signal without spanning the entire corpus.Unlike instruction fine-tuning, the concatenated system prompt, question, and answer are treated as part of the knowledge to ingest.
  • Token embeddings: DKL uses instruct-model token embeddings to mitigate mismatches between base-model training and instruct-model inference.The mismatch can arise from special tokens introduced during instruction fine-tuning that were not present when training the base model.

4 Experimental Setup

The experiments evaluate DKL on two technical RedBook corpora and the non-technical QuALITY dataset, comparing it with three knowledge-ingestion baselines. The setup uses synthetic QA at different corpus-coverage levels and selects DKL’s scaling hyperparameter by validation performance.

  • Datasets: Experiments use two technical RedBook datasets and QuALITY, with ten randomly sampled QuALITY articles serving as the knowledge base.The RedBook datasets include technical text and corresponding test question answers, while QuALITY contains long-form open-domain articles.
  • Baselines and evaluation: DKL is compared with RAFT, PA-RAG, and Chat-Vector, while a human study checks whether the LLM judge aligns with human judgments.RAFT and PA-RAG both rely on synthetically generated question-answer pairs for knowledge ingestion.
  • Synthetic training data: RAFT and PA-RAG generate synthetic QA whose total words equal twice the corpus size, with PA-RAG additionally requiring multiple answers per question.These baselines generate data intended to cover the entire corpus.
  • Synthetic training data: DKL selects synthetic QA totaling 50% of corpus words and therefore does not require full-corpus coverage.The selected question-answer pairs are randomly sampled from the generated material.
  • Training details: DKL trains knowledge vectors with LoRA adapters of rank r = 16 and chooses the scaling hyperparameter α from four candidates using validation performance.The candidates are α ∈ {0.25, 0.5, 0.75, 1.0}; training continues to convergence while α controls overfitting.

5 Experimental Results

The experiments evaluate DKL's knowledge ingestion, RAG behavior, synthetic-data dependence, corpus coverage, model robustness, and embedding choice. Across these settings, DKL generally outperforms baselines while requiring less synthetic QA data.

  • DKL is evaluated for parametric QA knowledge ingestion and for combining ingested knowledge with retrieved context in RAG.
  • Synthetic QA imbalance over-represents some concepts, causing RAFT and PA-RAG to answer about short-term risks when questions concern long-term risks.This failure can occur even when retrieved passages explicitly contain the long-term-risk information.
  • DKL consistently outperforms baselines across datasets by exposing the model to the full corpus during EPT and merging the result with the instruct model.In successful-retrieval RAG cases, reported gains reach 92.13 and 83.85.
  • DKL achieves near-optimal RAG performance with 0.5× synthetic data, while PA-RAG improves by more than 10% when scaling from 0.5× to 4×.The performance difference for DKL between 0.5× and 4× synthetic data is 2.24%.
  • Adding QA data from only chapters 1–3 improves DKL performance on chapters 4–5 compared with corpus-DKL trained without synthetic QA.This is the biased-DKL versus c-DKL control experiment.
  • DKL outperforms baselines across four model architectures and sizes, and instruct-model token embeddings improve performance in both QA and RAG setups.Without instruct embeddings, performance drops significantly under retriever failure cases and approaches PA-RAG.

6 Conclusion

DKL is presented as a lightweight approach that transfers knowledge learned through base-model extended pretraining into an instruction-tuned model. The experiments report stronger performance than SFT-based methods while using substantially less synthetic data.

  • DKL trains a knowledge adapter through extended pretraining on the base LLM and transfers it to the instruct LLM for knowledge infusion.
  • DKL consistently outperforms SFT-based methods such as RAFT and PA-RAG while requiring substantially less synthetic data.

7 Limitations

The paper identifies constraints on DKL's applicability and model-merging procedure. Its use depends on access to a suitable base checkpoint and on sensitive merge-hyperparameter selection.

  • DKL requires a base model that has not undergone instruction fine-tuning, but most open-weight models are released in instruction-tuned form.Instruction-tuned checkpoints are described as more brittle under unsupervised training, limiting the method's applicability.
  • The merging procedure requires extensive hyperparameter search because optimal knowledge-vector, task-vector, and checkpoint choices are sensitive to the data distribution.The paper states that no principled theoretical framework or automated selection method is currently available.

C.2 Human Annotation Results

The human annotation study evaluates agreement and dataset quality across QA and RAG setups. Automatic filtering removes many ill-formed QA pairs, while LLM-as-a-Judge aligns more closely with human judgments for RAG than QA.

  • α ≈0.66 for QA versus ≈0.92 for RAG, indicating lower human agreement for the more verbose QA responses.QA responses averaged 196 words versus 135 for RAG and often included hallucinations or extraneous details.
  • ∼71% of human-identified ill-formed QA pairs were removed by automatic filtering with Llama 3.1 70B Instruct.The judge rated examples on Accuracy, Relevance, Clarity, and Usefulness, retaining only pairs rated 10.
  • ∼84% accuracy on QA and ∼97% on RAG were achieved by the LLM-as-a-Judge against human majority labels.The stronger alignment in RAG reflects more grounded responses, whereas context-free QA generations were more verbose and ambiguous.
  • Synthetic test sets contained a non-trivial proportion of ill-formed QA pairs, motivating scalable filtering and evaluation with an LLM judge.The study used human annotations and then extended automatic filtering to the full test dataset.

E Results on Llama

This section presents the main Llama-model results and dataset statistics in the paper’s result tables. The supplied passages identify the tables but do not report their numerical outcomes.

  • Table 10 compares DKL with the paper’s baselines using the Llama 3.1 8B model.
  • Table 11 reports DKL and baseline results on RedBook 2.
  • Table 9 provides data statistics for the datasets used in the paper.

G Stopping Criteria Ablation

The stopping-criteria ablation examines how extended-pretraining duration and merge timing affect downstream performance. Performance peaks at intermediate checkpoints and declines with continued training, making convergence-based merging conservative.

  • The ablation trains the LLaMA 3.1 8B base model for 60 epochs and evaluates four merge weights from 0.25 to 1.0.Experiments use the Book 1 corpus and select the optimal merge according to LLMaJ Score under RAG.
  • Performance improves in early and middle training, peaks at intermediate checkpoints, and gradually declines as training approaches convergence.The trend is observed across the evaluated setups after selecting the best merge weight at each checkpoint.
  • Training until convergence before merging makes the main-article scores conservative estimates.More careful stopping criteria could further enhance performance.
  • DKL is reported to outperform all baselines across additional model architectures and sizes.The robustness experiment uses SmolLM2-1.7B-Instruct and Qwen3-0.6B on Redbook1.

I Token Swapping Ablation

The token-swapping ablation studies whether embedding consistency and selectively swapping divergent tokens improve DKL. The results favor swapping the full instruct embedding layer, with the most divergent tokens carrying much of the observed effect.

  • The Qwen3-0.6B ablation shows trends similar to those observed with Mistral.
  • The top 50 swapped embeddings achieve performance equivalent to the full DKL regime.Embeddings with the largest instruct–base differences influence the final trained model more than the remaining embeddings.
  • Comparable performance for bottom_k-50-DKL and e-DKL indicates that embeddings shared by the instruct and base models have little impact.e-DKL denotes the version without embedding swaps.
  • Increasing the selected embedding set from top_k-50 to top_p-0.9 causes a performance drop rather than an improvement.The authors attribute this behavior to token usefulness and increasing inconsistency within the embedding layer.
  • Full-layer swapping preserves embedding consistency and is advised over targeted-token swapping because it introduces no computational bottleneck.The authors posit that out-of-distribution tokens initially help but later interfere as inconsistency increases.

J Performance on general tasks

DKL is evaluated against Instruct, RAFT, and PARAG across several general benchmarks. It maintains competitive general-task performance, while RAFT and PARAG regress relative to Instruct.

  • Benchmark setup: DKL, Llama 3.1 8B Instruct, RAFT, and PARAG are compared across several general benchmarks.The benchmarks include language understanding, reasoning, graduate-level STEM, difficult mathematics, general knowledge, and multistep reasoning tasks.
  • Benchmark setup: The evaluation covers Big Bench Hard, GPQA, MATH-Hard, MMLU-Pro, and MUSR.These benchmarks span 23 challenging tasks, graduate-level STEM questions, math competition questions, diverse-field knowledge questions, and LLM reasoning capabilities.
  • Results: DKL maintains competitive performance across all general benchmarks.This result is reported in Table 15 for RedBook 1.
  • Results: RAFT and PARAG show regression on general tasks relative to Instruct.The comparison is reported alongside DKL's general-benchmark performance in Table 15.
Loading 2609.02685v1…