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

TL;DR

This paper exploits 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, which effectively distills the multi-step reasoning capability of LLMs to smaller LMs while addressing the aforementioned challenges.

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.

Figures (7)

• Figure 1: Comparison between (a) previous approaches of reasoning distillation and (b) Mentor-KD (ours). Our framework utilizes an intermediate-sized task-specific mentor model to complement the distillation sets of teachers.
• Figure 2: A general overview of our proposed framework, Mentor-KD. Mentor-KD is composed of three steps. First, CoT annotations are initially collected from the teacher LLM and filtered. Second, the preserved annotations are used to train the mentor model, and the trained mentor model augments multi-step rationales. Lastly, the student model is trained on annotations from the teacher and the student, as well as soft labels from the mentor model.
• Figure 3: Performances by differentiating the degree (number) of mentor-generated CoT rationales per question. We adopt FlanT5-large and FlanT5-small as mentor and student models, respectively.
• Figure 4: Comparison of (a) accuracy of our mentor model (FlanT5-large) and LLM baselines on teacher-incorrect samples, and (b) performances of student models trained with augmented distillation sets from LLM baselines and our mentor models.
• Figure 5: Comparison between Mentor-KD (Ours) and Vanilla-KD baseline on various distillation sets by differentiating the percentage of rationales being used.