SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {L} _ {\text {r e r a n k}} (\phi) = \frac {1}{n \cdot l} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {l} D _ {\mathrm {K L}} \left(Q \left(\boldsymbol {d} \mid \boldsymbol {r} _ {i j}\right) \| P _ {\phi} \left(\boldsymbol {d} \mid \boldsymbol {x} _ {i}\right)\right). \tag {3}
$$
Intuitively, the objective function guides the reranker to assign higher scores to passages that are similar to the rationale  $\pmb{r}_{ij}$ . Note that both objective  $\mathcal{L}_{\mathrm{distill - KB}}(\theta)$  and  $\mathcal{L}_{\mathrm{rerank}}(\phi)$  are independent; therefore, we do not need to jointly update both of the small LM and the reranker.
4.3 Inference
After training, we obtain the small LM with the parameter  $\theta^{*} \in \mathrm{argmin}_{\theta} \mathcal{L}_{\mathrm{distill - KB}}(\theta)$  and the reranker with the parameter  $\phi^{*} \in \mathrm{argmin}_{\phi} \mathcal{L}_{\mathrm{rerank}}(\phi)$ . At the test time, to answer the question  $x_{*}$ , we first get a set of candidate documents  $\tilde{\mathcal{D}}_{*} = \mathrm{topk}(\rho(\boldsymbol{d}|\boldsymbol{x}_{*};\mathcal{D}),\kappa^{*})$  with the retriever  $\rho$  and  $\kappa^{*} = 100$ . Then we re-rank all the document  $d \in \tilde{\mathcal{D}}_{*}$  with  $f_{\phi^{*}}$  and choose top- $k$  relevant documents w.r.t the question  $x_{*}$  as follows:  $\hat{\mathcal{D}}_{*} = \mathrm{topk}(\{f_{\phi^{*}}(\boldsymbol{d},\boldsymbol{x}_{*}) \mid \boldsymbol{d} \in \tilde{\mathcal{D}}_{*}\}, k)$ . Finally, we generate a rationale  $\boldsymbol{r}_{*} = \mathrm{argmax}_{\boldsymbol{r}} p_{\theta^{*}}(\boldsymbol{r}|\boldsymbol{x}_{*},\hat{\mathcal{D}}_{*})$  and an answer  $\boldsymbol{y}_{*} = \mathrm{argmax}_{\boldsymbol{y}} p_{\theta^{*}}(\boldsymbol{y}|\boldsymbol{r}_{*},\boldsymbol{x}_{*},\hat{\mathcal{D}}_{*})$ .
Task and Dataset In our experiments, we focus on knowledge-intensive reasoning tasks which require both the reasoning ability over the knowledge and the compound knowledge of the specific domain. As our primary benchmark, we use the medical multiple-choice question dataset — MedQA-USMLE [23]. The dataset contains 12,723 4-option multiple-choice question answering problems from US medical licensing exam. This dataset is the best fit to evaluate our method since  $98\%$  of the questions simulate the realistic clinical settings by presenting patient cases that require extensive professional domain-specific knowledge and complex reasoning ability over multiple evidence sources. To further validate our approach, we employ StrategyQA [14] dataset, which involves 2,780 yes/no questions that demand sophisticated multi-step reasoning skills and the ability to gather supporting evidence from various domains. We additionally validate our approach on commonsense reasoning with OpenbookQA [39] dataset, which consists of 5,957 elementary-level science questions with 4 multiple-choice options.
Baselines We compare our method against relevant baselines. Few-shot In-context Learning (ICL) utilizes a few training samples as a prompt to make a prediction [5]. Few-shot ICL + Chain-of-Thought (CoT) leverages chain-of-thought prompting to generate a rationale and generate an answer based on the rationale [28]. Fine-tuning refers to the one that fine-tunes a pre-trained model to generate an answer given only a question. The performance of the above baselines represents the capability of a small language model to solve knowledge-intensive reasoning tasks using only training data but without any extrinsic guidance on reasoning or external knowledge.
To assess the impact of external knowledge, we augment the above three baselines with documents retrieved from the knowledge base (Wikipedia), denoted as Knowledge-Augmented models. For the knowledge augmentation, we append retrieved passages along with the question at both training and inference time. These baselines help us understand how much external knowledge can improve the performance of each baseline. We also compare our KARD against the standard Reasoning Distillation without knowledge-augmentation [33, 17, 38, 12, 18].
Table 1: Experimental results on the MedQA-USMLE dataset with Flan-T5 [53] and OPT [20, 59] models. We report the mean and standard deviation of accuracy with 3 different runs for reasoning distillation methods.
Method	MedQA-USMLE (Flan-T5 [53])	MedQA-USMLE (OPT [20, 59])
Base (250M)	Large (780M)	XL (3B)	350M	1.3B-IML
Few-shot	23.49	31.50	35.66	27.42	29.14
Few-shot + Chain-of-Thought (CoT)	25.22	32.21	32.99	25.06	26.39
Knowledge-Augmented Few-shot + CoT	31.34	32.60	34.41	25.84	28.75
Fine-tuning	30.71	34.49	37.39	26.47	25.77
Knowledge-Augmented Fine-tuning	33.39	37.71	39.12	25.84	28.67
Reasoning Distillation	31.03±.40	39.62±.29	46.32±.36	29.43±1.13	34.30±.95
KARD (ours, BM25)	33.14±.23	41.87±.93	47.27±.67	30.79±.78	35.48±.37
KARD (ours, Reranker)	38.15±.39	44.59±.47	48.94±.32	32.86±1.12	38.83±.46
KARD (Silver Knowledge, Oracle)	40.30	49.80	53.50	35.90	42.18
CoT from ChatGPT (Teacher, Oracle)	61.59	65.51	67.16	-	50.27
Table 2: Experimental results on the StrategyQA and OpenbookQA dataset with T5 models [45].  $\dagger$  indicates experiments with Flan-T5 having the same size. We report experimental results as in Table 1.
Method	StrategyQA (T5 [45])	OpenbookQA (T5 [45])
Base	Large	XL	Base	Large	XL
Few-shot	48.47	48.47	51.67	23.00	27.60	25.00
Few-shot + CoT	48.47	48.33	48.76	27.60	27.40	27.80
KA Few-shot + CoT	48.47	48.91	48.76	27.60	27.60	27.80
Fine-tuning	52.26	56.33	51.53	54.00	62.00	74.60
KA Fine-tuning	52.11	58.81	53.38	53.80	64.60	73.80
Reasoning Distillation	55.36±.27	64.97±.55	68.41±.48	58.87±.50	66.13±.34	77.00±.59
KARD (ours, BM25)	55.90±.24	65.94±.12	68.8±1.08	55.93±.38	64.40±.71	76.00±.28
KARD (ours, Reranker)	56.57±.25	66.04±.60	70.55±.81	59.33±.74	66.40±.16	78.53±.25
KARD (Silver Kn., Oracle)	57.50	65.65	72.34	63.40	72.40	82.00
CoT from ChatGPT (Oracle)†	66.38	67.10	72.05	58.60	78.80	87.80
Figure 3: Experimental results on MedQA-USMLE, where we perform DAPT [15] on each dataset of x-axis with Flan-T5 Base.
As oracle models, we present a variant of KARD that receives better knowledge as input. In particular, at the inference time, we augment KARD with the silver document which is the passage retrieved with the gold rationale generated by the LLM as a query. This model represents an upper bound of the neural reranker performance. Additionally, we directly provide the small instruction fine-tuned language models (Flan-T5 [53] and OPT-IML [20]) with the rationale from the LLM in inference, to assess the upper bound of the performance gain on small models with high-quality rationales.
Language Models For all the experiments, we use the T5 models [45] including Flan-T5 [9], and OPT models [59] including OPT-IML [20]. For the reranker, we use LinkBERT models [56]. As for the teacher LLM, we employ GPT-3.5-turbo (ChatGPT) [42] through the proprietary API.
See Appendix B for experimental settings in detail.
Table 1 shows that KARD consistently outperforms all the baselines on the MedQA-USMLE dataset on both encoder-decoder (Flan-T5) and decoder-only (OPT) language models. Remarkably, KARD exhibits a substantial positive effect on smaller models, as evident from the significant performance gain of the Flan-T5 Base model, which has 250 million parameters, over a fine-tuning baseline on the MedQA-USMLE dataset. Regarding the analysis of the model size, please refer to Section 5.3. The impact of KARD decreases as the size of the model increases since larger models can better memorize knowledge during pre-training and fine-tuning. Moreover, we empirically show that knowledge augmentation consistently improves performance not only in reasoning distillation but also in few-shot chain-of-thought and fine-tuning. It is worth noting that this empirical evidence supports our theoretical analysis in Section 3 that knowledge augmentation enhances the performance of small models. Furthermore, our experimental results indicate that the reranker consistently improves the performance of models for all sizes, over the retrieval with BM25. From the experimental results with silver knowledge (oracle), there is room for improvement by retrieving more relevant documents, which can help the model generate a high-quality rationale.
We also present additional experimental results on StrategyQA and OpenbookQA datasets in Table 2. Once again, KARD outperforms all baselines in experiments with both datasets. Notably, compared to MedQA-USMLE, the few-shot methods on StrategyQA and OpenbookQA exhibit performance similar to random guessing, as T5 lacks the ability of in-context learning [9]. Furthermore, fine-tuning
Figure 4: (a) Efficiency on training data and (b) model size. On MedQA-USMLE, we compare KARD against the fine-tuning baseline by varying either the number of training data with Flan-T5 Large or the number of parameters, including the few-shot in-context learning performance of Flan-T5 XXL (11B). (c) Considering silver documents as ground truth, we measure Hits@k on the documents retrieved by BM25 and the reranker.
Table 3: Analysis on rationale diversity.
Rationales	BM25
Base	Large
l = 3	30.09	35.43
l = 5	32.13	39.04
l = 10	32.91	41.79
Table 4: Analysis on  ${\kappa }^{ * }$  .
Passages	Reranker
Base	Large
κ* = 20	36.45	43.91
κ* = 50	36.06	44.23
κ* = 100	36.76	45.25
Table 5: Analysis on  $k$  .
Passages	Flan-T5 Base
BM25	Reranker
k = 1	32.91	36.76
k = 2	32.84	37.71
k = 3	32.36	37.39
T5-XL on StrategyQA results in poor performance since it fails to generalize to the test data. On the other hand, reasoning distillation improves the performance of models across all different sizes on both datasets. Our KARD further yields performance improvement over the reasoning distillation baseline, demonstrating the effectiveness of knowledge augmentation in both datasets.
Experiments with DAPT Domain-Adaptive Pre-Training (DAPT) [15] is the useful strategy to adapt Pre-trained Language Models (PLMs) on the specific domain to effectively tackle the tasks on it, which is done by further pre-training the PLM on a large-scale domain-specific text corpus [2, 30, 37]. As it is interesting to observe whether the DAPT can enhance the capacity of PLMs for reasoning distillation in domain-specific knowledge-intensive tasks by further performing training on relevant domain-specific data before distillation, we conduct experiments with models from DAPT. Specifically, we further pre-train the Flan-T5 Base model on two moderate-scale biomedical corpora, PubMed abstracts and MedWiki [35], respectively. Then, we apply reasoning distillation and KARD to PLMs with further pre-trained parameters. In Figure 3, we observe that DAPT on PubMed marginally enhances the performance of reasoning distillation. On the other hand, KARD contributes more substantially to performance improvement than DAPT. This result indicates that KARD offers a distinct advantage in knowledge-intensive reasoning tasks compared to DAPT.
Efficiency on Dataset and Model Sizes To validate the efficiency of our KARD in terms of training data and model size, we measure the test accuracy on the MedQA-USMLE dataset while varying the number of training data and model parameters. As shown in Figure 4a, our KARD can effectively transfer the reasoning ability of the LLM with the proposed KARD mechanism, using only a small number of training data. Moreover, the gaps between the naive fine-tuning and our KARD become much larger when increasing the number of training data, which confirms that we can potentially increase the effectiveness of KARD with more training data for knowledge-augmented distillation from LLMs. Furthermore, it is worth noting that KARD is a sample-efficient. With  $25\%$  of the training data, KARD outperforms the same model fine-tuned on the full data.
For the efficiency in terms of the model size, as shown in Figure 4b, KARD with 250M parameters achieves higher accuracy than the fine-tuned model with 3B parameters (14 times larger). Moreover, KARD with 780M parameters outperforms the 11B in-context learning baseline. These results show the significant practical advantage of our KARD in resource-restricted settings since the small LM with KARD requires significantly less computational cost yet it outperforms the LLM.
Retrieval Performance To evaluate the performance of the reranker on MedQA-USMLE, we consider the top-3 silver documents retrieved with the rationales generated by LLM as the ground
Table 6: Rationale generation example of KARD, compared with Reasoning Distillation in MedQA-USMLE with Flan-T5 XL. The answer is in blue. We mark the hallucination in the generated rationale as yellow and highlight the factually correct statements in the generated rationale and its supporting fact in the passage as red.
Question: A 29-year-old woman presents to the clinic after several months of weight loss. She noticed a  $6.8\mathrm{kg}$  (15 lb) unintentional weight loss over the preceding several months. She has not changed her diet or exercise habits. She also reports feuding with her boyfriend over the temperature of their shared apartment, as she always feels warmer than he does. The vital signs include: heart rate  $110 / \mathrm{min}$  and blood pressure  $146 / 78\mathrm{mmHg}$ . The physical exam is notable for warm and slightly moist skin. Which of the following is the best single treatment option for this patient? (A) Glucocorticoids (B) Methimazole (C) Propranolol (D) Radioiodine therapy
Retrieved Passage	Outputs
Reasoning Distillation (without knowledge)	No knowledge	The patient's symptoms are consistent with hyperthyroidism, specifically Graves' disease. Glucocorticoids are the first-line treatment for Graves' disease. Methimazole, propranolol, and radioiodine therapy can be used as adjunctive treatment. Answer: A
KARD (Ours, reranker)	Hyperthyroidism. (...) The best single treatment option for Graves' disease is methimazole, which may temporarily help people while other treatments are having effect. Glucocorticoids and radioiodine therapy can also be used, but they are not the best option in this case. Propranolol can help control the heart rate and blood pressure, but it will not treat (...)	The symptoms and physical exam findings suggest hyperthyroidism, specifically Graves' disease. The best single treatment option for Graves' disease is methimazole, which works by inhibiting the thyroid hormone synthesis. Glucocorticoids and radioiodine therapy can also be used, but they are not the best option. Answer: B
truth, and measure Hits@k on the documents retrieved by BM25 and reranker with  $\kappa^{*} = 100$ . In Figure 4c, the reranker achieves significantly better Hits@k than BM25. This result indicates that the reranker successfully learns to prioritize passages that are helpful to generate correct rationale at the test time, which leads to performance improvement on the knowledge-intensive reasoning tasks.
The Number of Rationales During Training Following Ho et al. [17], we generate multiple rationales for each training sample in order to facilitate diverse reasoning in small language model training. In Table 3, we present the impact of rationale diversity during training on both Flan-T5 base and large models using the MedQA-USMLE dataset. As the number of rationales per training data increases, the performance also improves, demonstrating the benefit of employing multiple rationales. However, the performance gains become small when we increase the number of rationals from 5 to 10. This suggests that utilizing more diverse rationales beyond 10 may not yield significant further improvements, at least in the MedQA-USMLE dataset.
The Number of Candidate Documents for Reranker It is crucial to determine the size of the candidate document set  $(\kappa^{*})$  to which the reranker assigns the relevance scores w.r.t a question. In Table 4, we present the performance of both Flan-T5 base and large models on MedQA-USMLE, while varying  $\kappa^{*}$ . The results indicate that increasing the number of candidate documents tends to be beneficial, as it allows the reranker to consider a broader range of diverse candidate documents.
The Number of Passages Used for Inference Even LLMs tend to be easily distracted by irrelevant context [47]. Therefore, simply adding more passages during inference may not necessarily enhance performance if relevant knowledge is not selected. In Table 5, we present the impact of the number of passages used in KARD during inference ( $k$  in Section 4.3) on Flan-T5 Base and MedQA-USMLE. We observe that the performance of KARD (BM25) without the re-ranker decreases with increasing  $k$ . This result implies that using additional passages does not always result in generating better rationales. In contrast, using two passages ( $k = 2$ ) with the reranker is better than a single passage ( $k = 1$ ). This result indicates that the reranker effectively selects more suitable knowledge than BM25, thereby contributing to performance improvement in the MedQA-USMLE benchmark.
Qualitative Analysis In Table 6, we provide an example comparing the rationale generated by our KARD against the rationale by the baseline model with reasoning distillation but without knowledge augmentation. We choose one sample from the MedQA-USMLE dataset and generate the rationale using the Flan-T5 XL model. The model without knowledge augmentation generates the rationale that seems plausible based on the given instruction and question. However, it mistakenly generates the hallucination that Glucocorticoids are the first-line treatment for Graves' disease, which is incorrect. As a result, it fails to predict the correct answer (B) Methimazole. In contrast, when the model is fine-tuned with KARD, it generates a correct rationale that is supported by the retrieved knowledge indicating that methimazole is the best single treatment option for Graves' disease. Consequently, it successfully predicts the correct answer. This example highlights the effectiveness of our KARD method for generating accurate rationales by incorporating relevant knowledge, which leads to an improved question answering performance on knowledge-intensive reasoning benchmarks.
6.1 Comparison to Retrieval-augmented Generation
Retrieval-Augmented Generation (RAG) [32] primarily focuses on solving knowledge-intensive tasks (e.g., open-domain QA), where the accurate knowledge retrieval is important to achieve higher performance. In terms of the methodology, the key differences between KARD and RAG are that RAG utilizes the question as a query and jointly fine-tunes the generator and retriever. To quantitatively analyze the advantage of our KARD
Table 7: Experimental results including RAG on Reasoning Distillation (RD) with (Flan-)T5 base.
MedQA	StrategyQA
KA Fine-tuning	33.39	52.11
RAG + RD	24.84	54.24
KARD (Reranker)	38.15	56.57
against RAG in reasoning distillation, we conduct experiments with RAG on the reasoning distillation with two datasets that we used in the main experiment, where we use Flan-T5 base for MedQA-USMLE and T5 base for StrategyQA as base LMs and DPR [25] as the trainable retriever for RAG. In Table 7, experimental results show that using RAG in reasoning distillation achieves lower accuracy than KARD, showing that our KARD is more tailored approach to reasoning distillations.
6.2 Failure Case Analysis
In Table 1, we can see significant differences between KARD with reranker on the Flan-T5 XL and the ChatGPT in MedQA-USMLE. Our investigation focuses on understanding the cause of these gaps by examining samples where our method fails while ChatGPT succeeds. We collect 30 samples from corresponding cases and categorize them into two groups. The first group consists of cases where the reranker fails to obtain the document relevant to generating the correct rationale. The second group includes cases where the small language model fails to produce correct rationales and makes incorrect predictions, despite having access to relevant knowledge in the retrieved document. Out of 30 samples, 15 fall into the first category, while the remaining 15 belong to the second category. This observation indicates the need for further improvements in both retriever and distillation methods to enhance the performance of small language models in knowledge-intensive reasoning tasks.
We have shown substantial improvements in small LMs' performance on knowledge-intensive reasoning tasks through our KARD. However, it is important to acknowledge the limitations of our study. First, in terms of methodology, the effectiveness of our knowledge augmentation heavily relies on the quality of the document retrieved from the external knowledge base. As indicated in Table 1 and Figure 4c, our reranker substantially improves the performance of small models by retrieving better knowledge. Despite the diminishing performance gap between BM25 and the reranker as the model size increases, there is still a significant difference between the document retrieved by the reranker and the silver knowledge. This indicates that the reranker might miss important passages that can augment the knowledge of even large LMs. Therefore, further advancements in retrieval methods are necessary to generate better rationale, as this remains an important research challenge even for large language models [13]. Second, regarding experiments, we have tested our approach on relatively small LMs having under 3B parameters given our limited computational budgets. However, exploring the use of relatively larger language models like GPT-3 [5, 44] or LLaMA [50, 51] with KARD could be of great interests, which is a promising direction for future research.
In this work, we proposed Knowledge-Augmented Reasoning Distillation (KARD) which enhances the capabilities of small Language Models (LMs) on knowledge-intensive reasoning tasks that demand both knowledge and reasoning abilities. Our approach involves generating rationales from large LMs and fine-tuning small LMs on these rationales, while augmenting small LMs with external knowledge from a non-parametric memory. Our theoretical analysis motivates our method by demonstrating the effectiveness of external memory in reducing the memorization requirements of small LMs. Through empirical experiments, we showed that KARD outperforms traditional approaches such as fine-tuning and reasoning distillation, thereby providing a pathway to improve small LMs in reasoning tasks that require a comprehensive understanding of domain-specific knowledge.
This work was done while the first author was working at AITRICS. We would like to thank Kangwook Lee and the anonymous reviewers for their insightful comments and suggestions regarding this work, which helped us to make improvements to the paper. This work was supported by AITRICS, the Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. 2019-0-00075, Artificial Intelligence Graduate School Program (KAIST), No. RS-2022-00187238, Development of Large Korean Language Model Technology for Efficient Pre-training, and No.2022-0-0071), the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS-2023-00256259 and NRF-2018R1A5A1059921) and Samsung Electronics (IO201214-08145-01).