LLM 知识蒸馏生产实战：用 Teacher Logprob 契约、On-policy 回放与 Top-K Tail Bucket 控制学生模型退化

执行摘要

大模型的知识蒸馏（Knowledge Distillation, KD）经常被简化成”拿大模型生成答案，再对小模型做一次 SFT”。然而，真正将蒸馏推上生产环境时，真正棘手的问题往往不在训练脚本本身，而在于：教师概率分布如何版本化、温度和散度如何固定、Top-K 截断后如何保留尾部概率、学生自生成轨迹如何回放，以及学生模型何时可以替代教师模型。

本文从工程落地视角出发，系统梳理知识蒸馏的五个核心环节——Teacher Logprob 契约、Top-K Tail Bucket、On-policy 回放、Tokenizer 对齐和上线门禁，帮助团队避免蒸馏后静默退化。

---

背景问题：蒸馏不是一次普通 SFT

知识蒸馏的目标是让较小的 Student（学生模型）学习较大 Teacher（教师模型）的行为分布，从而降低推理显存、延迟和单位请求成本。经典方法把硬标签损失与教师软分布之间的散度组合起来：

L = (1 - α) · L_CE(student, label) + α · T² · KL(p_teacher^T || p_student^T)

其中 T 是温度，α 控制硬标签与蒸馏信号的权重。

公式本身并不复杂，生产问题却集中在公式之外：

1. 教师模型、Tokenizer、Chat Template 或生成参数变化后，旧 Logprob 是否仍可复用；
2. Top-K Logprob 经网络或对象存储传输后，是否仍代表原始分布；
3. 学生只模仿教师生成文本，是否会在自行生成时暴露训练未见的错误轨迹；
4. 学生模型总体分数接近教师，但是否在长文本、低频领域、格式遵循和概率校准上发生退化；
5. 教师调用、缓存构建和重复训练的成本是否可以追踪与复现。

因此，蒸馏系统首先要管理的是概率契约与数据血缘，其次才是训练脚本。

---

核心原理

一、Teacher Logprob 必须成为版本化制品

对每条训练记录，不应只保存 Teacher 输出文本。建议至少记录以下结构化信息：

{
  "record_id": "sha256(prompt_canonical_bytes)",
  "teacher_model_revision": "model-commit-or-checkpoint-hash",
  "tokenizer_revision": "tokenizer-artifact-hash",
  "chat_template_hash": "sha256(template)",
  "generation_config_hash": "sha256(config)",
  "temperature": 1.0,
  "loss_support": "teacher_top_k_plus_tail",
  "top_k": 32,
  "token_ids": [123, 456],
  "teacher_topk_token_ids": [[123, 789], [456, 321]],
  "teacher_topk_logprobs": [[-0.1, -2.3], [-0.2, -1.7]],
  "tail_logprob_mass": [-0.08, -0.12]
}

关键原则：只要模型修订、Tokenizer、模板、温度或 Top-K 规则发生变化，这批 Logprob 就应视为新制品，不能静默覆盖。这是整个蒸馏系统的数据血缘基础。

二、Top-K 截断必须保留 Tail Bucket

全词表 Logit 传输和存储代价高，工程上常只保存 Top-K。但若简单丢弃其余词并重新归一化，教师分布会被人为变尖。

更稳妥的方式是把未进入 Top-K 的概率质量合并为 Tail Bucket：

p_tail = 1 - Σ p_top_k

计算散度时，让学生 Top-K 之外的概率总和与 p_tail 对齐。这样既减少传输量，又不会假装尾部概率不存在。

散度类型	特点	对尾部概率的敏感度
Forward KL	覆盖教师全部支持集（mean-seeking）	高——缺失尾部会导致偏差
Reverse KL	聚焦学生高概率区域（mode-seeking）	中——尾部影响相对小
Generalized JSD	Forward 和 Reverse 的折中	中——通过 β 参数调节

需要注意，Forward KL、Reverse KL 和 JSD 对支持集合的要求不同。使用外部 Teacher Server 时，应把 loss_top_k、散度方向、Tail Bucket 规则写进训练配置指纹，而不是依赖框架默认值。

三、On-policy 回放解决训练—推理分布偏移

纯离线蒸馏通常让学生学习教师产生的序列。上线后，学生根据自己的历史 Token 继续生成，一旦早期发生偏差，后续状态就可能落入训练集未覆盖区域。这就是经典的 分布偏移（Distribution Shift） 问题。

On-policy Distillation 的做法是：

1. 让学生先生成一部分回答；
2. 把学生实际访问到的状态发送给教师；
3. 在这些状态上获取教师概率反馈；
4. 用教师分布纠正学生自己的错误轨迹。

这与只对教师答案做 SFT 的区别很大：前者让学生学习”自己犯错之后该如何恢复”，后者主要学习”教师在理想轨迹上会怎么写”。

四、Tokenizer 不一致不能直接做逐位置 Logit 对齐

Teacher 和 Student 若使用不同 Tokenizer，同一字符串可能产生不同长度和不同边界的 Token 序列。此时 teacher_logits[t] 与 student_logits[t] 没有天然的一一对应关系。

生产方案按复杂度从低到高通常是：

方案	复杂度	适用场景
共享 Tokenizer 与词表	低	最稳定，适合直接 Logit KD
序列级蒸馏	中	把教师文本当作训练目标，不做逐 Token 分布对齐
公共词表或 Token 映射	高	只在可靠对齐区域计算损失
跨 Tokenizer 表征对齐	很高	需要额外模型或对齐算法，必须单独验证

⚠️ 不要通过裁剪词表尺寸或按位置硬配对来”快速兼容”，这种错误往往不会报错，却会让损失失去意义。

---

工程落地

一、建立三层蒸馏流水线

建议拆成三个可独立重跑的阶段：

1. Trajectory Builder：生成教师轨迹、学生轨迹或混合轨迹；
2. Teacher Scorer：输出版本化 Logprob 制品；
3. Student Trainer：读取固定制品训练，不直接依赖在线教师。

这样做的好处是教师服务故障不会直接中断学生训练，Logprob 可以审计和复用，也能单独比较不同散度或温度配置。

若采用完全 On-policy 训练，可以保留外部 Teacher Server，但仍应把请求和响应写入可重放日志，并限制教师模型修订在一次实验内保持不变。

二、定义不可变实验指纹

每次实验至少固定以下配置：

teacher:
  model_revision: teacher-v3.2
  tokenizer_hash: sha256:...
  chat_template_hash: sha256:...
  server_image_digest: sha256:...

student:
  base_revision: student-v1.8
  tokenizer_hash: sha256:...

distillation:
  temperature: 1.0
  kd_ratio: 0.7
  divergence: generalized_jsd
  beta: 0.5
  on_policy_ratio: 0.6
  top_k: 32
  tail_bucket: true

训练结果必须能反查到教师制品、数据集快照、损失配置和代码版本。

三、对 Teacher Server 做契约测试

Teacher Server 不只需要可用性检查，还应验证以下契约：

• 返回的 Token ID 是否对应声明的 Tokenizer；
• Logprob 是否为自然对数；
• Top-K 是否按概率降序；
• Top-K 概率和 Tail Mass 是否接近 1；
• Temperature 是否在服务端还是训练端生效；
• EOS、特殊 Token 与 Padding 是否参与损失；
• 同一请求在固定配置下是否满足允许的数值误差。

四、缓存策略按”概率制品”而不是”文本结果”设计

Logprob Cache Key 应包含以下维度：

prompt_hash + teacher_model_revision + tokenizer_revision
+ chat_template_hash + generation_config_hash + scoring_position_hash

缓存命中后还要校验 Schema Version。禁止仅用 Prompt 文本作为 Key，否则教师升级后会继续读取旧分布，导致训练结果无法复现。

五、能力门禁要比较学生、教师和原始基线

至少保留三组对照：

• 原始 Student Base
• Distilled Student
• Teacher

发布门禁建议覆盖以下维度：

维度	检查项
领域能力	领域任务准确率或业务成功率
行为规范	指令遵循、格式约束和拒答行为
长度鲁棒性	长度分桶表现
长尾覆盖	低频领域与难例
概率校准	概率校准与置信度
性能指标	P50/P95 延迟、显存与吞吐
成本	单请求和单位有效输出成本

蒸馏目标不是让学生在所有指标上等于教师，而是在明确成本约束下达到可接受的能力边界。

---

适用场景

适合知识蒸馏的情况：

• 大模型效果可接受，但线上延迟或 GPU 成本过高；
• 业务集中在有限领域，学生不必保留教师全部通用能力；
• 有稳定教师模型和足够规模的高质量 Prompt 分布；
• 需要把模型部署到较小 GPU、边缘设备或高并发服务；
• 可持续收集学生真实失败轨迹并进行迭代蒸馏。

不适合的情况：

• 教师本身不稳定，数据分布频繁剧烈变化；
• Teacher 与 Student Tokenizer 无法可靠对齐；
• 业务要求学生完整继承教师的长尾通用能力。

---

常见误区

误区一：教师生成文本等于知识蒸馏

只训练教师生成文本更接近 Sequence-level KD 或 SFT。它丢失了教师对次优 Token 的概率结构，也无法直接表达教师的不确定性。

误区二：温度越高，教师信息越丰富

温度提高会拉平分布，但过高会放大尾部噪声。温度必须通过校准集选择，并与散度方向、Top-K 大小共同评估。

误区三：Top-K 越小，训练越便宜且质量不受影响

Top-K 太小会丢失教师在相近候选之间的排序信息。即便只保留少量 Token，也应保存 Tail Bucket，并对 K 值做消融实验。

误区四：总体平均分达标即可上线

蒸馏退化常集中在长文本、罕见格式、少数领域和置信度校准。平均分可能掩盖关键业务子集的明显回退。

误区五：教师升级后缓存仍可继续使用

教师修订、Tokenizer、模板或推理参数任一变化，都可能改变概率分布。旧缓存只能作为旧实验制品保留，不能默认继续混用。

---

上线检查清单

制品与契约

• Teacher、Student、Tokenizer、Template 和代码版本均有不可变指纹
• Logprob Schema、对数底、温度位置和 Top-K 规则已固定
• Top-K 之外的概率质量已通过 Tail Bucket 保留
• Teacher Server 契约测试已进入 CI

训练与数据

• On-policy 与 Off-policy 比例有明确配置
• 学生轨迹可重放，失败样本有稳定 Record ID
• Padding、Prompt、Assistant 与特殊 Token 的 Loss Mask 已验证
• 训练断点能恢复轨迹游标和 Logprob 制品版本

质量与发布

• 分领域、分长度、分难度回归已通过
• 学生与教师的概率校准差异已评估
• Latency、Throughput、显存和成本达到压缩目标
• 灰度期间保留教师回退路径
• 线上失败轨迹能够回流下一轮 On-policy 蒸馏

---

FAQ

离线 Logit Cache 和在线 Teacher Server 应该选哪种？

教师稳定、数据集固定且希望提高训练可复现性时，优先离线缓存。需要完全 On-policy、学生轨迹持续变化时，可使用外部 Teacher Server，但必须固定教师修订并记录全部请求、响应和概率契约。

蒸馏后是否还需要普通 SFT Loss？

通常需要。硬标签或高质量参考答案提供任务目标，教师软分布提供相对偏好。kd_ratio 应通过实验确定；纯 KD 可能继承教师偏差，纯 SFT 又无法充分利用教师概率信息。

---

参考资料

1. NVIDIA NeMo AutoModel, Knowledge Distillation
2. Hugging Face TRL, Distillation Trainer
3. On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
4. MiniLLM: Knowledge Distillation of Large Language Models
5. Distilling the Knowledge in a Neural Network
6. NVIDIA NeMo Framework, Distillation