Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation

TL;DR 针对采样 token 在线策略蒸馏的多样性瓶颈（pass@k 停滞），提出 IDA-OPD：利用 First-Order Local Entropy Influence 选择性保留熵增更新、收缩熵减更新，仅用采样 token 信息即提升 pass@k 且维持 pass@1。

问题

问题背景

在大模型后训练阶段，on-policy distillation (OPD) 已成为将教师模型能力高效迁移到学生模型的核心手段，并被集成到 GLM-5、KiMi、Qwen3 等主流系统中。为降低全词表教师概率的计算与存储成本，sampled-token OPD 只对 student 生成的 token 获取教师 log-probability，大幅提升了工程可行性。

现有方法局限

然而 sampled-token OPD 常出现 多样性蒸馏失败：学生模型的 pass@1 持续上升，但 pass@k 停滞不前，说明学生只能复现教师的最可能答案，却未能继承教师生成多样解的能力。现有解决方案要么依赖昂贵的全词表 Forward-KL 目标，要么需要额外教师信息或训练信号，难以在仅使用采样 token 教师概率的低信息条件下有效遏制熵收缩。

为什么这个问题难/重要

多样性对于推理任务尤为关键，例如数学题多解生成、代码多方案合成等场景，pass@k 是反映模型覆盖度的核心指标。在采样数据中，每个 token 携带的教师信息有限，难以准确区分哪些梯度更新会扩张熵、哪些会收缩熵；同时推理型后训练对成本敏感，无法简单引入全词表监督。该问题同时触及优化理论、信息约束和工程效率，是当前 LLM 后训练中亟待突破的瓶颈之一。

行业类比

类似 RLHF 中仅用稀疏奖励却需保持策略探索空间：若蒸馏一味强化教师最可能输出，模型会过早收敛到窄分布，损害后续自改进与多候选筛选能力。

核心洞察

• 多样性蒸馏失败的根源在于更新步骤中的局部熵影响：First-Order Local Entropy Influence 将每个 token 位置的熵变化效应分解为 teacher–student log-probability gap 与该位置学生概率分布的一阶局部结构项，从而揭示出 pass@k 停滞并非整体分布对齐不足，而是由特定负影响位置上的熵收缩累积导致。这一视角将问题从全局 KL 散度转移到逐 token 的一阶贡献分解，不同于直接优化 full-vocabulary Forward-KL 或简单采样重加权的方法，为低成本诊断和干预提供了可操作的信号。
• IDA-OPD 的核心在于将熵影响拆成两类更新并分别处理：保留熵扩张更新不变以保证多样性继承，对熵收缩更新应用 divergence-adaptive advantage shrinkage，仅凭 teacher 在采样 token 上的 log-probability 即可判定干预强度。与依赖 full-vocabulary teacher 分布的 Forward-KL 蒸馏不同，这种方法将干预成本压缩到与 vanilla OPD 相同，却能在 pass@k 上匹配最强 teacher-informed 方法，同时维持 pass@1，证明了 OPD 框架下可以通过细粒度更新筛选解决多样性瓶颈，而无需额外 teacher 信息。

方法

输入

• 学生模型 on-policy 采样生成的 token 序列。
• 教师模型在这些采样 token 上的 log-probability（仅需采样 token，无需全词表分布）。
• 学生模型自身在采样位置的局部概率结构（如 top-p 分布、熵梯度等）。

关键模块

1. First-Order Local Entropy Influence：计算每个 token 位置的一阶熵影响代理。该代理将每次梯度更新对熵的影响分解为两部分：教师-学生 log-prob gap（衡量知识差距）和学生局部概率结构（衡量当前分布形态）。得到带符号的值：正值表示该更新会促进熵扩展（增加多样性），负值表示熵收缩（减少多样性）。
2. 更新策略分派：对于熵扩展的 token 位置，保留 vanilla OPD 的 advantage 加权更新，继续促进多样性继承。对于熵收缩的位置，使用 divergence-adaptive advantage shrinkage：根据教师与学生在该 token 上的概率 divergence（例如 log-prob gap 或 KL 散度）自适应地缩放 advantage，降低对低熵行为的强化力度，避免多样性被进一步压缩。

输出

更新后的学生策略，其 pass@k 指标显著提升，成功继承教师多样性，同时 pass@1 基本维持 vanilla OPD 水平。

IDA-OPD 仅依赖教师采样 token 的 log-probability，无需全词表 Forward-KL 目标，因此计算和存储开销远低于传统全词表蒸馏方法。

与同类方法的差异点：IDA-OPD 通过熵影响导向的自适应 shrinkage 在不使用全词表教师信息的前提下，解决了采样 token on-policy 蒸馏的多样性瓶颈。

实验

实验设计

本研究在 reasoning-oriented distillation 场景下验证 IDA-OPD 的多样性传递效果。采用学生生成 token 的 on-policy 蒸馏流程，对比基线包括 vanilla OPD 及最强 teacher-informed 方法。核心评估指标为 pass@k（特别是 k>1 时反映多样性），同时监控 pass@1 确保单样本质量不降。

关键发现

• IDA-OPD 在多个推理任务上稳定提升 pass@k，有效继承教师分布多样性。
• 在保持 pass@1 与 vanilla OPD 相当的前提下，pass@k 不再停滞。
• 相比需要全词表 teacher 概率的同类方法，IDA-OPD 仅依赖 sampled-token 的 teacher log-probability，成本显著降低。

对比解读

Vanilla OPD 的更新方向仅由 advantage 决定，易导致熵收缩，从而多样性丧失。IDA-OPD 引入 First-Order Local Entropy Influence 作为一阶代理，对熵扩张更新予以保留，对熵收缩更新采用 divergence-adaptive advantage shrinkage 替代。这一机制使得学生能够在低资源条件下逼近 teacher-informed 方法的多样性继承效果。实际工程启示：在大规模 LLM 后训练中，可避免昂贵的 full-vocabulary Forward-KL 目标，同时缓解多样性瓶颈，适合以 on-policy 蒸馏为核心的训练管线。

行业影响

落地场景

• 推理模型蒸馏：数学解题、代码生成、逻辑推理等任务需要输出多个合理候选（pass@k）。IDA-OPD 在不增加教师全词表计算成本的前提下提升多样性，适用于构建多候选生成系统，如代码补全工具提供多种实现方案、智能解题助手给出不同思路。
• 对话与内容生成：聊天机器人、创意写作辅助等场景需要回答多样性避免重复；电商平台的商品推荐解释、客服自动回复也可受益于多样化的候选回复，提升用户覆盖。

商业价值

• 降本：仅需教师模型对采样 token 的 log-probability，消除全词表 Forward-KL 目标，显著降低大词表 LLM 蒸馏时的计算与存储开销。
• 体验提升 / 增收：保持 pass@1 的同时提升 pass@k，意味着生成多个候选后重排序或人工筛选更易得到高质量结果，直接改善用户体验；对于付费 API 场景，提供多个可选回复可增加用户粘性与调用量。

与现有工作流的接口

• 即插即用：IDA-OPD 可作为 OPD 训练中的一个损失函数替换模块，只需在原有采样 token 蒸馏流程中修改 advantage 收缩策略，无需改动数据管线、教师推理或分布式训练框架。
• 兼容主流 RL/蒸馏框架：与 PPO、GRPO 等 advantage 计算天然兼容，可集成至现有后训练栈，如 Hugging Face TRL、DeepSpeed 等，适合在已有 reasoning model 蒸馏任务中快速上线验证。

具体场景举例

• 电商智能客服：当用户询问商品对比时，模型生成多个不同角度的推荐理由（如价格、耐用性、售后），提升回答说服力与用户转化。
• 教育解题助手：数学题求解中提供多种解法，增强学习体验，减少“只会一种方法”的挫败感。

局限

• **任务范围较窄**：实验主要针对推理导向蒸馏（如数学 / 代码生成），是否适用于更广泛的指令跟随、对话生成或开放式文本生成任务尚未验证。作者未在非推理任务上充分评估，因此跨任务泛化性存疑，实际部署到不同领域可能需要额外适配。
• **依赖教师采样 token 的 log-probability**：虽然避免了 full-vocabulary 计算，但仍需对教师模型进行前向传播并获取采样 token 的 log-prob，在教师模型规模很大时依然存在不可忽略的推理开销。此外，divergence-adaptive shrinkage 中的阈值或系数等超参数可能对不同教师-学生容量组合敏感，增加了调参复杂度。
• **理论诊断基于一阶近似**：First-Order Local Entropy Influence 是一个 signed 一阶代理量，现实训练动态中高阶项、分布偏移或非局部效应可能导致其不能精确反映真实熵变化。文章未深入讨论该代理在更长训练周期或更激进学习率下的可靠性，以及它与真实熵变化方向不一致时的修正策略。