Mentor- KD : Making Small Language Models Better Multi-step Reasoners (2024)

Abstract: Large Language Models (LLMs) have displayed remarkable performances across various complex tasks by leveraging Chain-of-Thought (CoT) prompting. Recently, studies have proposed a Knowledge Distillation (KD) approach, reasoning distillation, which transfers such reasoning ability of LLMs through fine-tuning language models of multi-step rationales generated by LLM teachers. However, they have inadequately considered two challenges regarding insufficient distillation sets from the LLM teacher model, in terms of 1) data quality and 2) soft label provision. In this paper, we propose Mentor-KD, which effectively distills the multi-step reasoning capability of LLMs to smaller LMs while addressing the aforementioned challenges. Specifically, we exploit a mentor, intermediate-sized task-specific fine-tuned model, to augment additional CoT annotations and provide soft labels for the student model during reasoning distillation. We conduct extensive experiments and confirm Mentor-KD’s effectiveness across various models and complex reasoning tasks.

## introduction large language models (llms) have shown impressive emergent capabilities, showing their competence on a variety of reasoning tasks in the natural language processing (nlp) landscape @xcite @xcite . one particularly interesting strategy for this approach is chain-of-thought (cot) prompting, which elicits multi-step reasoning abilities of llms by explicitly generating intermediate reasoning steps for complex tasks @xcite . however, such reasoning abilities have been shown to only manifest in language models (lms) with over hundreds of billion parame-ters @xcite , which require significant computational resources or expensive api calls, restricting their deployment on resource-limited scenarios. to circumvent these deployment challenges, previous works @xcite @xcite have followed a knowledge distillation (kd) approach, reasoning distillation, which transfers the multi-step reasoning ability of llms to small lms. the kd pipeline generally applies in-context learning (icl) on the llm teacher model to generate outputs (e.g., multi-step rationales) as distillation sets, and then utilizes them to fine-tune the student model. previous studies have shown that reasoning distillation can significantly improve student performances and may even outperform their llm teachers on specific tasks @xcite . however, previous approaches to reasoning distillation have two challenges arising from insufficient distillation sets generated by llm teachers. first, as llms may not have access to task-specific data, the quality of the rationales for distillation can be low (e.g., only 58% accuracy on gpt-3.5 for strategyqa). the low quality of llm teacher rationales limits the number of reasoning rationales to only a small set of correct ones due to the exclusion of incorrect rationales that negatively affect student performances @xcite . second, because accessibility of black-box llm teachers is generally restricted, the student model cannot mimic the predictive behavior and knowledge from the soft labels @xcite . such oversights may lead to the student model being over-fitted on limited distillation sets from teacher models and undermine its generalization capabilities. to address these challenges, we propose mentor-kd, a novel reasoning distillation framework that effectively distills the multi-step reasoning capability of llms. our core idea is to introduce a mentor, an intermediate-sized task-specific model, that complements the llm teacher's knowledge during reasoning distillation. to this end, we first fine-tune the mentor models on specific tasks and generate both cot rationales and soft labels to augment distillation sets. by leveraging task-specific mentors whose power is concentrated toward a specific target ability, mentor-kd effectively addresses two issues through training on more diverse rationales and intrinsic knowledge from soft labels. we conduct extensive experiments on various types of complex reasoning tasks, including commonsense, arithmetic, logical, and symbolic reasoning tasks. the experimental results clearly demonstrate the superiority of our method over baselines leveraging knowledge only from llms. in addition, we verify that the mentor model can generate a substantial number of correct reasoning samples compared to other llm baselines, highlighting the effectiveness of our method as means of data augmentation. lastly, we demonstrate that our mentor-kd significantly improves student performances in low-resource scenarios, indicating its cost-efficiency. in summary, the contributions of this paper include the following: ## methodology we elaborate on the detailed implementations of our mentor-kd. the core idea is to augment the distillation training set by leveraging a task-specific intermediate-sized mentor model. to this end, we first generate cot annotations from llm teacher models (section 3.1). we then fine-tune the mentor model with the distillation set from the llm teacher, and the trained mentor model generates additional training sets, including both rationales and soft labels (section 3.2). by augmenting both signals from the mentor, we distill the knowledge to student models (section 3.3). figure 2 illustrates an overview of our framework. ## experiments in this section, we describe the experiment details and evaluate our mentor-kd on various complex reasoning tasks. ## analysis to delve into the benefits of our method, we perform a series of fine-grained analytical experiments with the following research questions (rqs): • rq1. can mentor-kd be generalized to the various sizes and types of student models? ( §5.1) • rq2. how does each component in mentor-kd contribute to its overall performance? ( §5.2) ## conclusion we have presented mentor-kd, a novel framework to transfer reasoning capabilities from llms to smaller lms. to this end, we have introduced a mentor model, a novel auxiliary model, for complementing the distillation sets from llms by augmenting multi-step rationales and providing soft labels for the student model. through extensive experiments, we have demonstrated that our mentor-kd significantly improves the effectiveness of reasoning distillation. specifically, our student models outperform existing reasoning distillation baselines with various sizes and types of models on complex reasoning tasks. furthermore, we have verified that our mentor model can generate effective reasoning samples and soft labels for training student models, resulting in consistent performance improvements.