Filter, Then Reweight: Rethinking Optimization Granularity in On-Policy Distillation

TL;DR FiRe-OPD 先过滤低质量轨迹，再对保留轨迹内的 token 进行软加权，联合优化在线策略蒸馏的轨迹和 token 粒度，提升稳定性和性能，在多个基准上显著领先。

问题

问题背景

大型语言模型（LLM）的后训练阶段，On-Policy 蒸馏（OPD） 作为一种缓解分布偏移的关键技术，正从全轨迹 KL 散度对齐向更细粒度的选择性监督范式演进。业界越来越关注如何从学生自生成的推理轨迹中筛选最有效的学习信号。

现有方法局限

当前方法通常孤立地在轨迹级或 token 级做选择：

• 轨迹级过滤：直接丢弃低质量 rollout，但硬性去除可能损失多样性和边缘正样本，且阈值设定对整体性能影响敏感。
• Token 级选择：常见如硬截断（仅对 KL 散度贡献大的 token 做对齐），忽略了 token 间软性差异，导致信息丢失和优化不稳定。 这些方法未能协同考虑两个粒度的信号质量，往往在提升效率的同时牺牲了训练稳定性或最终效果。

问题难点与重要性

核心挑战在于优化粒度的平衡：

1. 学生生成的推理链中，错误 token 可能引发连锁反应，但并非所有错误 token 都同等“有害”；简单过滤或硬性忽略会破坏分布匹配的连续性。
2. 教师监督信号本身带有噪声，需要一种机制动态评估 token 可靠性，而非二值门控。 这种细粒度信号重调对于资源敏感的大模型微调至关重要，直接影响蒸馏后学生模型的推理能力上限。例如，在 strong-to-weak 蒸馏或多教师融合场景中，信号选择策略的优劣会被急剧放大。

行业类比

好比在代码生成模型的强化学习训练中，先根据编译通过率过滤低质量生成样本，再对正确代码块中的关键语法 token 赋予更高权重——一种“先筛选整体质量，再细调局部关注”的实用策略。

核心洞察

• FiRe-OPD 首次在 on-policy 蒸馏中联合轨迹级过滤与 token 级软重加权，将优化粒度从单一维度拓展到双维度协同。此前工作要么只筛轨迹（如拒绝采样），要么只选 token（如硬门控），忽略了这两个层面的交互：低质轨迹中也可能存在有用 token，而轨迹内部 token 重要性天然不均。FiRe-OPD 通过先剔除明显差的轨迹，再在保留样本内用 soft weighting 放大关键 token 的监督信号，既减少了低效样本浪费算力，又避免了硬选择造成的信息丢失与梯度断层，为 OPD 提供了一种更稳定、更细粒度的优化范式。
• 软加权对比硬截断的核心优势在于保留信息的连续性。硬选择方法对 token 做二值 mask，会直接丢弃部分 token 的 KL 监督，这部分 token 可能仍携带微弱的分布信息，完全舍弃会损害优化稳定性，尤其在训练早期容易触发损失剧烈波动。FiRe-OPD 采用 soft reweighting，根据 token 的信息量分配连续权重（如基于 teacher-student 分歧度），让模型仍能接触到所有 token，但将学习重心导向最可迁移的知识。这种“软聚焦”减少了人工阈值调参的脆弱性，在 AIME 2024 和 Miner 等复杂推理 benchmark 上带来显著提升，且对超参数不敏感，工程落地更友好。

方法

输入与总体流程

FiRe-OPD 接收一个学生模型和至少一个教师模型，以及一批用于生成式训练的提示（prompts）。每个 prompt 首先由学生模型进行自回归采样，产生一条完整轨迹（rollout）。教师模型随后对轨迹中的每个 token 给出软标签（通常是对数概率分布），形成基础的蒸馏监督信号。

关键模块 1：轨迹级滤波 (Filter)

对批量生成的轨迹，FiRe-OPD 计算一个反映轨迹整体质量的标量分数（例如基于教师模型下的序列对数概率）。按分数将所有轨迹排序，仅保留前 K% 的高质量轨迹，直接丢弃低质量样本。这一步相当于在轨迹粒度上对训练数据进行筛选，避免学生从错误或噪声严重的自生成样本中学习。

关键模块 2：Token级软重加权 (Reweight)

在保留的高质量轨迹内部，对每个 token 计算一个重要性权重。该权重衡量教师分布与学生分布之间的信息差异（如 KL 散度、Jensen-Shannon 距离，或者直接使用教师概率的熵），反映该 token 提供的监督价值。权重通过一个软映射函数（如 sigmoid 缩放或归一化）平滑地调节 token 损失贡献，而非硬阈值截断。最终蒸馏损失为所有 token 损失的加权和：

loss = Σ w_i * KL(teacher_i || student_i)

其中 w_i 为 token i 的权值，KL 为前向 Kullback-Leibler 散度。

优化与输出

梯度仅从滤波后的轨迹和重加权的 token 回传，更新学生参数。由于使用了软加权，模型仍能从所有保留的 token 中获得连续梯度信号，避免了硬选择造成的梯度截断和信息丢失。

与同类方法的差异

现有 Token 级 OPD 方法（如基于 Top-k 选择、熵门控或硬掩码）倾向于离散地决定一个 token 是否参与训练，这会粗暴地丢弃部分信息，并引入优化不稳定性。FiRe-OPD 的软重加权策略保留了每个 token 的梯度，仅通过权值控制其影响，做到了更细粒度的信号利用，在保持训练平稳的同时提升信息提取效率。

实验

实验设计

FiRe-OPD 在三种知识蒸馏范式下评估：

• Strong‑to‑Weak：强教师模型蒸馏至较弱学生模型。
• Single‑Teacher：单教师标准蒸馏。
• Multi‑Teacher：多教师集成蒸馏，学生同时学习多个教师。

所有实验均基于语言模型推理任务，衡量学生模型在 AIME 2024（数学竞赛）和 Miner（多步推理）上的准确率。与近期聚焦 token 级选择的 OPD 方法对比，FiRe-OPD 引入两阶段调控：先按轨迹质量过滤低质量 rollout，再在保留轨迹内对信息量大的 token 施加软加权，而非硬截断。

关键发现

• 联合优化粒度带来显著增益：Strong‑to‑Weak 场景下 AIME 2024 准确率提升 +6.25，Multi‑Teacher 场景下 Miner 提升 +18.81，说明同时调整轨迹与 token 层级的监督信号能大幅超越仅做 token 选择的方法。
• 软加权缓解信息损失：相比硬 token 选择（直接丢弃部分 token），软重加权保留全部轨迹内 token，仅调整贡献，既保持梯度流稳定，又避免因过早丢弃导致的欠优化。
• 优化稳定性增强：过滤低质量轨迹减少噪声梯度，软加权则精细调控高信息量 token 的 KL 惩罚，两者互补使训练曲线更平滑，减少 loss 尖峰。

与基线对比的深度解读

本文挑选的方法均为近期 token 级 OPD 改进方案（如选择性 KL、策略梯度重加权等），它们只关注 token 维度，却忽视整个轨迹层面的质量差异。FiRe-OPD 的创新在于将“选哪些轨迹”与“轨迹内重点学哪些 token”联合优化：

• 轨迹过滤 相当于在批次内做难例挖掘，剔除错误率高的 rollout，确保监督源干净；
• 软加权 则是连续的表征重要性，比 0/1 选择保留更多结构信息，尤其对需要多个推理步骤的任务（如数学解题），中间关键 token 的温和强调能稳定传递推理链。

这一设计对实际工程的启示：在部署 OPD 微调管道时，不应仅关注 token 级损失裁剪，还需考虑数据质量的前置过滤与 token 重要性的连续刻画，二者搭配可获得更鲁棒的蒸馏效果，且计算开销可控（仅需维护轻量加权模块）。

行业影响

落地场景

FiRe-OPD 所提出的分层优化策略（轨迹过滤 + 软重加权）直接服务于大语言模型的 on-policy 蒸馏 管线。该技术可被广泛应用于需要将强教师模型能力迁移到轻量学生模型的产品场景，例如：

• 移动端 / 浏览器端运行的 对话助手 与 代码补全工具，通过蒸馏获得高性价比的小模型；
• 数学推理与逻辑解题 类应用（如自动判题、辅导系统），需要准确复制推理链；
• 多教师知识融合 场景，如从多个专家模型（逻辑、安全、领域知识）共同蒸馏到一个服务模型。

商业价值

在 降本 维度，轨迹过滤直接丢弃低质量样本，减少无效梯度计算，提升训练效率；软重加权避免硬截断带来的信息损失，使相同训练预算下学生模型性能更高（如在 AIME 2024 数学推理上提升 +6.25 分）。这直接转化为 更低的云计算 / 推理成本 与 更快的模型迭代周期。

在 增收与体验 上，更强的学生模型可支撑更多付费产品功能（如高级代码生成、金融分析问答），并在实时交互中给出更可靠、更连贯的推理过程，降低用户因错误而流失的风险。

与现有工作流的接口

FiRe-OPD 可作为一个轻量级模块插入现有 LLM 训练栈（如基于 PyTorch + Hugging Face 的 SFT / RLHF 流程），无需大幅改动架构：

• 在 student rollout 生成后，增加一个 轨迹评分与过滤 步骤（基于 reward 或 teacher likelihood）；
• 在 KL 散度损失计算时，将原来的 token 级别均匀权重重置为 基于信息量的软权重（如 teacher-student 概率差异、熵等），形成新的蒸馏损失。
• 该方法与主流的 PPO / GRPO 等 on-policy 算法正交，可叠加使用，且超参数（过滤分位数、重加权强度）易于调节。

具体落地 Use Case

1. 金融投资分析助手：将多个大型分析模型（宏观解读、财报分析、风险建模）蒸馏至一个低延迟学生模型，用于实时回答投资顾问的问题。通过轨迹过滤去掉逻辑矛盾或数据错误的思维链，再对关键推理步骤（如计算收益率、风险评估）进行软加权，保证输出严谨性，同时将推理延迟从 5 秒降至 1 秒以内。

2. 代码辅助编程平台（如 GitHub Copilot 风格）：从超大代码生成模型中蒸馏出离线可用的小模型，利用 FiRe-OPD 在训练时自动滤除语法错误或无关补全的轨迹，并对代码核心逻辑行给予更高权重，使小模型在补全准确率和上下文一致率上接近大模型，从而为免费版或边缘设备用户提供更高品质体验。

局限

• - **对教师模型质量的依赖**：FiRe-OPD 的有效性建立在教师提供可靠 logits 的基础上，当教师模型本身在特定任务上表现不佳或与学生的能力差距过大时，过滤轨迹和 Token 重加权可能引入偏差，导致监督信号不可靠，并未充分讨论教师模型误标对学生优化的负面影响。
• - **超参数敏感性与计算开销**：方法引入轨迹过滤百分位、重加权温度系数等多个超参数，尽管论文提供消融实验，但不同任务/模型尺度下需重新调参，缺乏自适应机制；同时轨迹过滤需要额外前向评估（如奖励计算），重加权涉及 token-wise 权重调整，增加了训练过程中的计算负担，文内未给出与硬选择方法、标准 OP D 的 wall-clock 时间对比。
• - **评估场景与模型的局限性**：实验主要集中于数学推理（AIME、Miner）和较小规模模型（如 Qwen2.5-Math 系列），未含更广泛的知识密集型或对话任务，也未在更大规模（>30B）或更主流通用模型（如 Llama-3）上测试，泛化性存疑；对比 baseline 未涵盖当下 SOTA 的在线蒸馏变体（如 ReST-EM、SPIN 等），公平性受限。