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

TL;DR

This paper introduces KARD, a method to enhance small language models' reasoning abilities by distilling knowledge from large models and external knowledge bases, significantly improving performance on knowledge-intensive tasks.

Contribution

The paper proposes KARD, a novel knowledge-augmented distillation approach that enables small LMs to better memorize and utilize external knowledge for reasoning tasks.

Findings

KARD improves small T5 and GPT models on reasoning datasets.

250M T5 models outperform larger 3B models with KARD.

Significant performance gains on MedQA-USMLE, StrategyQA, OpenbookQA.

Abstract

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.