Teaching-Assistant-in-the-Loop: Improving Knowledge Distillation from Imperfect Teacher Models in Low-Budget Scenarios

📄 arXiv: 2406.05322v1 📥 PDF

作者: Yuhang Zhou, Wei Ai

分类: cs.CL, cs.AI

发布日期: 2024-06-08

备注: Accepted by ACL 2024 Findings

💡 一句话要点

提出Teaching-Assistant-in-the-Loop框架，提升低预算下不完美教师模型的知识蒸馏效果

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

关键词: 知识蒸馏 大语言模型 模型压缩 自我一致性 不确定性评估 助教模型 低资源学习

📋 核心要点

  1. 大语言模型（LLM）的知识蒸馏面临高昂的查询成本和教师模型输出不完美的双重挑战。
  2. 核心思想是利用学生模型的自我一致性以及助教模型评估的不确定性作为信号，辅助学生模型的训练。
  3. 实验结果表明，提出的两阶段框架在复杂推理任务上显著优于现有方法，相对提升高达20.79%。

📝 摘要（中文）

本文提出了一种三组件框架，旨在提升资源受限且教师模型不完美场景下的知识蒸馏效果。该框架利用三种信号：学生模型的自我一致性（通过多次输出的一致性来衡量置信度），以及通过引入“助教”（TA）模型评估学生和教师输出的不确定性，作为另外两种信号用于学生训练。此外，提出了一个两阶段训练方案，首先用少量数据预热学生模型，以更好地利用学生自身的信号。实验表明，在四个复杂的推理任务上，该框架优于现有方法，平均而言，相比于没有任何信号的微调，所提出的两阶段框架带来了高达20.79%的相对改进。

🔬 方法详解

问题定义：论文旨在解决低预算场景下，从不完美的教师模型（如GPT-4）中进行知识蒸馏时，学生模型训练效率低下的问题。现有方法通常需要大量的教师模型输出作为训练数据，且教师模型的不完美输出会对学生模型的学习产生负面影响。

核心思路：论文的核心思路是利用学生模型自身的知识和教师模型的辅助信息，通过引入“助教”模型来评估学生和教师输出的不确定性，并将其作为训练信号，从而提高样本效率和鲁棒性。同时，采用两阶段训练策略，先用少量数据预热学生模型，使其具备一定的自我评估能力。

技术框架：该框架包含三个主要组件：学生模型、教师模型和助教模型。整体流程如下：1) 学生模型和教师模型对输入进行预测；2) 助教模型评估学生模型和教师模型输出的不确定性；3) 利用学生模型的自我一致性和助教模型评估的不确定性作为信号，指导学生模型的训练。整个训练过程分为两个阶段：预热阶段和蒸馏阶段。

关键创新：该论文的关键创新在于引入了“助教”模型来评估学生和教师输出的不确定性，并将其作为训练信号。这种方法能够有效地利用学生自身的知识和教师模型的辅助信息，提高样本效率和鲁棒性。此外，两阶段训练策略也能够更好地利用学生自身的信号。

关键设计：助教模型的设计细节未知，论文中可能没有详细描述其结构和训练方式。损失函数的设计可能包括衡量学生模型预测结果与教师模型预测结果之间差异的损失项，以及衡量学生模型自我一致性的损失项。两阶段训练策略中，预热阶段可能使用较小的学习率和较少的数据，蒸馏阶段则使用较大的学习率和更多的数据。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在四个复杂的推理任务上，所提出的两阶段框架相比于没有任何信号的微调，带来了高达20.79%的相对改进。这表明该框架能够有效地提高知识蒸馏的效率和鲁棒性，尤其是在教师模型不完美和资源受限的场景下。

🎯 应用场景

该研究成果可应用于各种需要将大型语言模型的知识迁移到小型模型的场景，例如在资源受限的边缘设备上部署智能应用、构建特定领域的知识库等。通过提高知识蒸馏的效率和鲁棒性，可以降低模型部署的成本，并提升模型的性能。

📄 摘要（原文）

There is increasing interest in distilling task-specific knowledge from large language models (LLM) to smaller student models. Nonetheless, LLM distillation presents a dual challenge: 1) there is a high cost associated with querying the teacher LLM, such as GPT-4, for gathering an ample number of demonstrations; 2) the teacher LLM might provide imperfect outputs with a negative impact on the student's learning process. To enhance sample efficiency within resource-constrained, imperfect teacher scenarios, we propose a three-component framework leveraging three signal types. The first signal is the student's self-consistency (consistency of student multiple outputs), which is a proxy of the student's confidence. Specifically, we introduce a ``teaching assistant'' (TA) model to assess the uncertainty of both the student's and the teacher's outputs via confidence scoring, which serves as another two signals for student training. Furthermore, we propose a two-stage training schema to first warm up the student with a small proportion of data to better utilize student's signal. Experiments have shown the superiority of our proposed framework for four complex reasoning tasks. On average, our proposed two-stage framework brings a relative improvement of up to 20.79% compared to fine-tuning without any signals across datasets.