将知识蒸馏 (knowledge distillation)(KD)应用于大型语言模型(LLM)时,一个重要选择是知识迁移的范围:学生模型是应在多种潜在应用中复制教师模型的通用能力(任务通用),还是应优化为在单一、特定的下游任务上模仿教师模型(任务专用)?这个选择很大程度上影响了蒸馏过程、最终学生模型的特点,以及它在各种部署情境中的适用性。
任务通用蒸馏:构建一个通用型学生模型
任务通用蒸馏旨在创建一个较小的、通用型LLM,该模型保留了大型教师模型广泛的能力。目标不局限于任何单一的下游应用,而是寻求在大范围的输入分布上近似教师模型的行为,通常与教师模型预训练 (pre-training)所用的数据相似。
方法与目标
最常见的方法是训练学生模型以匹配教师模型在大型通用语料库上的输出分布(logits或概率)。通常使用Kullback-Leibler (KL) 散度来衡量学生模型和教师模型对每个输入token的概率分布之间的差异。
设T为教师模型,S为学生模型。设x是从通用数据分布Dgeneral中提取的输入序列。任务通用蒸馏损失LKD通常被公式化为最小化教师模型输出分布pT(y∣x,τ)与学生模型输出分布pS(y∣x,τ)之间的KL散度,并在数据集上取平均值:
LKD=Ex∼Dgeneral[DKL(pT(y∣x,τ)∣∣pS(y∣x,τ))]
此处,τ是用于软化概率分布的温度参数 (parameter),使得学生模型能够从教师模型赋予的相对概率中学习,即使是对于不正确的token。匹配层间的中间表示或注意力模式也可以加入,但主要目标仍是广泛的知识迁移。
优点
- 灵活性: 所得学生模型是通用型模型。它可以作为较小的核心模型,准备好针对各种下游任务进行微调 (fine-tuning),类似于原始教师模型的调整方式。
- 规模效率: 如果需要多个专用模型,蒸馏出一个通用学生模型再进行微调可能效率更高,而非对每个任务单独进行任务专用蒸馏,特别是如果微调学生模型比完整蒸馏成本低很多。
缺点
- 任务性能次优: 学生模型作为通用型模型,在特定任务上的表现可能不如直接为该任务蒸馏的模型。通用知识的迁移可能会稀释与特定应用相关的专业知识。
- 数据需求: 需要访问代表教师模型预训练分布的大型多样化数据集,处理起来计算成本高昂。
- 知识稀释: 在小得多的学生模型中捕捉大型教师模型的全部能力本身就存在困难。一些能力或不太常见的知识模式可能会丢失。
应用场景示例: 创建像Llama 3 70B这样的核心模型的较小版本,例如得到一个7B参数的模型,它仍然具备强大的通用语言理解和生成能力,适用于在不同应用范围中进行进一步微调。
任务专用蒸馏:打造一个专用型学生模型
任务专用蒸馏专注于迁移教师模型在特定下游任务上的专业知识,例如文本分类、问答或摘要。这里,教师模型通常首先在目标任务上进行微调 (fine-tuning),成为一个“专用教师模型”。目标是创建一个在这个任务上表现出色的学生模型,模仿微调后的教师模型在任务专用数据集上的行为。
方法与目标
蒸馏过程使用与目标任务关联的数据集(Dtask)。目标函数通常结合了标准的任务专用损失(例如分类的交叉熵损失,摘要的序列到序列损失)以及KD损失。KD部分鼓励学生模型专门在与任务相关的示例上匹配(经过微调的)教师模型的输出。
设Ltask为任务的标准监督损失(例如学生预测与真实标签ytrue之间的交叉熵)。组合损失Ltotal通常是一个加权和:
Ltotal=αLtask(pS(y∣x),ytrue)+(1−α)LKD(pT(y∣x,τ),pS(y∣x,τ))
期望是在任务专用数据集Dtask上计算的,pT现在表示微调后教师模型的输出分布。超参数 (parameter) (hyperparameter)α平衡了直接从真实标签学习与从教师模型的软化输出学习。
优点
- 高任务性能: 在特定目标任务上通常能实现接近大型微调教师模型的性能,甚至可能超出仅从任务数据从头训练的学生模型的表现。
- 专业化: 生成高度优化、紧凑的模型,非常适合部署在资源受限的环境中,只需要一种特定能力的场景。
- 定向知识迁移: 将蒸馏工作集中在与应用最相关的知识上,可能在该特定范围实现更有效的迁移。
缺点
- 通用性不足: 所得学生模型高度专业化,在其蒸馏范围之外的任务上表现可能不佳。
- 需要任务专用教师模型: 需要一个已经针对目标任务微调过的教师模型。如果需要多个专用学生模型,这就需要多个微调过的教师模型和单独的蒸馏运行。
- 设置复杂性: 管理针对各种任务的不同微调教师模型和蒸馏流程会增加操作开销。
应用场景示例: 构建一个高精度、低延迟的情感分析API。首先将一个大型模型微调用于情感分析,然后将其知识蒸馏 (knowledge distillation)到一个小得多的类BERT模型中,专门用于此任务,并优化以部署在边缘设备或无服务器功能上。
策略比较
任务通用蒸馏与任务专用蒸馏之间的选择在很大程度上取决于项目目标和限制条件。
任务通用蒸馏与任务专用蒸馏流程的比较,突出了教师模型、数据来源、目标和最终学生模型特点的差异。
决策因素:
- 最终目标: 目标是单一的高性能应用(倾向于任务专用)还是一个适用于多种潜在用途的通用小型模型(倾向于任务通用)?
- 性能要求: 在特定任务上实现接近教师模型性能的程度有多重要?任务专用通常能为目标应用带来更好的结果。
- 可用资源: 你是否能获取用于任务通用蒸馏的大型通用语料库?你是否有足够的计算资源来微调 (fine-tuning)教师模型并为每个所需的特定任务运行蒸馏?
- 生成能力: 蒸馏通用生成能力(如流畅性、连贯性、创造力)通常更难,并倾向于使用大型数据集的任务通用方法。蒸馏特定生成风格(例如,以特定格式进行摘要)可能通过任务专用方法有效处理。
- 与其他技术的结合: 任务专用蒸馏模型可以通过量化 (quantization)或剪枝进一步优化,这些技术专门针对其较窄的操作范围。
也可以采用混合策略。例如,可以先进行一次任务通用蒸馏,以创建一个通用的紧凑模型,然后对这个已经紧凑的学生模型进行任务专用微调,甚至进行第二轮任务专用蒸馏。
最终,理解在创建通用型学生模型与专用型学生模型之间的权衡,对于成功应用知识蒸馏 (knowledge distillation)进行LLM压缩并根据特定的性能和部署要求调整结果是重要的。