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

📄 arXiv: 2410.09037v1 📥 PDF

作者: Hojae Lee, Junho Kim, SangKeun Lee

分类: cs.CL, cs.AI

发布日期: 2024-10-11

备注: EMNLP 2024

💡 一句话要点

提出Mentor-KD，通过中间导师模型提升小语言模型的多步推理能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 多步推理 思维链 语言模型 软标签

📋 核心要点

  1. 现有推理蒸馏方法在利用LLM教师模型时，面临数据质量和软标签提供不足的挑战。
  2. Mentor-KD利用中间规模的、特定任务微调的“导师”模型，生成额外的CoT标注并提供软标签。
  3. 实验结果表明，Mentor-KD能够有效提升小语言模型在复杂推理任务上的性能。

📝 摘要（中文）

大型语言模型(LLMs)通过思维链(CoT)提示在各种复杂任务中表现出卓越的性能。最近，一些研究提出了知识蒸馏(KD)方法，即推理蒸馏，通过微调由LLM教师模型生成的多步推理过程的语言模型，来迁移LLM的这种推理能力。然而，它们没有充分考虑到来自LLM教师模型的蒸馏数据集不足的两个挑战，即1)数据质量和2)软标签提供。在本文中，我们提出了Mentor-KD，它有效地将LLM的多步推理能力提炼到较小的LM，同时解决了上述挑战。具体来说，我们利用一个中间规模的、特定于任务的微调模型（导师），来扩充额外的CoT注释，并在推理蒸馏期间为学生模型提供软标签。我们进行了广泛的实验，并证实了Mentor-KD在各种模型和复杂推理任务中的有效性。

🔬 方法详解

问题定义：现有方法在知识蒸馏中，直接使用大型语言模型（LLM）作为教师模型，生成用于训练小型语言模型（SLM）的推理数据。然而，LLM生成的数据可能存在质量问题，且难以提供高质量的软标签，限制了SLM推理能力的提升。

核心思路：Mentor-KD的核心思想是引入一个中间规模的、特定任务微调的“导师”模型，作为LLM和SLM之间的桥梁。导师模型负责生成高质量的CoT标注数据，并为SLM提供更准确的软标签，从而提升SLM的推理能力。这种设计旨在弥补LLM直接蒸馏的不足，提高蒸馏效率。

技术框架：Mentor-KD的整体框架包含以下几个主要步骤：1) 使用LLM生成初始的CoT数据；2) 使用特定任务的数据微调一个中间规模的模型作为导师模型；3) 导师模型利用初始CoT数据进行增强，生成更丰富、更高质量的CoT标注数据；4) 导师模型为SLM提供软标签；5) 使用增强的CoT数据和软标签训练SLM。

关键创新：Mentor-KD的关键创新在于引入了中间“导师”模型，该模型通过特定任务的微调，能够生成更符合任务特点、质量更高的CoT数据，并提供更准确的软标签。这与直接使用LLM进行蒸馏的方法不同，后者可能由于LLM的通用性而无法提供针对特定任务的最佳指导。

关键设计：导师模型的选择至关重要，需要根据具体任务选择合适的模型结构和规模，并使用相关数据进行充分的微调。软标签的生成方式也需要仔细设计，例如可以使用导师模型的输出概率分布作为软标签。此外，损失函数的设计也需要考虑软标签的特性，例如可以使用KL散度损失来衡量SLM和导师模型输出分布的差异。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Mentor-KD在多个复杂推理任务上显著优于直接使用LLM进行蒸馏的方法。例如，在某些任务上，Mentor-KD可以将小模型的性能提升超过10%。此外，实验还验证了Mentor-KD在不同模型架构和数据集上的泛化能力。

🎯 应用场景

Mentor-KD可应用于各种需要多步推理能力的场景，例如问答系统、对话系统、代码生成等。通过将大型语言模型的推理能力迁移到小型模型，可以在资源受限的环境中部署高性能的推理系统，例如移动设备或边缘计算设备。该方法还可以用于构建更高效、更可解释的AI系统。

📄 摘要（原文）

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.