MiniLLM: On-Policy Distillation of Large Language Models

TL;DR

MiniLLM以反向KLD和在策略蒸馏训练小模型;LLaMA-7B在SelfInst达73.1 GPT-4分、23.2 Rouge-L。

cs.CL 🔴 高级 2023-06-14 28 次浏览
Yuxian Gu Li Dong Furu Wei Minlie Huang
知识蒸馏 大语言模型 反向KLD 在策略优化 指令跟随

核心发现

方法论

MiniLLM将传统前向KL[p||q]改为反向KL[q||p],使学生聚焦教师的主要模式而非覆盖低概率文本区域。作者用Policy Gradient推导在策略训练,并结合Single-Step Decomposition、Teacher-Mixed Sampling和Length Normalization,另加预训练语言建模损失LPT维持通用能力。

关键结果

  • 在GPT-2、OPT和LLaMA三族、120M至13B学生规模上,MiniLLM几乎全面超过SFT、word-level KD与SeqKD。OPT-1.3B在SelfInst获47.0 GPT-4分、14.8 Rouge-L,明显高于SeqKD的36.6与12.7。
  • LLaMA-7B学生在VicunaEval达到64.1 GPT-4分、20.7 Rouge-L,SelfInst为73.1与23.2;后者接近LLaMA-13B教师的75.5与23.4。GPT-2-760M在UnNI获37.7 Rouge-L,高于教师相关基线。
  • 暴露偏差分析显示,MiniLLM在150个以上生成词后ExAccErr趋于稳定;校准上,SST2的ECE/准确率为0.099/89.7,优于KD的0.191/84.7。去除长度归一化后Dolly验证Rouge-L由24.6降至14.7。

研究意义

论文回应了开放源代码LLM压缩中的关键问题:如何利用白盒教师分布,而不只是模仿采样文本。它说明生成任务不能简单套用分类模型的前向KL,因为容量有限的学生会为教师几乎不支持的区域分配概率。MiniLLM降低了推理成本,并提升长文本、分布外指令和校准表现,对边缘部署、企业私有化和低延迟服务具有直接价值。

技术贡献

核心贡献是把反向KLD正式用于LLM白盒蒸馏,并通过策略梯度将不可直接优化的序列目标转化为在策略学习。Eq.7同时利用词表级单步期望和序列级优势项;教师混合分布ep抑制奖励欺骗,近似重要性权重降低方差,长度归一化消除短输出偏置,LPT则保护基础语言能力。

新颖性

相较于SeqKD把教师生成的全部样本当作必须记忆的目标,MiniLLM允许学生在自身容量内选择教师偏好的高概率模式。新颖性不只在于采用reverse KLD,还在于针对文本生成提出完整的低方差、抗奖励欺骗在策略实现,并在120M—13B范围验证可扩展性。

局限性

  • 实验主要基于databricks-dolly-15K和指令跟随,教师也在同一任务数据上微调;结论对复杂推理、事实性、安全性和多语言任务的外推仍有限。
  • 反向KLD具有模式寻求倾向,可能忽略教师分布中的有价值多样性;训练需反复采样、计算教师概率和重要性权重,成本高于普通SFT或SeqKD。
  • 论文报告了整体指标,但对错误类型、事实幻觉和不同教师质量下的稳定性分析不足。

未来方向

后续可研究自适应KL方向、温度与教师混合系数,结合偏好反馈、安全约束和检索增强;还应在数学推理、多语言、代码、长上下文及黑盒教师场景验证,并设计更高效的缓存、蒸馏采样和多教师融合策略。

AI 总览摘要

大型语言模型能力强,却昂贵、缓慢且难以部署。知识蒸馏通常让小模型模仿教师输出,但传统前向KL会要求容量有限的学生覆盖教师的所有可能文本;在开放式生成中,这会把概率放到教师几乎不支持的“空白区域”,造成低质量样本。MiniLLM针对这一结构性问题,研究白盒LLM教师分布到小模型的压缩。

方法的关键转折是以反向KLD KL[q||p]替代KL[p||q]。学生从自身生成文本出发,接受教师概率反馈,因而更倾向于学习主要模式。作者用Policy Gradient实现on-policy distillation,并加入Single-Step Decomposition降低方差、Teacher-Mixed Sampling缓解重复文本导致的reward hacking、Length Normalization避免学生偏爱空短回答,同时保留LPT语言建模损失。

实验覆盖GPT-2、OPT、LLaMA,学生规模从120M到7B、教师最高13B,数据来自databricks-dolly-15K,评测包括DollyEval、SelfInst、VicunaEval、S-NI和UnNI。LLaMA-7B在SelfInst达到73.1 GPT-4分和23.2 Rouge-L,接近13B教师的75.5和23.4;OPT-1.3B在SelfInst达到47.0和14.8。MiniLLM还改善暴露偏差和校准,SST2 ECE为0.099。结果表明,生成式蒸馏的关键不是让学生记住更多文本,而是让它在自身能力范围内更可靠地选择文本。

深度分析

研究背景

KD源自Hinton等人的模型压缩,过去主要服务白盒分类模型;近年SeqKD和API蒸馏则用教师生成文本训练小型生成器。开放源LLM提供完整输出分布,使白盒蒸馏更有价值,但生成空间多模态、长尾且远比分类标签复杂,容量有限的学生难以覆盖教师分布。

核心问题

标准目标KL[p||q]惩罚学生未覆盖教师的每个模式,可能迫使其给教师低概率或零概率区域赋予过高概率。自由生成时,这些错误会累积为暴露偏差、低质量长文本和较差校准。因此需要既利用教师概率,又适配学生自身生成轨迹的目标。

核心创新

  • �� 反向KLD:聚焦教师主要模式,减少void-region概率。
  • �� On-policy Policy Gradient:从学生采样并用教师对数概率提供反馈。
  • �� Single-Step Decomposition:直接计算词表期望,降低长序列方差。
  • �� Teacher-Mixed Sampling:以α=0.2混合教师与学生分布,抑制奖励欺骗。
  • �� Length Normalization与LPT:分别消除短输出偏置、保持基础语言能力。

方法详解

  • �� 先在Dolly指令—回答数据上监督微调学生,选择验证损失最低的检查点。
  • �� 目标为L=KL[qθ||p]=−Eqθ log(p/q)。序列梯度由Policy Gradient给出,逐步奖励rt=log[p(yt|·)/qθ(yt|·)]。
  • �� 用ep=αp+(1−α)qθ采样,并用近似逐 token重要性权重控制方差。
  • �� Single部分对整个词表求Σy′qθ(y′)log(p(y′)/qθ(y′));Long部分使用归一化未来奖励RNorm。
  • �� 将两部分梯度与∇LPT相加更新参数,并采用clipping稳定训练。

实验设计

训练数据为databricks-dolly-15K,约12.5K训练、1K验证、0.5K测试;预训练语料为OpenWebText或RoBERTa corpus。模型包括GPT-2 120M/340M/760M、OPT 1.3B/2.7B/6.7B、LLaMA-7B,教师为GPT-2-1.5B、OPT-13B、LLaMA-13B。基线为SFT、word-level KD和SeqKD;指标为Rouge-L、GPT-4评分、人评、ECE和ExAccErr。

结果分析

MiniLLM在多数模型和数据集上领先。LLaMA-7B在VicunaEval获64.1 GPT-4、20.7 Rouge-L,在UnNI获40.2 Rouge-L;OPT-2.7B在SelfInst获52.7 GPT-4、17.2 Rouge-L。教师规模增大时学生表现持续提升。消融显示去除teacher mix-in或length normalization后Dolly验证Rouge-L分别降至20.4和14.7,证明工程组件重要。

应用场景

适合将开放源大模型压缩为本地助手、企业知识问答、边缘设备和低延迟API。前提是可访问教师输出分布、任务提示数据和足够采样算力。相较纯文本SeqKD,它尤其适合要求长回答、较好校准和跨数据集泛化的场景。

局限与展望

方法依赖白盒教师概率,不能直接用于只有API文本的黑盒教师;on-policy采样和教师前向计算也可能抵消部分部署收益。反向KL的模式寻求特性可能损失多样性,并不自动保证事实性或安全性。未来应扩展到推理、代码、多语言和更长上下文,研究自适应目标与高效采样。

通俗解读 非专业人士也能看懂

把大模型想成经验丰富的厨师,小模型是厨房里的学徒。传统做法要求学徒把师傅做过的每一道菜、包括偶尔失败或奇怪的菜,都原样记住。学徒能力有限,结果反而会把不该学的味道也混进菜单。MiniLLM换了规则:学徒先自己做菜,再让师傅评价哪些步骤和成品更像优秀作品。

学徒不会被迫复制师傅所有可能答案,而是集中练习最可靠、最常用的做法。为了避免学徒靠重复一句话骗取高评价,训练时让师傅偶尔参与抽样;为了防止学徒只做几口的小菜,又按菜品长度重新计算奖励。训练过程中还让学徒继续阅读普通文章,避免只会回答训练题。

结果显示,这种“边做边听建议”的方式比死记答案更好。它能产生更准确的长回答,遇到没见过的问题也更稳,并且更少在回答后半段失控。代价是训练时需要师傅反复打分,成本和实现复杂度都更高。

简单解释 像给14岁少年讲一样

想象你在玩游戏:超级高手是13B教师,小角色是120M或1.3B学生。普通训练像把高手所有走位录像都塞给小角色,甚至连高手偶尔乱按的动作也要照抄。小角色学不会这么多,于是比赛时会做出奇怪操作。

MiniLLM让小角色自己打,再让高手告诉它“这个决定靠谱”还是“别往这里走”。它重点学习高手最常用、最成功的路线,而不是地图上每个稀有角落。训练还安排高手偶尔带路,防止小角色发现某个重复动作能骗分;同时按比赛时长调整分数,避免它为了得分立刻结束。

研究者在Dolly、SelfInst、VicunaEval、S-NI和UnNI上测试,模型从GPT-2到OPT、LLaMA都有。LLaMA-7B在SelfInst得到73.1 GPT-4分,接近LLaMA-13B教师的75.5;OPT-1.3B也达到47.0。

所以重点不是“把大模型整本背下来”,而是“学会在关键时刻做对选择”。不过高手必须开放自己的概率信息,训练也更贵;而且更可靠不等于永远正确,事实错误和安全问题仍要另外检查。

术语表

Reverse KLD(反向KL散度)

衡量学生分布相对教师分布的差异,形式为KL[q||p]。它倾向保留教师的主要模式而忽略低概率区域。

MiniLLM的核心蒸馏目标。

On-Policy Distillation(在策略蒸馏)

学生使用自身或教师混合分布生成样本,再依据教师概率更新自己。它减少训练时教师强制与推理时自由生成的差异。

用于实现反向KLD。

Policy Gradient(策略梯度)

通过采样动作并按奖励调整其概率的优化方法。论文把每个token的教师相对概率视为奖励。

推导序列级反向KLD梯度。

Exposure Bias(暴露偏差)

训练时使用正确历史、推理时使用模型自身历史所造成的分布不一致。错误可能在长文本中逐步累积。

MiniLLM用on-policy采样缓解。

Teacher-Mixed Sampling(教师混合采样)

按ep=αp+(1−α)q混合教师和学生分布进行采样。它抑制重复文本等奖励欺骗。

实验中α固定为0.2。

ExAccErr

衡量训练—解码差异造成的额外累积错误。数值越低,表示暴露偏差越小。

用于长文本生成分析。

开放问题 这项研究留下的未解疑问

  • 1 反向KLD在数学推理、代码和多语言任务中是否仍优于前向KL,论文尚未回答;需要跨任务、跨语言和更强教师的系统对照。
  • 2 模式寻求可能牺牲创造性与事实覆盖率;如何自适应平衡准确性、多样性、安全性和校准,仍缺少理论与可靠指标。

应用场景

近期应用

企业本地助手压缩

企业可用开放源白盒教师、内部指令数据和MiniLLM训练较小模型,部署到私有服务器。预期降低显存和延迟,同时保留问答、摘要及流程响应能力。

边缘端指令模型

设备厂商可将OPT或LLaMA教师蒸馏到数亿参数学生,用于离线设备控制和个人助手。需准备任务提示、教师概率访问及额外采样训练资源。

远期愿景

可验证的模型压缩

未来可把MiniLLM与偏好反馈、检索和安全约束结合,形成面向事实性、合规性和低成本推理的压缩管线,服务大规模个性化AI。

原文摘要

Knowledge Distillation (KD) is a promising technique for reducing the high computational demand of large language models (LLMs). However, previous KD methods are primarily applied to white-box classification models or training small models to imitate black-box model APIs like ChatGPT. How to effectively distill the knowledge of white-box LLMs into small models is still under-explored, which becomes more important with the prosperity of open-source LLMs. In this work, we propose a KD approach that distills LLMs into smaller language models. We first replace the forward Kullback-Leibler divergence (KLD) objective in the standard KD approaches with reverse KLD, which is more suitable for KD on generative language models, to prevent the student model from overestimating the low-probability regions of the teacher distribution. Then, we derive an effective on-policy optimization approach to learn this objective. The student models are named MiniLLM. Extensive experiments in the instruction-following setting show that MiniLLM generates more precise responses with higher overall quality, lower exposure bias, better calibration, and higher long-text generation performance than the baselines. Our method is scalable for different model families with 120M to 13B parameters. Our code, data, and model checkpoints can be found in https://github.com/microsoft/LMOps/tree/main/minillm.

cs.CL cs.AI