PaperlayerView exact revision on arXiv

Source-linked AI summary

Knowledge-Augmented Reasoning Distillation for Small Language Models in Knowledge-Intensive Tasks

Minki Kang, Seanie Lee, Jinheon Baek, Kenji Kawaguchi, Sung Ju Hwang

arXiv:2305.18395v2cs.CLcs.AIcs.LG

TL;DR

Small LMs have limited capacity to memorize the knowledge required for knowledge-intensive reasoning, while LLM deployment can be computationally expensive and raise privacy concerns. KARD fine-tunes small LMs on LLM-generated rationales with external retrieved knowledge and uses a neural reranker to prioritize rationale-relevant passages. KARD significantly improves small LMs on MedQA-USMLE, StrategyQA, and OpenbookQA; KARD with 250M models achieves higher accuracy than fine-tuned 3B models.

  • Problem

    Small LMs have limited capacity to memorize the knowledge required for knowledge-intensive reasoning, while LLM deployment can be computationally expensive and raise privacy concerns.

    Source: §1 · pp. 1–2

  • Method

    KARD fine-tunes small LMs on LLM-generated rationales with external retrieved knowledge and uses a neural reranker to prioritize rationale-relevant passages.

    Source: §1 · p. 2

  • Results

    KARD significantly improves small LMs on MedQA-USMLE, StrategyQA, and OpenbookQA; KARD with 250M models achieves higher accuracy than fine-tuned 3B models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    KARD provides a pathway to improve small LMs on reasoning tasks that require knowledge and reasoning abilities.

    Source: §7 · p. 10

    KARD’s effectiveness heavily relies on retrieved-document quality, and its reranker may miss important passages compared with silver knowledge.

    Source: §6.3 · p. 10

Abstract

from arXiv · show

Large Language Models (LLMs) have shown promising performance in knowledge-intensive reasoning tasks that require a compound understanding of knowledge. However, deployment of the LLMs in real-world applications can be challenging due to their high computational requirements and concerns on data privacy. Previous studies have focused on building task-specific small Language Models (LMs) by fine-tuning them with labeled data or distilling LLMs. However, these approaches are ill-suited for knowledge-intensive reasoning tasks due to the limited capacity of small LMs in memorizing the knowledge required. Motivated by our theoretical analysis on memorization, we propose Knowledge-Augmented Reasoning Distillation (KARD), a novel method that fine-tunes small LMs to generate rationales obtained from LLMs with augmented knowledge retrieved from an external knowledge base. Moreover, we further propose a neural reranker to obtain documents relevant to rationale generation. We empirically show that KARD significantly improves the performance of small T5 and GPT models on the challenging knowledge-intensive reasoning datasets, namely MedQA-USMLE, StrategyQA, and OpenbookQA. Notably, our method makes the 250M T5 models achieve superior performance against the fine-tuned 3B models, having 12 times larger parameters, on both MedQA-USMLE and StrategyQA benchmarks.

1 Introduction

Knowledge-intensive reasoning requires both domain knowledge and reasoning, but small models have limited capacity to memorize the knowledge such tasks need. KARD augments rationale distillation with retrieved knowledge and a reranker, and improves small-model performance across three QA benchmarks.

1 Introduction

  • LLM deployment can be difficult because of computational expense and potential privacy leakage from black-box, API-only access.The text gives 326GB of GPU memory to load GPT3-175B as an example.Source: §1 · p. 1
  • Small LMs lack the memorization capacity for knowledge-intensive tasks, motivating external knowledge as non-parametric memory alongside distilled reasoning.The paper theoretically shows that this memory can reduce the bits of training data a model must memorize to perform well.Source: §1 · p. 2
  • KARD fine-tunes a small LM to generate an LLM-obtained rationale and answer from a question and relevant passages retrieved from an external knowledge base.The described knowledge base may be Wikipedia.Source: §1 · p. 2
  • A neural reranker prioritizes passages useful for rationale generation when inference uses the question, because question-based retrieval may otherwise return irrelevant passages.During training, a rationale query retrieves pertinent knowledge; at inference, the question is available instead.Source: §1 · p. 2
  • KARD significantly improves small OPT and T5 language models on MedQA-USMLE, StrategyQA, and OpenbookQA versus few-shot in-context learning, fine-tuning, and reasoning distillation without knowledge augmentation.KARD with 250M models achieves higher accuracy than fine-tuned 3B models.Source: §1 · p. 2

2 Related Works

Prior work transfers LLM reasoning to smaller models, but knowledge-intensive reasoning also requires relevant factual evidence. This paper addresses that gap with a reranker trained to prioritize rationale-relevant passages.

2 Related Works

  • LLMs offer strong knowledge-intensive reasoning, but their computational expense and black-box access pose challenges in offline and privacy-sensitive settings.The paper frames smaller models as an alternative for these settings.Source: §2 · p. 3
  • Reasoning distillation transfers LLM-generated rationales to small LMs, but has been less effective when factual knowledge is important to rationale generation.Prior work cited in the passage found this limitation for knowledge-intensive tasks.Source: §2 · p. 3
  • Existing retrieval-augmented LMs have not been well explored for complex reasoning over knowledge, and prior work found a retriever insufficient for these tasks.KARD’s reranker uses LLM rationales to prioritize passages relevant to reasoning rather than relying on plain queries alone.Source: §2 · p. 3

3 Motivation: Effect of Knowledge-Augmentation on Memorization

The section motivates external knowledge as a way to reduce small models’ memorization burden. Its analysis contrasts the information required without augmentation with a lower bound under knowledge augmentation.

3 Motivation: Effect of Knowledge-Augmentation on Memorization

  • Small models may perform worse on knowledge-intensive tasks because they cannot memorize the training data that the task requires.The analysis builds on prior results that memorization capacity grows with model size and that memorization can be necessary for good language-task performance.Source: §3 · p. 3

3.1 Background without Knowledge-Augmentation

  • Under Brown et al.’s setting, any learner attaining ε-suboptimal error must memorize Ω(nd) bits of training data.Here d corresponds to the size of the knowledge base.Source: §3.1 · p. 4

3.2 Memorization with Knowledge-Augmentation

  • With knowledge augmentation, a learner can attain ε-suboptimal error while memorizing O(min(N, n)m) bits, where m is defined in the theorem.The theorem states the existence of a pair of inference and learning algorithms with this bound.Source: §3.2 · p. 4
  • When n=N and ε=o(1), the augmented memorization bound becomes O(n log2(N + R)) bits rather than nd bits.The analysis treats R as the number of extra references in the knowledge base that are irrelevant to the task.Source: §3.2 · p. 4

4 Knowledge-Augmented Reasoning Distillation

KARD distills LLM-generated rationales into small models while supplying retrieved knowledge, then trains a reranker to select rationale-relevant passages at inference. Inference uses the question to retrieve candidates, reranks them, and generates a rationale and answer.

4 Knowledge-Augmented Reasoning Distillation

  • KARD combines rationale distillation with reranker training: small models learn from LLM-generated rationales and retrieved knowledge, while the reranker prioritizes rationale-relevant passages.The overview presents training and inference for both components.Source: §4 · p. 4, Figure 2 · p. 5

4.1 Teach Small Models to Generate Rationales with External Knowledge

  • For each training question, LLMs generate rationales, and the small LM is fine-tuned to generate the rationale followed by the answer using retrieved passages.Passages are retrieved from a large external knowledge base with the LLM rationale as the query; the question and retrieved documents condition fine-tuning.Source: §4.1 · pp. 4–5

4.2 Training Neural Reranker for Rationale Generation

  • Because inference cannot use the rationale as a query, the reranker reorders question-retrieved candidates to surface passages useful for rationale generation.Training uses rationale-retrieved passages to guide the reranker’s passage scores, and its objective is independent of the small LM’s objective.Source: §4.2 · pp. 5–6

4.3 Inference

  • At test time, the retriever obtains 100 candidates for the question, the reranker selects the top-k passages, and the small LM generates a rationale and then an answer.The rationale and answer are conditioned on the question and selected documents.Source: §4.3 · p. 6

5 Experiments

The experiments evaluate KARD across three knowledge-intensive QA benchmarks and analyze its data, model-size, retrieval, and rationale-design behavior. Results favor KARD over baselines, with gains persisting under smaller training sets and models.

5.1 Experimental Setting

  • Experiments cover MedQA-USMLE, StrategyQA, and OpenbookQA, comparing few-shot, CoT, fine-tuning, knowledge-augmented baselines, reasoning distillation, and oracle variants.MedQA-USMLE contains 12,723 questions; StrategyQA has 2,780 yes/no questions, and OpenbookQA has 5,957 science questions.Source: §5.1 · pp. 6–7

5.2 Experimental Results

  • KARD outperforms all baselines on MedQA-USMLE, StrategyQA, and OpenbookQA, including encoder-decoder and decoder-only MedQA models.The reported results also show gains over reasoning distillation without knowledge augmentation on StrategyQA and OpenbookQA.Source: Table 1 · p. 7, Table 2 · p. 7, §5.2 · pp. 7–8

5.3 Analysis

  • PubMed DAPT marginally improves reasoning distillation, while KARD contributes more substantially to performance improvement.The DAPT experiments further pre-train Flan-T5 Base on PubMed abstracts or MedWiki before applying reasoning distillation or KARD.Source: Figure 3 · p. 7, §5.3 · p. 8
  • With 25% of training data, KARD beats the same model fine-tuned on full data; 250M KARD beats 3B fine-tuning, and 780M KARD beats 11B few-shot ICL.The 3B fine-tuned model has 14 times more parameters than the 250M KARD model.Source: §5.3 · p. 8, Figure 4
  • The reranker achieves significantly better Hits@k than BM25 and consistently improves model performance over BM25 retrieval.The retrieval comparison uses top-3 silver documents as ground truth and κ*=100.Source: §5.2 · p. 7, §5.3 · pp. 8–9, Figure 4
  • More rationales improve MedQA-USMLE performance, with small gains from five to ten; larger candidate-document pools also tend to help the reranker.Rationale diversity was evaluated on Flan-T5 Base and Large; candidate-pool size was varied for both models on MedQA-USMLE.Source: Table 3 · p. 8, Table 4 · p. 8, §5.3 · p. 9
  • On MedQA-USMLE, adding inference passages without reranking can reduce performance, whereas reranked k=2 outperforms k=1.The comparison uses Flan-T5 Base and indicates that additional passages alone do not necessarily yield better rationales.Source: Table 5 · p. 8, §5.3 · p. 9
  • In one MedQA-USMLE example, reasoning distillation hallucinates a Graves’ disease treatment, while KARD produces a knowledge-supported rationale and the correct answer.The example uses Flan-T5 XL; the baseline predicts incorrectly, while KARD predicts Methimazole.Source: Table 6 · p. 9, §5.3 · p. 9

6 Discussion

The discussion compares KARD with RAG, analyzes cases where it fails, and identifies limitations in retrieval quality and evaluated model scale.

6.1 Comparison to Retrieval-augmented Generation

  • KARD achieves lower accuracy than KARD’s approach in reasoning distillation when using RAG, which uses question queries and jointly fine-tunes the generator and retriever.The comparison uses Flan-T5 base on MedQA-USMLE and T5 base on StrategyQA; Table 7 reports the experimental results with RAG and (Flan-)T5 base.Source: §6.1 · p. 10, Table 7 · p. 10

6.2 Failure Case Analysis

  • In 30 MedQA-USMLE cases where Flan-T5 XL with KARD failed but ChatGPT succeeded, 15 involved retrieval misses and 15 involved incorrect rationales and predictions despite relevant documents.The authors identify both retrieval and distillation methods as needing improvement.Source: §6.2 · p. 10

6.3 Limitations

  • KARD’s effectiveness depends heavily on retrieved-document quality, and the reranker still trails silver knowledge; experiments covered only models under 3B parameters because of computational-budget limits.The performance gap between BM25 and the reranker diminishes as model size increases, but the gap between reranker documents and silver knowledge remains significant.Source: §6.3 · p. 10

7 Conclusion

KARD combines LLM-generated rationales with external non-parametric memory to improve small LMs on knowledge-intensive reasoning. The appendix discusses the memorization costs and conditions in its theoretical analysis.

7 Conclusion

  • KARD fine-tunes small LMs on LLM-generated rationales while augmenting them with external knowledge in non-parametric memory, which the theoretical analysis says reduces memorization requirements.The paper reports that KARD outperforms traditional fine-tuning and reasoning distillation.Source: §7 · p. 10

A.1 Additional Details and Discussion

  • Because the knowledge base S is unordered and the inference algorithm does not know which document is useful for each input, some training information still must be learned and memorized.The appendix’s setup defines reference strings and training samples for its abstracted language problem.Source: §A.1 · p. 17
  • Theorem 2 states that the memorization increase with N + R possible retrieved documents is only at rate log2(N + R), allowing improvement when KB size d exceeds that logarithm.N and R count useful and unuseful documents; N + R is the total number of possible retrieved documents.Source: §A.1 · p. 17
  • When training size n grows faster than the number of useful documents N, Theorem 2 removes dependence on n and gives I(X; A(X)|P) = O(N log2(N + R)).The proof does not rely on having a singleton training sample for each reference string, unlike the cited prior proof.Source: §A.1 · p. 17

A.2 Proof of Theorem 2

The proof analyzes error by subpopulation and training-sample events, then bounds the memorization needed by constructing an inference procedure that can use the knowledge base.

A.2 Proof of Theorem 2

  • The error analysis conditions on the test subpopulation and separates cases according to whether training data exists there, includes a sufficiently long example, and has one longer than the test input.The proof defines E0 as prediction error and Qj as the event that the test sample belongs to subpopulation j.Source: §A.2 · pp. 17–18
  • When a training example provides a matching document prefix, the inference procedure can identify the reference string and predict its next symbol if no other KB document shares that prefix.The proof therefore bounds error in this case by the probability that two distinct documents in S share their first m bits.Source: §A.2 · p. 19
  • The construction memorizes at most min(N, n)(m + 1) bits, yielding I(X; A(X)|P) = O(min(N, n)m).It stores up to m bits and a subpopulation identifier for at most min(N, n) training samples.Source: §A.2 · pp. 18, 20

B Implementation Details

The implementation uses LLM-generated rationales, retrieved knowledge, and a neural reranker to train and run small models. In the retriever and knowledge-base analysis, PubMed did not improve performance over Wikipedia, while DPR degraded it.

B Implementation Details

  • GPT-3.5-turbo generates multiple rationales per training sample, and incorrect rationales are filtered using Flan-T5 base.MedQA-USMLE uses instruction and 5-shot examples; StrategyQA and OpenbookQA use a chain-of-thought prompt.Source: B Implementation Details · p. 20
  • Models are fine-tuned for three epochs; T5 replaces Flan-T5 for StrategyQA and OpenbookQA to avoid potential test-set contamination.KARD uses k = 1 documents for MedQA-USMLE and StrategyQA and k = 3 for OpenbookQA.Source: B Implementation Details · p. 20
  • At inference, rationale-generating methods use multiple rationale-prediction samples and majority voting; the retriever uses Wikipedia and BM25.The datasets use their specified official or in-house splits.Source: B Implementation Details · p. 20
  • The reranker uses ColBERT-style scoring with BioLinkBERT-base for MedQA-USMLE and LinkBERT-base for StrategyQA; LoRA is used for training.The implementation sets κ1 = 8, κ2 = 0 for MedQA-USMLE and κ1 = 4, κ2 = 4 for StrategyQA and OpenbookQA.Source: B Implementation Details · p. 20
  • Domain-adaptive pre-training uses PubMed abstracts or biomedical Wikipedia passages, with Flan-T5 base trained for three epochs.The experiment assesses the impact of further pre-training before reasoning distillation.Source: B Implementation Details · p. 21

C.1 Retriever and Knowledge Base

  • PubMed did not improve performance compared with Wikipedia, while DPR degraded performance; Table 8 analyzes the retriever choice.The authors speculate that reranking PubMed passages could help, but leave that analysis for future work.Source: §C.1 · p. 21, Table 8 · p. 21

C.2 κ1 and κ2 for Reranker Training

The reranker is trained to score passages relative to one another in a candidate set, whose composition can combine question- and rationale-retrieved passages. The analysis finds that question-only candidates are not viable for learning rationale relevance.

C.2 κ1 and κ2 for Reranker Training

  • Softmax-normalized passage scores train the reranker to learn relative scores within a computationally limited candidate set.The candidate set must therefore be selected from passages retrieved using either the rationale or the question.Source: §C.2 · p. 21
  • Combining passages retrieved with the rationale and question can help the reranker distinguish question-similar from rationale-relevant passages.This setup trains it to prioritize rationale-relevant passages among those retrieved by question at inference.Source: §C.2 · pp. 21–22
  • Question-only candidate passages are not a viable training set composition because they hinder the reranker’s learning of rationale relevance.The candidate set size is fixed at eight while κ1 and κ2 control its composition.Source: §C.2 · p. 22

C.3 Knowledge used in KARD training ( ˆDij)

Using question-retrieved passages during KARD training instead of rationale-retrieved passages leads to performance degradation.

C.3 Knowledge used in KARD training ( ˆDij)

  • Training with passages retrieved using the question rather than the rationale leads to performance degradation because they are unlikely to contain the knowledge needed for rationale generation.The analysis tests question-query retrieval against rationale-query retrieval.Source: §C.3 · p. 22

C.4 Examples of Failure Cases

The failure cases illustrate two limits: the reranker may miss relevant knowledge, or the small LM may fail to reason from a retrieved passage. StrategyQA failures mostly fall into the retrieval-failure category, with questions often requiring multiple knowledge sources.

C.4 Examples of Failure Cases

  • A MedQA-USMLE example shows that a reranker retrieval failure can leave the small LM generating a factually incorrect treatment rationale.The example suggests warm compresses for allergic conjunctivitis.Source: §C.4 · p. 22
  • Even after retrieving a relevant passage, the small LM can fail to reason from it and produce an incorrect rationale.In the example, the passage concerns pulmonary contusion, while the relevant distinction involves cardiac contusion.Source: §C.4 · p. 22
  • Most StrategyQA failure cases are retrieval failures, consistent with questions requiring complex, multiple-source knowledge.One example retrieves a passage about Japan’s flag instead of the information needed about Azerbaijan’s flag and the Powerpuff Girls.Source: §C.4 · p. 22
  • A rare StrategyQA case retrieves relevant human-sacrum information, but the LM fails to comprehend it and the question also needs canine-sacrum knowledge.The example requires comparing the number of fused vertebrae in human and Alaskan Malamute sacrums.Source: §C.4 · p. 22

D Broader Impact

KARD is described as beneficial for knowledge-intensive tasks requiring professional knowledge, including medicine. The paper also urges caution in real-world clinical use because small models may generate factually incorrect statements.

D Broader Impact

  • KARD is beneficial in domains requiring professional knowledge, such as medicine.Source: D Broader Impact · p. 23
  • Small language models used in real-world clinical applications may generate factually incorrect statements, so their use warrants caution.Source: D Broader Impact · p. 23
  • Large language models also warrant caution in clinical settings because they have not been extensively validated there and may generate inaccurate information.Source: D Broader Impact · p. 23
Loading 2305.18395v2…