Mentor-KD 方法解决LLM推理蒸馏挑战，高效将多步推理能力灌输给小模型！

AIGC 先锋科技

发布于 2025-02-28 14:53:32

910

大型语言模型（LLMs）通过链式思维（CoT） Prompt 在各种复杂任务上表现出令人瞩目的性能。最近，有研究提出了知识蒸馏（KD）方法，即推理蒸馏，通过微调LLM教师生成的多步推理语言模型的语言模型，将LLM的推理能力进行传递。然而，他们没有充分考虑LLM教师模型中存在的两个挑战，即1）数据质量和2）软标签提供。在本文中，作者提出了一种名为Mentor-KD的方法，该方法在解决上述挑战的同时，有效地将LLM的多步推理能力蒸馏到较小的语言模型中。具体而言，作者利用一个导师，中间大小的特定任务微调模型，在推理蒸馏过程中为学生的模型增加额外的CoT标注并为其提供软标签。作者进行了大量实验，并确认了Mentor-KD在各种模型和复杂推理任务中的有效性。

1 Introduction

大型语言模型（LLMs）展示出了惊人的新兴能力，在自然语言处理（NLP）领域中展现出了多种推理任务的能力。Brown等人（2020年）、Rae等人、Hoffmann等人和Chowdhery等人的研究都证实了这一点。这种方法中一个特别有趣的方法是连续思维（CoT） Prompt ，通过明确生成复杂任务的中间推理步骤来诱发LLM的多步推理能力。然而，这种推理能力只体现在具有数百亿参数的语言模型（LMs）上，这需要大量的计算资源或昂贵的API调用，限制了它们在资源受限场景中的部署。

为应对这些部署挑战，之前的工作Ho等人；Li等人（2023）；Magister等人（2023）采用了知识蒸馏（KD）方法，即推理蒸馏，将LLM的多步推理能力转移到小型的LLM中。KD Pipeline 通常在LLM教师模型上应用即席学习（ICL）来生成输出（例如，多步推理），并将其作为蒸馏集，然后利用它们来微调学生模型。

以前的研究表明，推理蒸馏可以显著提高学生性能，甚至可能在特定任务上超越他们的LLM教师Ho等人；Chen等人。

然而，推理蒸馏方法之前存在两个由LLM教师生成的不足蒸馏集所带来的挑战。首先，由于LLM可能没有访问特定任务的数据，推理集的质量可能较低（例如，在GPT-3.5上的策略性问答的准确率仅为58%）。

LLM教师的推理质量有限制，因为它们排除了那些负面影响学生表现的错误推理，这些错误推理来自于Ho等人（2023年）的研究。其次，由于黑盒LLM教师的可用性通常受到限制，学生模型无法模拟来自软标签的预测行为和知识。这样的疏忽可能导致学生模型过度适应于来自教师模型的有限蒸馏集，从而损害其泛化能力。

为了应对这些挑战，作者提出了一种名为Mentor-KD的新颖推理浓缩框架，该框架可以有效地浓缩LLM的多步推理能力。作者的核心思想是引入一个mentor，一个中间大小的特定任务模型，在推理浓缩过程中补充LLM教师的知识。为此，作者首先在特定任务上微调mentor模型，并生成CoT合理性和软标签来扩充浓缩集。通过利用特定任务的mentor，其能力集中于特定目标能力，Mentor-KD通过在更多样化的合理性和软标签的训练以及内在知识的利用，有效地解决了两个问题。

作者在各种复杂的推理任务上进行了广泛的实验，包括常识推理、算术推理、逻辑推理和符号推理任务。实验结果明确显示，作者的方法在利用仅限于LLM的 Baseline 上的知识方面具有优越性。

此外，作者还验证了mentor模型可以生成比其他LLM Baseline 更多的正确推理样本，突显了作者方法作为数据增强的有效性。最后，作者展示了作者的mentor-KD在低资源场景下显著提高了学生的性能，表明了其成本效益。

总的来说，本文的贡献包括以下几点：

作者提出了一种名为Mentor-KD的新颖推理蒸馏框架，该框架在考虑语言模型教师提供的不足的蒸馏集限制下，提高了小型语言模型的推理能力。
作者提出了一种导师模型，可以生成补充LLM教师有限训练数据的有理性和软标签样本。
作者通过大量实验表明，Mentor-KD在各种类型的推理和学生模型上，提高了推理蒸馏的有效性。

2 Related Works

Chain-of-Thought Prompting

因果 Prompt （CoT prompting）是一种通过ICL（Wei等人，2022年）来激发语言模型多步推理能力的方法。CoT的本质在于，它为语言模型提供逻辑进展的指导，帮助其分解和解决复杂的推理任务。因此，它使得语言模型在复杂的推理任务上表现出色，这些任务传统的小样本学习方法往往难以解决。近期的工作通过提高推理步骤的质量来进一步改进 CoT Prompt 。Madaan等人（2023年）通过自我反馈让语言模型迭代自 refined推理，而Gou等人（2024年）则利用外部工具获取反馈。Trivedi等人（2023年）；Zhao等人（2023年）将信息检索系统集成到语言模型推理中，以增强其推理的可信度。

尽管取得了成功，但之前的工作报告，当预训练语言模型（PLMs）扩展到数百亿参数时，基于CoT的推理的优点才会出现。为了解决这些问题，作者的工作重点是通过推理蒸馏使小规模PLMs实现CoT推理。

Knowledge Distillation for LLMs

KD (Hinton等人，2015年)被证明是一种有前景的压缩LMs的方法，通过将大型LMs的预测行为（例如软标签）或内部知识（例如隐藏表示）转移到较小规模的LMs中。然而，现有的KD方法针对预训练的LMs，涉及提取软标签（Sanh等人，2019年；Gu等人，2024年）或表示，需要访问教师模型的内部参数。这些需求对利用KD中的LLMs造成了很大的挑战，因为它们具有黑盒性质和不切实际性。

反过来，最近的研究实践了推理蒸馏，使得较小的 LMs（学生）通过使用 LLM 教师生成的理由来执行多步推理，类似于 LLMs，而不是软标签。例如，Ho 等人（2023）； Magister 等人（2023）；

李等人（2023）在 LLM 生成的多步推理的基础上微调了学生。同样，Shridhar 等人（2023）让学生通过让 LLM 生成子问题来分解一个复杂问题。王等人（2023）迭代地使用 LLM 提供了实时反馈，该反馈专门针对学生的生成。康等人（2023）；

赵等人（2024）利用信息检索系统来增强学生在知识密集型任务上的推理的真实性。最近，朱等人（2024, 2024）将多步推理以 LLM 生成的代码格式集成到提高学生的算术推理技能。

与此同时，周和爱等人（2024）也利用中等大小的模型进行 LLM 蒸馏。作者的工作不同之处在于，作者使用中等大小的模型来补充教师模型蒸馏信号，而不是用于过滤标注。

虽然大多数先前的研究都是通过利用LLM提供的知识蒸馏集来改进推理蒸馏，但作者认为这些方法可能不足够，可能会削弱学生的能力。因此，作者的工作与LLM教师在这方面存在的不足相辅相成。

3 Methodology

作者详细阐述了作者的Mentor-KD的实现细节。核心思想是利用特定任务的中间大小的mentor模型来增强蒸馏训练集。为此，作者首先从LLM教师模型中生成CoT标注（第3.1节）。然后，作者将LLM教师模型提供的蒸馏集与mentor模型进行微调，训练好的mentor模型生成额外的训练集，包括理由和软标签（第3.2节）。通过从mentor中增强两种信号，作者将知识蒸馏到学生模型（第3.3节）。图2说明了作者的框架概述。

Chain-of-Thought Annotations

Mentor Model

Reasoning Distillation

4 Experiments

在本节中，作者将描述实验的详细信息，并评估作者在各种复杂推理任务上的Mentor-KD。

Experiment Setup

任务与数据集。参考Wei等人（2022年）、Kojima等人（2022年），作者在四个复杂的推理任务类别上评估作者的Mentor-KD，这些类别包括常识推理、算术推理、逻辑推理和符号推理。具体而言，作者针对每个任务采用最多三个数据集，以便在相同任务类型的各种数据集上评估作者的框架。本文使用的数据集包括策略问答Geva等人（2021年）、常识问答Talmor等人（2019年）用于常识推理，GSM8K Cobbe等人（2021年）、ASDiv Miao等人（2020年）和SVAMP Patel等人（2021年）用于算术推理，追踪移动目标、日期理解Srivastava等人（2023年）用于逻辑推理，以及最后字母连接Wei等人（2022年）；Kojima等人（2022年）用于符号推理。更多详细信息请参见附录A。

语言模型。作者使用 OpenAI API 上的 gpt-3.5-turbo 作为作者的教师模型。对于导师和学生模型，作者主要使用 Chung 等人的 FlanT5-XXL 和 FlanT5-XL（2022）作为作者的导师和学生模型。为了进行额外的分析，作者使用了各种大小的 FlanT5 和 T5，包括大型、基础和小型模型，Raffel 等人的论文（2020）。

链式思维标注。对于GSM8K、ASDiv、SVAMP和CommonsenseQA，作者使用Chen等人（2023年）提供的CoT标注。这些标注是通过使用GPT-3.5-Turbo和零样本CoT Prompt 收集的，这与作者在第3.1节中提到的方法相同。其他数据集是由作者的研究所新 Prompt 和收集的。

Baseline 。对于 Baseline ，作者整合了之前的方法进行推理蒸馏。具体来说，作者实现了一种通用的推理蒸馏方法，即Vanilla-KD，它对学生在教师模型的生成理由Ho等人（2023）；Magister等人（2023）和MCC-KD（Chen等人，2023）上的模型进行微调。作者还比较了Mentor-KD与Fu等人（2023）的性能，该方法旨在将LM的推理能力针对特定任务进行专业化。作者通过Zero-shot-CoT（ZS-CoT） Prompt 报告教师模型的性能。

实现. 作者在两块NVIDIA RTX A6000 GPU上采用HuggingFace Wolf等人（2020）提供的模型。具体来说，作者为XXL/XL大小的模型训练了18个周期，为大型模型训练了10个周期，为 Baseline 模型（如Chen等人，2023；Ho等人，2023）训练了20个周期。在整个实验中，最大序列长度都设置为512，并遍历批处理大小{2, 4, 6, 8}。为了加速训练和节省内存使用，作者在主要实验中应用了混合精度（bfloat16和LoRA），参考了Chen等人（2023）的相关配置。此外，作者使用AdamW优化器（Loshchilov和Hutter，2019），学习率分别为{1e-4, 2e-4, 3e-4, 5e-4}。作者将损失插值超参数λ设置为0.3，并将蒸馏温度τ设置为{1.0, 2.0}。作者报告了来自四个随机种子的平均测试准确率结果。

Main Results

对于公平的比较，作者主要将Mentor-KD在三个算术推理任务和一项常识推理任务上与使用FlanT5-XL模型的MCC-KD进行比较，这些任务在推理蒸馏Ho等人（2023年）和Chen等人（2023年）中广泛使用。主要结果如表1所示。作者观察到，作者的Mentor-KD在四个不同的推理数据集上实现了最先进的性能。具体来说，作者的模型在平均准确率方面比MCC-KD提高了约2.0%。这些结果表明，Mentor-KD在解决复杂的推理任务方面具有有效性，包括算术推理和常识推理。

5 Analysis

为了深入探讨作者方法的优势，作者进行了以下一系列精细分析实验，研究问题如下（RQs）：

RQ1. mentor-KD 是否可以推广到各种大小和类型的学生模型？(SS5.1)
RQ2： Mentor-KD中的每个组件如何共同贡献其整体性能？（SS5.2）
RQ3： 导师模型能否为学习者生成有用的浓缩集？(SS5.3)
RQ4： Mentor-KD在低资源场景下是否提供改进？（SS5.4）
RQ5. 导师模型的大小是否会影响学生模型的性能？（SS5.5）

Various Student Models (RQ1)

为了进一步研究作者Mentor-KD的泛化性，作者在各种类型的学生模型上进行了实验，这些模型的规模不同。特别地，作者通过添加逻辑和符号推理任务进一步扩展了实验范围。具体来说，作者利用了广泛应用于LLM消融学习的T5和FlanT5（Ho等人，2023；Chen等人，2023）。作者将T5和FlanT5的大型变体作为作者的mentor模型，将{base, small}变体作为作者的学生模型。

结果显示在表2和表3中。作者观察到，作者的Mentor-KD在各种学生模型上的复杂推理任务的四种类别中，始终优于其他 Baseline 。特别是，Mentor-KD在常识和逻辑推理任务上取得了显著的性能改进，甚至可能超过LLM教师的性能（即GPT-3.5）。这些结果表明，作者的任务特定mentor模型可以成功地补充LLM教师的不足知识，从而通过传递更多的信息浓缩信号，使各种学生模型获得更好的性能。

Ablation Studies (RQ2)

作者进行了消融实验来探究作者方法中每种技术的贡献。具体而言，作者关注了从导师模型中提取的推理蒸馏（RD）和软标签蒸馏（SLD）的影响。作者观察到省略RD和SLD对两种模型类型和数据集都有显著影响。这些结果强调了RD对训练样本和软标签的重要性，这增强了教师提供的知识不足。

Impact of Data Augmentation (RQ3)

为了进一步研究提出的mentor模型数据增强方法的有效性，作者还从数量和质量两个角度分析了其有效性。

增强数据集的数量。作者首先分析了导师生成的推理集数量对学生的表现的影响，通过多样化导师在每题上生成的理由数量。结果如图3所示。总体上，作者观察到学生的表现随着推理集数量的增加而提高。这表明作者的导师模型成功地生成了有助于学生模型学习多步推理的论据。

然而，作者也观察到在六个增强之后，性能通常饱和，并且当引入更多的推理集时开始下降，这可能是因为模型产生的噪声（Liu等人，2022）。

增强数据集的质量。为了调查作者增强的数据集的质量，作者将作者的导师模型（即 FlanT5-large）与可能成为增强的导师模型的各种LLM进行比较（例如 GPT-3.5-Turbo3，Llama-3-8B-Instruct4 和 Vicuna-7B（Chiang等人，2023））。首先，作者比较了增强导师在其他LLM教师预测的错误样本上生成的增强的准确率与 Baseline （通过零样本-CoT Prompt ）。

然后，作者报告了在每个增强上训练的学生（即 FlanT5-small）的表现，以分析任务特定的导师是否能为学生提供有信息量的集合。

如图4所示，结果如下。尽管导师模型包含的参数比LLM（如Llama3-8B-Instruct）小（例如，比Llama3-8B-Instruct小10倍），但它们生成的合理性比其他LLM Baseline 更准确，表明能够为学生模型提供更多样化的合理性。此外，作者还观察到，使用导师模型生成的蒸馏集训练的学生确实比使用LLM教师生成的集获得更高的性能。这些结果表明，导师可以生成比LLM教师更高的质量的合理性。总体结果强调了针对特定任务的导师模型微调的优势。

Low-resource Scenarios (RQ4)

在推理蒸馏中，由于黑盒LLM的API调用成本，收集足够大的蒸馏集可能代价高昂。因此，作者研究了在资源受限的情况下Mentor-KD的有效性，即只收集原始数据集的一部分来生成蒸馏集。具体而言，作者将Vanilla-KD Baseline 与作者的Mentor-KD进行比较，变化LLM教师模型生成的蒸馏集比例。结果如图5所示。

作者观察到，Mentor-KD 也允许在资源有限的场景下提高学生模型的性能，只要导师模型提供有用的推理集和软标签。特别是，Vanilla-KD Baseline 在高限制的蒸馏信号上表现出性能退化，而作者的 Mentor-KD 表现出对有限数据集的鲁棒性。这些结果表明，作者的导师模型可以从有限蒸馏信号中缓解过拟合问题，并以成本效益的方式将 LLM 教师的知识蒸馏。作者在附录 C 中详细讨论了这个问题。

Effects of Mentor Sizes (RQ5)

为了更深入地研究Mentor-KD的有效性并验证作者的设计选择，作者通过改变mentor模型的规模进行了另一项实验。在这里，作者将FlanT5-small用作学生模型，将FlanT5-{XL, large, base, small}用作mentor模型。对于从小型到小型模型，作者采用自蒸馏，参考了先前的研究（Allen-Zhu和Li，2023；Zhu等人，2024）。

图6显示了结果。通常，作者观察到，在推理蒸馏过程中，当使用更大的导师模型时，学生模型表现更好。使用最小的导师模型会导致性能下降，但作者观察到这种场景仍然优于表3中的 Baseline 。结果表明，使用性能更好的较大模型有助于提高小学生的表现，这与以前的研究发现一致，即学生的表现与他们对应模型的表现相关（Ho等人，2023）。