模型部署与推理优化

知识蒸馏:从 logits 到 on-policy

kp-011核心25 分钟03-剪枝与蒸馏

一句话定义

知识蒸馏(Knowledge Distillation)让小模型(学生)不只学数据的硬标签,还模仿大模型(教师)输出的软分布,从而用更少的参数继承教师的能力,是"造小模型"的核心技术。

为什么重要

端侧与低延迟场景需要的往往不是 70B 而是 1–3B 模型——而同级模型中"蒸馏过的"普遍强于"从零训的",因为软标签携带了类间相似度等暗知识,数据效率更高。对推理工程师,蒸馏还是剪枝后的恢复手段(kp-010)与投机解码 draft 模型的来源(kp-018)。

前置知识

  • softmax 与交叉熵;kp-005(logits/概率的语言)。

核心概念

  • 软标签与温度 T:教师 logits 除以 T 再 softmax,T>1 把分布"软化",暴露次优选项的相对概率——"3 和 7 长得像"这类信息是硬标签没有的。
  • 经典 KD 损失:学生损失 = 任务交叉熵 + λ × T² × KL(教师软分布 ‖ 学生分布);T² 因子抵消梯度随 T 缩放。
  • 特征蒸馏(FitNets 一系):对齐中间层表示,适合结构与教师差异大的学生。
  • on-policy 蒸馏(GKD 一系):学生自己生成序列、教师给逐 token 分布当老师,避免"教师分布与学生实际会走到状态不匹配"(exposure bias),是 LLM 时代比逐词模仿更有效的路线。
  • 数据蒸馏:用教师生成/筛选训练数据再训学生,工程上最易落地的变体。

原理与机制

蒸馏有效的机制是"暗知识"传递:教师的软分布携带类间相似度与决策边界形状,监督信号密度远高于 one-hot 硬标签,学生因此能用更少数据学到更平滑的决策面。LLM 时代进一步发现:学生在"自己实际会生成的状态"上向教师分布对齐(on-policy),比逐词模仿教师生成的文本更能消除 train/inference 的分布错位。

公式与模型

\mathcal{L} = (1-λ) CE(y, ps) + λ T2 KL(pT(·|x;T) ‖ ps(·|x;T)), p = softmax(z/T)

KL 项逐 token 求平均。直觉:CE 教"正确答案",KL 教"正确答案与错误答案之间的距离感"。

直观类比

跟着老师做题不只抄标准答案(硬标签),还要听老师讲"为什么 B 选项最像对的"(软分布)——考场上遇到新题时,后者才是迁移能力。

实例或案例

为投机解码准备 draft 模型(kp-018):用 70B 教师蒸馏一个 1B 学生,要求学生与教师分布相关性高(接受率 α 高)比学生绝对能力更重要,因此蒸馏目标里常加大 KL 项权重并使用同源数据。另一例:端侧 1.5B 助手 = 7B 教师 + on-policy 蒸馏 + 量化(kp-012 流水线)。

常见误区

  • "蒸馏学生必弱于同尺寸从零训练":恰好相反,软标签提升数据效率,同级学生通常更强。
  • "蒸馏=让学生输出和教师逐字一样":目标是分布匹配不是字符串复制;逐字对齐反而限制泛化。
  • "logits 蒸馏必须有教师概率接口":闭源 API 无 logits 时,可用数据蒸馏/生成式 KD 变体。
  • 温度 T=1 一定最好:任务相关,硬任务常用 T≈2 暴露更多暗知识。

自测题

  1. 写出经典 KD 损失并解释 T² 因子。

要点:CE + λT²·KL(teacher‖student);补偿温度放大后 KL 梯度的量纲,使两项可加权平衡。

  1. 温度 T 在蒸馏中的作用?

要点:软化教师分布,暴露类间相似度等暗知识。

  1. on-policy 蒸馏解决什么问题?

要点:学生在自己会到达的状态上向教师学习,缓解 train/inference 分布不匹配。

与其他知识点的关系

kp-010 的剪后恢复训练用本页技术;kp-012 编排三者;kp-018/kp-019 的 draft 模型常由蒸馏产出。

延伸阅读

《Distilling the Knowledge in a Neural Network》(Hinton 等,2015)——领域奠基作,正文仅数页;《On-Policy Distillation of Language Models》(Agarwal 等,2024)——LLM 时代的修正。

#蒸馏#KD#温度#on-policy#教师模型