On-Policy Distillation：大模型后训练中的策略蒸馏原理与实践

1. 项目概述On-Policy DistillationOPD是什么最近在折腾大模型后训练特别是对齐和微调阶段发现一个挺有意思的现象大家都在用蒸馏Distillation来压缩模型或者迁移知识但很多方法效果不稳定尤其是在从更强的教师模型Teacher Model向学生模型Student Model传递能力时经常出现“学歪了”或者“消化不良”的情况。这让我开始关注一种更贴近模型自身行为逻辑的蒸馏方式——On-Policy DistillationOPD直译过来就是“在策略上的蒸馏”。这个“On-Policy”是关键它指的是蒸馏过程发生在学生模型自己生成的“轨迹”或“行为数据”上而不是在教师模型的输出或者一个静态的、外部的数据集上。简单来说OPD的核心思想是让学生模型在自己“走”出来的路上向教师模型学习如何“走”得更好。这听起来有点抽象我打个比方。传统的离线蒸馏Off-Policy Distillation好比是让一个学武术的学生反复观看宗师教师模型打一套固定拳法的录像带然后模仿。学生可能学得很像但一旦上了擂台面对瞬息万变的真实对手动态的、开放域的输入可能就手足无措了。而OPD则是让学生自己先上擂台比划生成自己的回答或行为序列然后宗师在旁边实时观看并针对学生此时此刻的招式即学生模型当前的输出分布给出更精妙的修正建议教师模型的指导。学生是在自己的实战情境中学习针对性极强。为什么大模型的后训练阶段特别需要OPD呢因为大模型尤其是经过指令微调或RLHF对齐后的模型面临的任务不再是简单的分类或填空而是开放域的对话、推理、创作。它的“策略”——即根据输入生成下一个词的概率分布——是高度复杂和上下文依赖的。用一个固定的、由教师模型生成的“标准答案”数据集去教学生很容易导致分布偏移Distribution Shift。学生学到的可能只是如何在教师模型偏好的数据分布上表现好一旦遇到自己策略分布下的输入这恰恰是它部署后要面对的真实情况性能就会下降。OPD通过强制蒸馏信号与学生模型当前的策略分布对齐确保了学习的稳定性和有效性让学生模型在其“主场”上变得更强。2. OPD的核心原理与必要性拆解要理解OPD为什么有效我们需要深入两个层面一是传统蒸馏在大模型后训练中的局限二是OPD如何从原理上解决这些问题。2.1 传统离线蒸馏的“阿喀琉斯之踵”在大模型场景下传统的知识蒸馏通常是这样操作的用一个更大、更强的教师模型比如GPT-4在一个庞大的指令数据集上生成回答得到一系列输入 教师输出配对。然后我们用这些配对数据来训练一个较小的学生模型比如一个7B模型目标是最小化学生模型输出与教师模型输出之间的差异通常用KL散度等损失函数。这个方法看似直接但存在几个根本性问题暴露偏差Exposure Bias学生模型在训练时看到的永远是教师模型生成的“完美”序列。但在它自己推理时是从头开始自回归生成的。一旦它在某个时间步产生了一个与教师模型稍有不同的词后续的生成就会进入一个它从未在训练中见过的状态序列导致错误累积生成质量下降。这就像学开车时只看教练完美驾驶的录像但自己一握方向盘遇到第一个小失误就不知道如何纠正了。分布不匹配Distribution Mismatch教师模型生成答案时是基于其自身的策略分布。这个分布与学生模型初始的、或者训练中的策略分布是不同的。强迫学生去拟合一个不属于自己策略空间的目标就像让一个习惯短跑的运动员去模仿马拉松运动员的跑姿不仅低效还可能受伤模型崩溃或性能下降。在对话中表现为学生模型生成一些生硬的、不符合其自身语言风格的“教师腔”回答。探索受限静态的教师-输出数据集固定了学习的边界。学生模型没有机会在那些教师模型可能未充分探索的、但对学生模型自身进化有益的方向上进行学习和改进。2.2 OPD的机制在自我轨迹中对齐与优化OPD的流程巧妙地规避了上述问题。其核心步骤可以概括为一个循环学生模型生成轨迹给定一个提示Prompt让学生模型以其当前的参数运行生成一个完整的回答序列或称为“轨迹”。这个轨迹完全反映了学生模型当前的策略和能力。教师模型进行评估与“润色”将同一个提示以及学生模型生成的轨迹或其中间状态输入给教师模型。教师模型的任务不是从头生成一个新答案而是对学生生成的轨迹进行评估、打分或提供改进版本。具体形式可以是生成改进序列教师模型基于学生的输出生成一个质量更高、更符合要求的版本。提供逐词或逐片段反馈比如通过计算每个生成位置的学生输出概率与教师输出概率的差异作为细粒度的学习信号。给出序列级奖励教师模型对整个学生生成的序列进行评分例如通过一个奖励模型或自身判断。基于对齐信号的策略更新利用步骤2中得到的信号改进的序列、概率差异或奖励构建损失函数来更新学生模型的参数。关键点在于这个损失函数是基于学生模型自身采样到的状态和动作生成的词来计算的。这个过程与强化学习中的“On-Policy”算法如PPO思想同源。在PPO中智能体根据当前策略与环境交互收集数据然后用这些数据来更新同一策略。OPD同理学生模型用自己的策略生成模型与环境提示交互产生数据再用教师模型对这些数据的评价来优化自己。这保证了训练目标教师提供的信号与数据分布学生策略始终保持一致极大缓解了分布偏移问题。为什么这对大模型后训练至关重要在大模型对齐的后期我们往往不是在教模型“事实知识”这已在预训练中学到而是在塑造其“行为风格”、“价值观”和“推理方式”。这些是高度策略性的。OPD允许模型在展现自身行为倾向的语境下获得指导使得对齐和能力的提升更加精准和稳定。例如在训练一个模型使其回答更加安全时OPD能直接针对模型自己可能产生的不安全回答倾向进行修正而不是泛泛地学习一些通用的安全模板。3. OPD的典型实现方案与实操要点理论讲完了我们来点实际的。如何实现一个OPD训练循环这里我结合常见的实践拆解一个基于序列级改进的OPD方案这也是目前很多研究采用的主流思路。3.1 整体训练架构设计一个典型的OPD训练流程包含以下几个核心模块学生模型待训练和提升的模型参数为 θ。教师模型提供指导信号的、能力更强的模型参数固定通常是冻结的。数据流一个提示池包含需要学习的各类任务或对话提示。训练循环从提示池采样一个批次Batch的提示x。将x输入学生模型使用其当前的生成策略例如采样或贪婪解码生成完整的输出序列y_s ~ π_θ(y|x)。这就是学生的“轨迹”。将相同的提示x和学生生成的y_s有时可能连同x一起提交给教师模型。教师模型的任务是生成一个针对x的、质量更高的参考序列y_t。这里y_t可以是对y_s的改写、润色也可以是教师从头生成的一个更优答案。计算损失。一个常用的损失函数是序列级蒸馏损失它鼓励学生模型输出与教师模型改进后的输出分布对齐L_OPD D_KL( π_θ(y | x) || π_teacher(y | x, y_s) )其中π_teacher(y | x, y_s)可以近似为以y_t为目标的分布例如使用教师模型对y_t中每个token的logits或直接使用y_t作为硬标签。更简单的实现是使用交叉熵损失让学生模型预测教师模型生成的改进序列y_tL_CE -Σ log π_θ(y_t_i | x, y_t_i)反向传播更新学生模型参数 θ。可选将新生成的x,y_s或x,y_t数据加入提示池实现某种形式的数据扩充或课程学习。3.2 关键超参数与配置选择在实操中以下几个点的选择直接影响OPD的效果教师模型的调用方式完全生成教师模型无视学生的y_s独立为提示x生成答案y_t。这种方式简单但可能与学生轨迹关联性较弱。条件生成教师模型以(x, y_s)为条件生成一个改进版y_t。例如提示可以是“以下是一个AI助手对用户问题的回答。请评估这个回答并提供一个更准确、更安全的版本。用户问题[x]。原回答[y_s]。改进后的回答”。这种方式更能体现“针对性指导”。偏好打分教师模型不生成完整答案而是对学生生成的多个候选答案进行排序或打分然后使用偏好优化如DPO损失来更新学生模型。这更节省教师模型的推理成本。学生轨迹的生成策略采样Sampling使用温度采样Temperature Sampling或核采样Top-p Sampling来生成y_s。这能带来更多的探索产生多样化的轨迹让教师模型看到学生更全面的行为但可能包含较多低质量输出。贪婪解码Greedy Decoding生成当前概率最高的词。这产生的轨迹更稳定、质量可能更高但探索性不足。实操建议在训练初期可以采用较高的温度如0.8-1.0进行采样鼓励探索在训练中后期逐渐降低温度或转向贪婪解码专注于利用已学到的良好行为。损失函数的设计序列级交叉熵最直接将教师输出y_t视为硬标签。优点是稳定易于实现。KL散度损失更“软”匹配整个输出分布。需要获取教师模型在y_t每个位置上的logits输出层未归一化的分数计算成本稍高但理论上更合理。结合原始任务损失如果学生模型本身还在进行指令微调可以将OPD损失与标准的指令跟随损失如根据(x, y_t)计算的交叉熵加权结合防止模型遗忘基础能力。教师模型的选择同构更强模型最理想的情况是教师模型是学生模型的同架构、更大参数量的版本例如学生是7B教师是70B。知识迁移最顺畅。异构专家模型教师模型可以是专门针对某个领域如代码、数学精调的模型即使总参数量不大但在特定技能上超越学生。这时OPD用于技能蒸馏。提示工程教师模型不一定是微调过的。一个足够强大的基础模型如GPT-4通过精心设计的提示词也能扮演优秀的“指导者”角色。这在基于API的蒸馏中很常见。注意使用教师模型尤其是通过API调用的大型商业模型时需密切关注数据隐私和成本。生成大量轨迹并进行评估会产生可观的推理费用。对于敏感数据务必使用本地部署的教师模型。4. 实操流程构建一个简易的OPD训练循环下面我将用一个简化的代码框架展示如何为一个开源大模型例如使用Hugging Face Transformers库的模型实现一个基础的OPD训练循环。我们假设使用序列级交叉熵损失教师模型通过条件生成提供改进答案。import torch from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments from datasets import Dataset import random # 1. 初始化模型和分词器 student_model_name your-student-model # 例如 “meta-llama/Llama-2-7b-chat-hf” teacher_model_name your-teacher-model # 例如一个更大的模型或API的本地代理 tokenizer AutoTokenizer.from_pretrained(student_model_name) tokenizer.pad_token tokenizer.eos_token # 设置填充token student_model AutoModelForCausalLM.from_pretrained(student_model_name) teacher_model AutoModelForCausalLM.from_pretrained(teacher_model_name, torch_dtypetorch.float16, device_mapauto) teacher_model.eval() # 冻结教师模型 # 2. 准备提示池示例 prompts [ 解释一下量子计算的基本原理。, 写一首关于春天的五言绝句。, 用Python实现一个快速排序算法。, # ... 更多提示 ] # 3. 定义数据生成函数核心OPD步骤 def generate_opd_dataset(prompts, student_model, teacher_model, tokenizer, num_samples1000): data {prompt: [], teacher_response: []} for prompt in prompts[:num_samples]: # 控制生成数量 # 学生生成轨迹 (On-Policy) student_inputs tokenizer(prompt, return_tensorspt).to(student_model.device) with torch.no_grad(): student_outputs student_model.generate( **student_inputs, max_new_tokens256, do_sampleTrue, # 使用采样进行探索 temperature0.7, top_p0.9, ) student_response tokenizer.decode(student_outputs[0][len(student_inputs[input_ids][0]):], skip_special_tokensTrue) # 教师提供改进答案 (条件生成) # 构建给教师的提示让其基于学生回答进行改进 teacher_prompt f用户问题{prompt} AI助手的初始回答{student_response} 请评估上述回答并提供一个更准确、完整、有帮助的改进版本。改进后的回答 teacher_inputs tokenizer(teacher_prompt, return_tensorspt).to(teacher_model.device) with torch.no_grad(): teacher_outputs teacher_model.generate( **teacher_inputs, max_new_tokens300, do_sampleFalse, # 教师通常使用贪婪解码以获得稳定输出 temperature0, ) # 提取教师生成中“改进后的回答”之后的部分 teacher_full_text tokenizer.decode(teacher_outputs[0], skip_special_tokensTrue) # 这里需要根据你的提示格式进行解析假设我们能提取出最终答案 teacher_response teacher_full_text.split(改进后的回答)[-1].strip() data[prompt].append(prompt) data[teacher_response].append(teacher_response) return Dataset.from_dict(data) # 4. 生成训练数据集 train_dataset generate_opd_dataset(prompts, student_model, teacher_model, tokenizer, num_samples500) # 5. 数据预处理将提示和教师回答拼接并tokenize def preprocess_function(examples): # 将提示和教师回答拼接成标准指令格式 texts [f### 指令{p}\n### 回答{r} for p, r in zip(examples[prompt], examples[teacher_response])] model_inputs tokenizer(texts, max_length512, truncationTrue, paddingmax_length) # 设置标签通常对于因果语言模型标签就是输入ids的移位 model_inputs[labels] model_inputs[input_ids].copy() return model_inputs tokenized_dataset train_dataset.map(preprocess_function, batchedTrue) # 6. 定义训练参数并开始训练 training_args TrainingArguments( output_dir./opd_finetuned_model, num_train_epochs3, per_device_train_batch_size4, gradient_accumulation_steps4, warmup_steps100, logging_steps50, save_steps500, evaluation_strategyno, save_total_limit2, fp16True, # 如果硬件支持 ) trainer Trainer( modelstudent_model, argstraining_args, train_datasettokenized_dataset, data_collatorlambda data: {input_ids: torch.stack([d[input_ids] for d in data]), attention_mask: torch.stack([d[attention_mask] for d in data]), labels: torch.stack([d[labels] for d in data])}, ) trainer.train()关键操作解析generate_opd_dataset函数这是OPD循环的核心。它模拟了学生生成、教师评估改进的过程。注意这里教师提示词的设计 (teacher_prompt) 至关重要它决定了教师指导的“风格”和“侧重点”。学生采样策略代码中使用了do_sampleTrue, temperature0.7, top_p0.9这是在训练初期鼓励探索的典型设置。教师生成策略教师使用do_sampleFalse贪婪解码旨在提供稳定、高质量的参考答案。数据格式预处理中将提示和教师回答拼接成固定的指令格式便于模型学习指令跟随。损失函数是标准的语言模型交叉熵损失目标是让学生的输出分布逼近“教师改进后的回答”。这个流程是一个高度简化的示例。在实际生产中你需要考虑更复杂的提示工程、处理多轮对话、集成奖励模型进行偏好优化、以及实现一个持续更新的动态提示池。5. 常见问题、挑战与调优技巧实录在实际部署和调优OPD时我踩过不少坑也总结出一些让训练更稳定、效果更好的技巧。5.1 典型问题与排查思路学生模型“摆烂”或退化训练一段时间后学生模型的生成变得非常简短、通用如“我不知道”、“这是一个好问题”或者开始重复教师的某些固定句式。原因这通常是模式坍塌的迹象。教师模型提供的改进答案多样性不足或者损失函数权重过大导致学生过度拟合少数几种“安全”模式丧失了生成多样性。排查与解决检查教师输出随机抽样查看教师模型生成的改进答案y_t是否过于单一或模板化尝试优化给教师的提示词增加“请提供多样化的、富有创造性的回答”等指令。引入多样性正则化在损失函数中加入鼓励输出多样性的项例如最大化生成序列的信息熵或者加入对抗性损失。调整损失权重降低OPD损失在总损失中的权重同时保留一部分原始的下一个词预测损失预训练损失以维持模型的基础语言能力。混合数据源不要只使用OPD生成的数据。将OPD数据与高质量的指令微调数据集如Alpaca、ShareGPT混合训练防止分布过度偏移。训练不稳定损失震荡剧烈。原因学生模型自身生成的轨迹y_s质量波动大导致教师模型提供的指导信号y_t也随之剧烈波动。这相当于学习目标在不断跳跃。排查与解决平滑学生轨迹在训练初期可以降低学生生成时的采样温度甚至使用贪婪解码先获得一批相对稳定的轨迹进行学习。随着训练进行再逐步提高温度引入探索。使用轨迹池不直接用当前批次学生生成的轨迹而是维护一个经验回放池Replay Buffer从中采样过去生成的轨迹进行学习。这能稳定数据分布。教师信号平滑不使用原始的教师输出y_t作为目标而是使用教师模型输出的logits与学生当前logits的加权平均即软目标或者对多个教师改进版本取平均。计算成本过高。原因每一轮训练都需要前向传播学生模型生成和教师模型评估尤其是教师模型通常更大推理成本是主要瓶颈。排查与解决蒸馏教师logits如果教师模型是本地的可以在生成改进序列y_t时一并保存其每个位置的输出logits。在后续训练中直接使用这些缓存的logits计算KL散度损失避免反复调用大教师模型进行前向传播。小教师或委员会使用一个参数较小但性能尚可的模型作为“常驻教师”或者使用多个较小模型组成的“委员会”来提供综合指导信号仅在关键步骤调用大教师。非自回归生成对于教师评估可以尝试非自回归或更高效的生成方式减少解码步数。5.2 高级调优技巧与心得课程学习Curriculum Learning不要一开始就用很难的提示。可以从简单的、封闭式的问答开始让学生模型先学会在简单任务上模仿教师。然后逐步增加提示的复杂性、开放性和长度。这能显著提升训练的稳定性和最终效果。迭代式OPD将OPD过程迭代进行。用第一轮OPD训练得到的学生模型作为新的“学生”再次进行OPD。或者用上一轮训练中生成的优质(x, y_t)对作为下一轮训练提示池的一部分。这能实现数据的自我进化和模型能力的持续提升。结合强化学习OPD天然可以与RLHF基于人类反馈的强化学习结合。教师模型提供的改进序列或偏好评分可以作为强化学习中的奖励信号。具体来说可以用教师模型对(x, y_s)的打分作为奖励用PPO等算法更新学生模型。这形成了“On-Policy”的强化学习比使用静态奖励模型更适应学生策略的变化。关注“负样本”除了让学生学习教师的好答案也可以有意地让学生生成一些错误答案例如通过引导或对抗性提示然后让教师模型指出错误并提供纠正。这种“从错误中学习”的方式能帮助模型更清晰地理解边界和禁忌。我个人最深刻的体会是OPD不是一个“一劳永逸”的魔法。它的效果极度依赖于教师模型的质量、提示词的设计以及训练流程的稳定性。它更像是一个精密的反馈调节系统。你需要像教练一样不断观察学生的表现模型评估调整训练难度课程设计和指导方式提示工程才能引导模型稳健地朝着期望的方向进化。最开始我直接套用简单实现效果时好时坏。后来花了大量时间在分析教师-学生交互数据上调整生成策略和损失混合比例才让训练曲线平滑下来最终得到的模型在保持自身风格的同时安全性和有用性都有了可感知的提升。这个过程虽然繁琐但当你看到模型在自己“选择”的道路上因为获得了精准的指导而越走越好时那种成就感是单纯跑通一个离线蒸馏脚本无法比拟的。