SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Minki Kang $^{1,2*}$ , Seanie Lee $^{2}$ , Jinheon Baek $^{2}$ , Kenji Kawaguchi $^{3}$ , Sung Ju Hwang $^{2,4}$ ,  $^{1}$ KRAFTON,  $^{2}$ KAIST,  $^{3}$ National University of Singapore,  $^{4}$ DeepAuto.ai {zzxc1133, lsnfamily02, jinheon.baek}@kaist.ac.kr, kenji@comp.nus.edu.sg, sjhwang82@kaist.ac.kr
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.
Large Language Models (LLMs) [5, 8] have excelled at various tasks across diverse domains with in-context learning. Recently, scaling up the number of parameters of LLMs has been shown to significantly improve their knowledge encoding and reasoning capability [54, 24]. Moreover, such LLMs have achieved remarkable performance on knowledge-intensive tasks in professional domains which are highly challenging, since they require a considerable depth of domain knowledge and reasoning [34, 48]. For example, in Figure 1 top, answering a medical question requires both domain knowledge and reasoning ability. The LLM should understand that the patient likely has ALS based on the symptoms and recognize SOD1 is the main cause of motor neuron diseases. Furthermore, it needs to reason over the knowledge that a mutation in SOD1 is highly associated with the symptoms.
Despite its effectiveness, deploying LLMs can still be challenging, especially in real-world applications. Firstly, utilizing LLMs to make predictions is computationally expensive. It requires 326GB GPU memory to load the GPT3-175B model [11]. Moreover, deployment of the LLM potentially poses a risk of privacy leakage since most of the production-grade LLMs [5, 44, 8, 43] operate in a black-box manner. That is, users cannot access the parameters of LLMs but only their output via some Application Programming Interfaces. Consequently, the need for white-box Small Language Models tailored to address problems requiring domain-specific knowledge will continue to gain
*Work done at AITRICS. †Code is available at https://github.com/Nardien/KARD.
Question: A 57-year-old man presents with a 2-month history of right upper and lower extremity weakness. (...) His right upper extremity shows forearm atrophy and depressed reflexes while his right lower extremity is hypertonic with a positive Babinski sign. Which of the following is most likely associated with the cause of this patient's symptoms? A) HLA-B8 haplotype B) HLA-DR2 haplotype C) Mutation in SOD1 D) Viral infection
Rationale: The symptoms suggest motor neuron disease and among the options, mutation in SOD1 is the most common cause of inherited motor neuron diseases. A viral infection can cause peripheral neuropathy, but the symptoms described suggest a central nervous system disorder.
Figure 1: Concept. An example of a knowledge-intensive reasoning task (medical QA [23]) on the top. On the bottom, we provide the conceptual illustration of our KARD, compared to existing reasoning distillation. On the right, we provide examples of passages retrieved with rationale and question from the external KB.
prominence. To tackle the above challenges of deploying models, previous works [33, 17, 38, 12, 18] have proposed to transfer the reasoning ability of large models to small models through reasoning distillation (See Figure 1 left). In particular, they leverage the LLM to generate high-quality rationales and fine-tune a small LM to generate the rationale obtained from the LLM. This reasoning distillation improves the performance of small LMs on tasks that require complex reasoning ability (e.g., arithmetic and symbolic reasoning [10, 55]). Based on this observation, we pose a research question: "Is it possible to transfer both the domain knowledge and reasoning ability of LLMs through reasoning distillation, for tasks requiring specific knowledge to reason for answering a question?"
Existing reasoning distillation is suboptimal to solve such knowledge-intensive reasoning tasks since small, distilled LMs are limited in their capacity to memorize the knowledge that is necessary to solve the tasks due to the small number of parameters. This motivates us to develop a method that distills the reasoning ability of LLMs into smaller LMs while injecting the specific task-relevant knowledge. Specifically, we augment a small LM with the knowledge retrieved from an external Knowledge Base (KB) as a non-parametric memory, and we theoretically show that the non-parametric memory can reduce the number of bits to memorize training data for performing well.
Based on this intuition and the theoretical analysis, we propose Knowledge-Augmented Reasoning Distillation (KARD) which enables to transfer the reasoning ability of an LLM to a small LM while injecting the knowledge, for knowledge-intensive reasoning tasks. Specifically, we utilize a retriever [46] to obtain passages containing relevant knowledge for generating a rationale from an external knowledge base (e.g., Wikipedia). We then fine-tune the small LM to generate the rationale, obtained from the LLM, based on the question and the retrieved document, and predict the answer.
During training, using a rationale as a query helps retrieve pertinent knowledge for generating rationales. However, during inference, relying on the question as a query may result in poor retrieval. As shown in Figure 1, the passage retrieved with the question is not relevant to generating the rationale. To mitigate the issue, we introduce a neural reranker to prioritize passages useful for rationale generation, ensuring a retrieval of relevant documents even with the question as the query.
To verify the efficacy of KARD, we empirically show that it significantly improves the performance of small LMs (OPT [20, 59] and T5 [45, 53]) on medical Question Answering (QA) (MedQA-USMLE [23]), multi-step factual QA (StrategyQA [14]), and commonsense reasoining (Open-bookQA [39]) datasets compared to few-shot in-context learning, fine-tuning, and reasoning distillation without knowledge augmentation. Also, our extensive analyses demonstrate that our KARD is efficient in terms of both the training data and the model size. Specifically, KARD with 250M models achieves higher accuracy than the fine-tuned 3B models, and KARD outperforms the fine-tuning only with a quarter of the full training data in 780M models.
Our findings and contributions are as follows:
- We demonstrate that fine-tuning small LMs to generate rationales from large LMs is insufficient for knowledge-intensive reasoning tasks and a non-parametric external knowledge base plays a crucial role in complementing the lack of knowledge in small LMs.
- Moreover, we address the limitations of the existing retriever method by introducing a reranker, in order to obtain pertinent passages for generating rationales in knowledge-intensive reasoning tasks.
- In widely-used medical, multi-step factual, and commonsense QA benchmark datasets, we empirically show that the proposed KARD significantly improves the performance of small LMs.
Large Language Models Large Language Models (LLMs) have shown impressive capabilities across various tasks. One of their notable strengths is their ability to memorize knowledge and leverage that knowledge to solve knowledge-intensive reasoning tasks. For example, LLMs like GPT-3.5 [44], Med-PaLM [48], ChatGPT [29], and GPT-4 [43] have shown the promising performance on the challenging medical question answering task, the United States Medical Licensing Examination (USMLE) [23], even surpassing the passing score by a large margin [41]. However, deploying LLMs in offline and privacy-sensitive environments is still challenging since most of these models are in black-box (accessible via APIs), and computationally expensive. Thus, we need alternative solutions that can leverage the capabilities of LLMs for knowledge-intensive reasoning tasks.
Reasoning Distillation from LLMs Recent works [33, 17, 38, 12, 18] have attempted to distill the reasoning ability of LLMs into small LMs, where the reasoning ability is an emergent property which enables LLMs to perform better in reasoning tasks through Chain-of-Thought (CoT) prompting (e.g., Let's think step-by-step) [28, 55]. Unlike arithmetic or symbolic reasoning tasks, however, previous works [33, 17] have shown that reasoning distillation is less effective for knowledge-intensive reasoning tasks [14] where factual knowledge is important to generate accurate rationale. Therefore, we augment small LMs with documents retrieved from the external knowledge base so that the models can leverage knowledge to generate better rationales that lead to correct answers.
Knowledge-Augmented LMs Knowledge-augmented LMs have utilized an external Knowledge Base (KB) to supplement their intrinsic knowledge [16, 32, 3, 22, 58]. One common approach to incorporate external knowledge is by retrieving relevant passages from a KB, such as Wikipedia, based on the input query [7]. Retrieving the correct evidence is crucial to generate accurate answers and factually grounded rationales. However, previous works usually have not explored the use of knowledge-augmented LMs for tasks that require complex reasoning over knowledge. Recently, BehnamGhader et al. [1] examined the reasoning ability of existing retrieval-augmented LMs and found that the existing retriever [25] is insufficient for retrieving relevant passages to solve the knowledge-intensive reasoning tasks. To address this limitation, we propose a re-ranker for rationale generation that prioritizes passages relevant to the rationale generated by LLMs given the query. This approach can be seen as a form of knowledge distillation for the retriever, as we use the rationale to guide the reranker to retrieve more relevant passages for reasoning, instead of using plain queries.
3 Motivation: Effect of Knowledge-Augmentation on Memorization
Large language models are known to memorize its training data [6, 49] and the memorization capacity is proven to increase as the size of the model increases [27, 57]. The previous work [4] showed that the memorization of training data is indeed necessary to perform well in a language problem. These results suggest that the reasoning distillation with a small language model (without knowledge augmentation) will degrade the performance because of (1) the incapability of memorizing training data and (2) the necessity of the memorization to perform well. In this section, we demonstrate that using an external Knowledge Base (KB) as a non-parametric memory with a retriever reduces the amount of the memorization needed to perform well and thus allows us to use small models.
3.1 Background without Knowledge-Augmentation
We adopt the exact same problem setting used in Brown et al. [4]. A task distribution  $P \sim q$  is drawn from meta-distribution  $q$ . Given a  $P$ , the training dataset  $X = ((Z_i, Y_i))_{i=1}^n$  and the test sample  $(Z, Y)$  are drawn as  $X \sim P^{\otimes n}$  and  $(Z, Y) \sim P$ . Here,  $Z$  is the input (i.e., the sequence of symbols) and  $Y$  is the label (i.e., the next symbol to be predicted). The overall error of a learning algorithm  $\mathcal{A}$  on the meta-distribution  $q$  with sample size  $n$  is defined by
$$
\operatorname{err}_{q,n}(\mathcal{A}) = \operatorname *{Pr}_{\substack{P\sim q,X\sim P^{\otimes n},\\ (Z,Y)\sim P}}(M(Z)\neq Y\text{where} M = \mathcal{A}(X)).
$$
Given  $q$  and  $n$ , there exists an optimal learner  $\mathcal{A}_{\mathrm{OPT}}$  that minimizes this overall error, which will be used as our reference. We adopt the abstracted language problem, i.e., the next-symbol prediction problem with  $N$  reference strings  $\{c_j\}_{j=1}^N$  where  $c_j \sim \text{Uniform}(\{0,1\}^d)$ , considered in the main
text of Brown et al. [4] with no symbol corruption (see [4] or Appendix A.1 for the details). Under this setting, Brown et al. [4] proved that any algorithm  $\mathcal{A}$  needs to memorize the  $nd$  bits of training data to achieve  $\epsilon$ -suboptimality where  $I$  denotes the mutual information:
Theorem 1 (Brown et al. [4]). Let  $N = n$ . Then, any learning algorithm  $\mathcal{A}$  that satisfies  $\mathrm{err}_{q,n}(\mathcal{A}) \leq \mathrm{err}_{q,n}(\mathcal{A}_{\mathrm{OPT}}) + \epsilon$  for  $\epsilon = o(1)$  also satisfies  $I(X; \mathcal{A}(X)|P) = \Omega(nd)$ .
3.2 Memorization with Knowledge-Augmentation
In Theorem 1,  $d$  corresponds to the size of KB. Thus, it shows that if the size of KB is small, then a small model can just memorize all KB by memorizing  $\Omega(nd)$  information to perform well. However, if the size of KB is large, then a small model cannot memorize  $\Omega(nd)$  information and hence the performance is expected to drop significantly when replacing a large model with a small model. In this subsection, we show that knowledge-augmentation reduces the memorization requirement of  $\Omega(nd)$  bits to that of  $O(n\log_2(N + R))$  bits, allowing the use of small models.
We consider an inference algorithm  $\varphi$  that uses a KB with a non-parametric retriever as follows:
$$
\operatorname{err}_{q,n}^{\varphi}(\mathcal{A}) = \operatorname *{Pr}_{\substack{P\sim q,X\sim P^{\otimes n},\\ (Z,Y)\sim P}}\left(\varphi (Z,M,S)\neq Y\text{where} M = \mathcal{A}(X)\right).
$$
An inference algorithm  $\varphi$  has no learnable parameters and makes prediction based on both the result of learning algorithm  $M = \mathcal{A}(X)$  and a KB denoted by  $S$ , which is defined as follows. Given a task instance  $P \sim q$ , we choose a KB such that  $|S| = N + R$  and  $\{c_j\}_{j=1}^N \subseteq S$  where  $R$  is the number of extra references that are irrelevant to this task  $P$ ; i.e.,  $R = 0$  in the best scenario.
Theorem 2 shows that the knowledge-augmentation reduces the amount of memorization to achieve  $\epsilon$ -suboptimality, from the  $nd$  to  $\min(N, n)m$  bits, under the same problem setting as Theorem 1:
Theorem 2. There exists a pair of inference and learning algorithms  $(\varphi, \mathcal{A})$  such that for any  $\epsilon > 0$ ,  $\mathrm{err}_{q,n}^{\varphi}(\mathcal{A}) \leq \mathrm{err}_{q,n}(\mathcal{A}_{\mathrm{OPT}}) + \epsilon$  and  $I(X; \mathcal{A}(X)|P) = O(\min(N, n)m)$  where  $m = \log_2((1 - (\frac{N-1}{N})^n) \frac{(N+R)^2 - (N+R)}{2\epsilon}) \leq \log_2(\frac{(N+R)^2}{2\epsilon})$ .
With  $n = N$  and  $\epsilon = o(1)$ , we have  $I(X; \mathcal{A}(X)|P) = O(\min(N, n)m) = O(n\log_2(N + R))$  (see Appendix A.2 for proof). Thus, it shows that knowledge-augmentation allows the reduction from the  $nd$  bits to  $n\log_2(N + R)$  bits for the amount of memorization needed to perform well.
4 Knowledge-Augmented Reasoning Distillation
We propose Knowledge-Augmented Reasoning Distillation (KARD), which consists of two learning processes: (1) reasoning distillation where we leverage Large Language Models (LLMs) to generate a rationale with black-box APIs and then fine-tune small models to generate both rationale and answer given a question and knowledge, in which the knowledge is retrieved from Knowledge Base (KB) with the rationale as a query; (2) reranker training to retrieve relevant passages for the question as a query at the inference time, for generating effective rationales. Our approach is illustrated in Figure 2.
4.1 Teach Small Models to Generate Rationales with External Knowledge
Rationale Generation with LLMs In our problem setup, we assume that training dataset  $((\pmb{x}_i,\pmb{y}_i))_{i = 1}^n$  for the target task is given, where  $\pmb{x}_i$  is input sequence (question in QA) and  $\pmb{y}_i$  is label (answer in QA). Additionally, there are LLMs accessible through black-box APIs [5, 8, 44, 43, 42]. In other words, the parameters and the architecture of the LLM are unknown and we can only access text sequences generated by the LLM. Since the ability to generate high-quality rationale is known as the emergent ability of LLMs [55, 28], we want to transfer such ability to a small language model with reasoning distillation. Firstly, we leverage the chain-of-thought prompting [28] to elicit the proper  $l$  rationales for each training data point with LLMs:  $\pmb{r}_{ij} = \mathsf{LLM}(\pmb {p},\pmb{x}_i,\pmb{y}_i)$  for all  $i\in [n]\coloneqq \{1,\dots ,n\}$  and  $j\in [l]$ , where  $\pmb{r}$  is the generated rationale and  $\pmb{p}$  is the chain-of-thought prompt [55, 28, 48].
Fine-tuning Small Models on Rationales Then we fine-tune a small language model  $p_{\theta}$  with trainable parameters  $\theta$  to generate both rationale  $r_{ij}$  obtained from the LLM and answer  $y_{i}$ , given the question  $x_{i}$ . In other words, we minimize the negative log-likelihood of the sequence of rationale  $r_{ij}$  and the answer  $y_{i}$  where the rationale must be generated first prior to the answer generation:
$$
\mathcal {L} _ {\text {d i s t i l l}} (\theta) = - \frac {1}{n \cdot l} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {l} \log p _ {\theta} \left(\boldsymbol {r} _ {i j}, \boldsymbol {y} _ {i} \mid \boldsymbol {x} _ {i}\right). \tag {1}
$$
Figure 2: Overview of KARD. (Left, § 4.1) Illustration of training (top) and inference (bottom) of knowledge-augmented reasoning distillation, where, during training, the small LM learns to generate rationales given the training data and the retrieved knowledge by the rationale. (Right, § 4.2) Illustration of reranker training (top) and inference (bottom). Reranker learns to prioritize the passage which has knowledge relevant to the rationale.
Intuitively, the rationale provides a deeper and more comprehensive understanding of the reasoning behind the answer associated with the question, which better guides the small model to correctly answer the question [18]. Although previous works [33, 17, 38, 12, 18] have also leveraged the rationales generated by LLMs to make small models excel at diverse reasoning tasks, generating rationales for knowledge-intensive tasks with a small LM requires additional care. As previously described in Section 3, the reasoning distillation with a small model but without knowledge augmentation may degrade the quality of the rationale generation due to the incapability of memorizing training data with the small model [57, 27] and the necessity of the memorization for better performance in language tasks [4]. Therefore, the rationale generation should be evidenced by extrinsic knowledge from external memory to enhance the capability of the small LM for generating a high-quality rationale.
Integrating External Knowledge Base Motivated by Theorem 2, we propose to retrieve a passage from an external Knowledge Base (KB) which is a corpus of over millions of documents  $\mathcal{D} = \{d_1,\dots ,d_K\}$  to support memorization capacity of the small LM. Note that the acquisition of the relevant document from KB is crucial for training the small LM to generate high-quality rationale which leads to correct answers for given questions. As done in open-domain QA task [7], we retrieve a set of relevant passages for a given query with the sparse retriever BM25 [46]. In order to obtain the document which is the most relevant to the rationale  $r_{ij}$  generated by the LLM, we utilize the rationale as a query to retrieve a set of passages  $\hat{\mathcal{D}}_{ij} = \mathrm{topk}(\rho (\pmb {d}|\pmb{r}_{ij};\mathcal{D}),k)\subset \mathcal{D}$ , where  $\rho$  denotes a retriever scoring the document  $d\in \mathcal{D}$  based on relevance to the query  $r_{ij}$  and topk yields the  $k$  passages with the top-  $k$  highest relevance scores. Finally, we utilize the retrieved documents  $\hat{\mathcal{D}}_{ij}$  for fine-tuning the small LM to generate the rationale  $r_{ij}$  and answer  $\pmb{y}_i$  for the question  $\pmb{x}_i$  as follows:
$$
\mathcal {L} _ {\text {d i s t i l l - K B}} (\theta) = - \frac {1}{n \cdot l} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {l} \log p _ {\theta} \left(\boldsymbol {r} _ {i j}, \boldsymbol {y} _ {i} \mid \boldsymbol {x} _ {i}, \hat {\mathcal {D}} _ {i j}\right), \tag {2}
$$
where the rationale and answer are sequentially generated as we did in Equation 1.
4.2 Training Neural Reranker for Rationale Generation
The remaining issue is that we cannot use the rationale as a query at the inference time. As an alternative, we can use the question  $\pmb{x}_i$  instead of the rationale  $r_{ij}$  as a query to retrieve a set of passages with the retriever. However, there is no guarantee that the top-  $k$  passages retrieved by the input  $\pmb{x}_i$  as a query contain relevant information to generate correct rationales. In detail, based on the question as a query, the retriever can obtain a set of passages that contain relevant documents for generating rationales with a sufficiently large  $k$  but  $k \ll K$ . However, the target documents we want for rationale generation may be assigned with low rankings and thus they may not be chosen for knowledge augmentation at the inference time. To remedy this issue, we propose to leverage a neural reranker  $f_{\phi}$  [26] with parameter  $\phi$  to re-rank the set of passages retrieved by the retriever  $\rho$  so that we can acquire more relevant documents for generating rationale at the inference time.
In order to train the neural reranker, we might manually construct a ground truth passage for each question. However, we assume a realistic setting where the ground truth passage for reranker training
is not given. Instead, we train the reranker to imitate how the retriever scores the passage  $\pmb{d} \in \mathcal{D}$  with the rationale  $\pmb{r}_{ij}$  as a query. Specifically, we first utilize the retriever  $\rho$  to obtain a set of passages from  $\mathcal{D}$  with the rationale  $\pmb{r}_{ij}$  as a query as follows:  $\tilde{\mathcal{D}}_{ij} = \mathrm{topk}(\rho(\pmb{d}|\pmb{r}_{ij};\mathcal{D}),\kappa_1)\bigcup \mathrm{topk}(\rho(\pmb{d}|\pmb{x}_i;\mathcal{D}),\kappa_2)$  where  $\kappa_1$  and  $\kappa_2$  are the number of candidate documents (Figure 2 is the case where  $\kappa_2 = 0$ ). Then, we normalize the score  $\rho(\pmb{d}|\pmb{r}_{ij};\mathcal{D})$  of the document from  $\tilde{\mathcal{D}}_{ij}$ , denoted as  $Q(\pmb{d}|\pmb{r}_{ij})$ . Similarly, we use the reranker  $f_{\phi}$  to score each document in  $\tilde{\mathcal{D}}_{ij}$  with the given question  $\pmb{x}_i$  and normalize the score denoted as  $P_{\phi}(\pmb{d}|\pmb{x}_i)$ . We use softmax for normalization with hyperparameters  $\tau_1,\tau_2 > 0$  as follows:
$$
Q (\pmb {d} | \pmb {r} _ {i j}) = \frac {\exp \left(\rho (\pmb {d} | \pmb {r} _ {i j} ; \mathcal {D}) / \tau_ {1}\right)}{\sum_ {\pmb {d} ^ {\prime} \in \tilde {\mathcal {D}} _ {i j}} \exp \left(\rho (\pmb {d} ^ {\prime} | \pmb {r} _ {i j} ; \mathcal {D}) / \tau_ {1}\right)}, P _ {\phi} (\pmb {d} | \pmb {x}) = \frac {\exp (f _ {\phi} (\pmb {d} , \pmb {x} _ {i}) / \tau_ {2})}{\sum_ {\pmb {d} ^ {\prime} \in \tilde {\mathcal {D}} _ {i j}} \exp (f _ {\phi} (\pmb {d} ^ {\prime} , \pmb {x} _ {i}) / \tau_ {2})},
$$
where  $\pmb{d} \in \tilde{\mathcal{D}}_{ij}$ . Finally, we minimize the KL divergence between  $Q(\pmb{d}|\pmb{r}_{ij})$  and  $P_{\phi}(\pmb{d}|\pmb{x}_i)$ :