【EMNLP2024】基于多轮课程学习的大语言模型蒸馏算法 TAPIR

近日，阿里云人工智能平台PAI与复旦大学王鹏教授团队合作，在自然语言处理顶级会议EMNLP 2024 上发表论文 《Distilling Instruction-following Abilities of Large Language Models with Task-aware Curriculum Planning》。文章提出了一个名为 TAPIR 的知识蒸馏框架，TAPIR 通过多任务课程规划来蒸馏黑盒大语言模型的指令回答能力，在蒸馏和多轮迭代过程中，使用教师 LLM 做为裁判找出对于学生 LLM 来说难以回答的指令，进行难度重采样。同时，TAPIR 调整多任务配比，进行训练集中的任务多样性分布的重采样，并根据相应多任务特点自动优化教师模型的回答风格。

背景

大语言模型在回答开放领域通用任务的指令上取得了很大地进步。指令微调是微调预训练模型，使其从文本补全模型成为强大的对话模型的关键。尽管已有研究探索了使用强大的黑盒教师模型（如GPT-4， Qwen-max）来自动蒸馏和标注指令的方法，但这些研究往往忽视了微调训练集中任务的多样性分布，以及训练集中指令难度的差异，这可能导致学生LLMs知识能力的不平衡和解决复杂任务的能力的不足。为了解决这些挑战，文章提出了一个名为TAPIR的新框架，它通过多任务课程规划来蒸馏黑盒大语言模型的指令回答能力，从而提高学生小模型的指令遵循能力。

算法流程

文章中提出的TAPIR（Task-Aware Curriculum Planning for Instruction Refinement）框架的算法流程是一个多轮次的蒸馏方法，旨在提升学生大型语言模型（LLMs）遵循指令的能力。整个流程从初始化一个预训练的学生模型开始，然后通过以下步骤进行：

1. 数据集难度过滤：使用一个开源的指令数据集（如Alpaca数据集）作为基础，通过计算模型拟合难度（MFD）分数来筛选出对学生模型来说较难的指令对，过滤得到种子数据集。
2. 多任务规划指令蒸馏：根据设定的任务类型配比，利用一个教师模型（如ChatGPT）扩展种子数据集，生成更多具有相似难度水平的指令-响应对，并提升推理类任务的采样概率，以更好的缓解能力冲突问题。
3. 多任务回答风格增强：对于某些任务，使用特定的提示重写响应，以便从教师模型获得更精细、更详细的回答，或者是特定任务格式的回答（如思维链，代码注释），这有助于学生模型更好地理解和学习复杂任务。
4. 模型多轮优化迭代：通过多轮训练，利用裁判模型得到学生模型的回答质量反馈奖励分数，采样得到新的蒸馏种子数据集。逐步增加新一轮蒸馏种子数据集中挑战性指令的比例，实现从易到难的泛化。

TAPIR框架通过这种逐步提升任务难度和均衡任务类型的策略，使学生模型能够在较少的训练数据下超越更大的模型，显示出更好的性能，并在多个基准测试中取得了显著的性能提升。

难度重采样

难度重采样旨在解决训练集中任务难度分布不均的问题。难度重采样的目标是确保学生大型语言模型在蒸馏微调过程中能够接触到难度逐渐增大的任务，从而在困难的任务中泛化。我们通过计算模型拟合难度（Model Fitting Difficulty, MFD）分数来评估每个指令对学生模型的难度。MFD分数是通过比较学生模型生成的响应与教师模型生成的响应之间的质量差异来确定的。我们使用教师模型来作为裁判打分。

根据MFD分数，筛选出对学生模型来说较难的指令对，即分差大于阈值的指令。这些指令对将被纳入种子数据集。

任务重采样

在TAPIR框架中，任务重采样旨在解决训练集中任务分布不均的问题。其目的是提升训练集的多样性。在均衡的任务配比下为微调学生模型，以缓解微调过程中的能力冲突和灾难性遗忘问题。

首先，我们训练了一个指令任务分类模型（Deberta v3）识别和分类训练集中的任务类型，给每条指令打上显示的任务标签。然后通过任务标签重采样，使数据集中的任务分布更均衡，并且增强逻辑推理和编程任务的占比。基于我们的采样概率，教师模型扩展种子数据生成了新指令问答对，这些新数据与原有数据在难度上相近。

设指令对的任务采样概率为，则学生模型微调的自回归损失可以写作：

我们针对任务特点增强了教师模型标注的回答格式。如下所示：

多轮迭代优化

在多轮迭代的过程中，我们可以不断更新计算学生模型在新的微调数据集上的模型拟合难度来动态调整新一轮的蒸馏种子数据集难度配比。如下面的公式所示，当 设置为 1 时，整个训练语料库仅由这些“困难”样本作为种子蒸馏。通过逐渐增加 , 系统地提高学习任务的复杂性。同时，为了保证指令的多样性，在每一轮中通过教师模型扩展难度重采样后的数据集，并将扩展后的数据集表示为。第 r 轮的损失函数定义如下：

在每轮之后，更新规则为：

其中是一个预定义的常数，用来逐渐增加学习任务的难度。

实验结果

实验结果表明，使用TAPIR框架训练的学生语言模型在较少的训练数据下，其性能超过了更大的指令调整模型和其他蒸馏基线方法。具体地说，TAPIR训练的模型在AlpacaEval 2.0基准测试中取得了7.80的胜率，优于Vicuna 13B和LLaMA2-Chat 13B，即便其培训数据和参数量仅为后者的一半。此外，在MT-Bench基准测试中，该模型在角色扮演、推理、数学、编程和人文学科等子任务中，表现优于LLaMA2 7B Chat 基线模型。为验证TAPIR框架在不同规模模型上的一致性，我们在Qwen1.5-Chat系列模型上做了实验，结果显示TAPIR能有效提升模型的指令遵循能力。

TAPIR-7B模型例子如下所示。在角色扮演任务中，语言模型扮演体育解说评论员。TAPIR-7B 生动地描述了比赛的最后胜利时刻并表现出色，而 Lion-7B 只是提供了如何评论的分析，没有完全执行任务，LLaMA2-Chat则误解了指令。

参考文献

• Li, M., Chen, L., Chen, J., He, S., Huang, H., Gu, J., & Zhou, T. Reflection-Tuning: Data Recycling Improves LLM Instruction-Tuning. ArXiv, abs/2310.11716.
• Song, C., Zhou, Z., Yan, J., Fei, Y., Lan, Z., & Zhang, Y. Dynamics of Instruction Tuning: Each Ability of Large Language Models Has Its Own Growth Pace. ArXiv, abs/2310.19651.
• Jiang, Y., Chan, C., Chen, M., & Wang, W. Lion: Adversarial Distillation of Proprietary Large Language Models. EMNLP 2023.

论文信息

论文名字：Distilling Instruction-following Abilities of Large Language Models with Task-aware Curriculum Planning
论文作者：岳元浩、汪诚愚、黄俊、王鹏
论文pdf链接：https://arxiv.org/pdf/2405.13448

阿里云人工智能平台PAI长期招聘研究实习生。团队专注于深度学习算法研究与应用，重点聚焦大语言模型和多模态AIGC大模型的应用算法研究和应用。简历投递和咨询：[e-mail]。