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

大模型向量数据库数据中台

备注好友:

方向-学校/公司-姓名/昵称

【AIGC 先锋科技】交流群

picture.image

picture.image

大型语言模型(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在低资源场景下显著提高了学生的性能,表明了其成本效益。

总的来说,本文的贡献包括以下几点:

  1. 作者提出了一种名为Mentor-KD的新颖推理蒸馏框架,该框架在考虑语言模型教师提供的不足的蒸馏集限制下,提高了小型语言模型的推理能力。
  2. 作者提出了一种导师模型,可以生成补充LLM教师有限训练数据的有理性和软标签样本。
  3. 作者通过大量实验表明,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说明了作者的框架概述。

picture.image

Chain-of-Thought Annotations

作者使用LLM通过零样本CoT(Kojima等人,2022年)从CoT标注中获取答案,这些标注由一个推理和一个最终预测组成。这种方法采用两阶段策略,包括推理和答案提取阶段,因此作者首先诱导LLM生成一个CoT推理,然后生成一个最终预测。

具体而言,作者首先将"Let's think step by step"添加到问题中, Prompt LLM获取推理。接下来,作者通过将已获得的推理纳入 Prompt ,再次 Prompt LLM以诱导其最终预测。正式而言,从一个数据集

中,其中

表示一个问题,

表示黄金标签,作者的目标是使LLM生成一个分步推理

和最终预测

,给定

作为输入。 Prompt 模板的形式为:"Q:

{r^{t}_{i}}. 因此,答案是

"。

后来,作者过滤由LLM生成的标注。参考Li等人(2023);Magister等人(2023);Fu等人(2023);Lee等人(2024)之前的工作,作者在最终预测

与样本的黄金答案

相匹配的标注上保留标注。然后,将标注重新格式化为问题和标签格式,以便用于训练导师和学生模型。更正式地,对于所有

的标注

,作者将一个数据样本

重新格式化为

,其中

采用 "

" 的形式。因此,答案为 "

"。因此,作者最后构建了

Mentor Model

在这里,作者描述了如何训练作者的导师模型以集中力量于特定任务,并利用它们来补充LLM教师的不足之处。

训练.对于训练导师模型,作者直接在之前构建的

上进行微调。具体来说,导师模型接收

作为输入,

作为标签,并采用标准的语言建模目标进行训练。

合理增强。然后,训练好的导师模型用于训练数据增强。对于来自

的数据样本,作者让导师模型逐步标注理由,以

作为输入。导师模型返回一个标签

,它包括一个逐步理由和对其本身的预测。作者通过像过滤教师的标注一样过滤导师的标注,并保留

的数据样本。通过这一阶段,作者为每个数据集构建

结合来自导师(

)和导师(

)的标注,作者最终构建了用于训练学生模型的

,其定义如下:

Reasoning Distillation

为了训练学生模型,作者通过通过导师模型(软标签蒸馏)获得的对数值,同时融合微调(论据蒸馏)和知识蒸馏。这样可以使学生模型共同学习如何在符号形式逐步推理 Ho等人(2023);李等人(2023);马吉斯特等人(2023),以及模仿更大模型的预测行为 Hinton等人(2015)。相应地,作者的训练目标包括两个损失函数。

合理蒸馏。与训练导师模型相同,通过微调学生模型与教师和导师提供的问标签对,可以逐步提取学生的推理能力。更具体地说,通过微调学习多步推理能力的方式如下定义:

学生在进行AI学习时,通常采用模型

,而方括号则表示字符串的连接。

软标签蒸馏。由于其黑盒性质或巨大尺寸,利用LLM教师的内部知识在实际中可能不切实际。因此,作者采用作者的导师模型来提供软标签。软标签通过前向传播和softmax函数获得,其中q作为输入。正式地,作者从第k个位置的logit值z_k处,通过以下方程获得导师和学生的软标签(概率分布)p_k:

其中

表示软化分布的温度超参数。在获得导师(

)和学生(

)的概率分布之后,作者采用Kullback-Leibler散度损失来最小化两个分布之间的散度。这使得学生模型能够模拟预测行为并学习大型模型的内部知识。软标签蒸馏的训练目标定义如下:

联合学习。最后,作者有了学生模型,它将联合学习上述两个目标。训练学生模型的损失函数如下:

λ 是用于插值两个损失函数的超参数。

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在解决复杂的推理任务方面具有有效性,包括算术推理和常识推理。

picture.image

5 Analysis

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

  1. RQ1. mentor-KD 是否可以推广到各种大小和类型的学生模型?(SS5.1)
  2. RQ2: Mentor-KD中的每个组件如何共同贡献其整体性能?(SS5.2)
  3. RQ3: 导师模型能否为学习者生成有用的浓缩集?(SS5.3)
  4. RQ4: Mentor-KD在低资源场景下是否提供改进?(SS5.4)
  5. 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教师的不足知识,从而通过传递更多的信息浓缩信号,使各种学生模型获得更好的性能。

picture.image

picture.image

Ablation Studies (RQ2)

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

picture.image

Impact of Data Augmentation (RQ3)

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

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

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

picture.image

增强数据集的质量。为了调查作者增强的数据集的质量,作者将作者的导师模型(即 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教师更高的质量的合理性。总体结果强调了针对特定任务的导师模型微调的优势。

picture.image

Low-resource Scenarios (RQ4)

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

picture.image

作者观察到,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)。

picture.image

6 结论

作者已经提出了Mentor-KD,这是一个新颖的框架,可以将LLM的推理能力转移到较小的LLM中。为此,作者引入了一个mentor模型,这是一个新颖的辅助模型,通过增加多步推理理由并为学生模型提供软标签来补充LLM的蒸馏集。

通过广泛的实验,作者证明了作者的Mentor-KD显著提高了推理蒸馏的有效性。

具体来说,作者的学生模型在复杂推理任务上超过了各种大小和类型的现有推理蒸馏 Baseline 。

此外,作者还验证了作者的mentor模型可以生成有效的推理样本和软标签用于训练学生模型,从而导致一致的性能改进。

参考文献

[0]. Mentor-KD: Making Small Language Models Better Multi-step Reasoners.

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论