SOD: Step-wise On-policy Distillation for Small Language Model Agents

📄 arXiv: 2605.07725v1 📥 PDF

作者: Qiyong Zhong, Mao Zheng, Mingyang Song, Xin Lin, Jie Sun, Houcheng Jiang, Xiang Wang, Junfeng Fang

分类: cs.CL, cs.AI

发布日期: 2026-05-08

🔗 代码/项目: GITHUB

💡 一句话要点

提出SOD：一种基于步进式在线策略蒸馏的框架，以提升小语言模型智能体的推理能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 小语言模型 在线策略蒸馏 工具集成推理 强化学习 模型压缩 智能体

📋 核心要点

  1. 现有在线策略蒸馏在工具集成推理中存在失效模式，即错误的工具调用会产生级联效应，导致学生模型与教师模型间的偏差随推理步骤不断放大。
  2. SOD框架通过引入步进式自适应重加权机制，根据每一步的偏差程度动态调整蒸馏强度，从而过滤掉高偏差区域的误导性教师信号。
  3. 实验表明，SOD显著提升了小语言模型在复杂推理任务中的表现，0.6B模型在AIME 2025基准上取得了26.13%的优异成绩，较基线有显著提升。

📝 摘要（中文）

工具集成推理（TIR）因长程交互的不稳定性及模型容量限制，难以扩展至小语言模型。尽管组相对策略优化（GRPO）等强化学习方法可行，但其仅提供稀疏的结果级奖励。近期，在线策略蒸馏（OPD）通过在学生生成的轨迹上提供教师模型的密集Token级监督而受到关注。然而，实验表明将OPD应用于TIR会导致严重的失效模式：错误的工具调用会在后续推理步骤中级联，导致学生与教师模型间的偏差不断放大，使教师的Token级监督变得不可靠。为此，本文提出了SOD（Step-wise On-policy Distillation），这是一种针对小语言模型智能体的步进式在线策略蒸馏框架，通过基于步级偏差自适应调整蒸馏强度。SOD能够在高偏差区域减弱误导性的教师信号，同时在对齐良好的状态下保留密集指导。在数学、科学和代码基准测试中，SOD较次优基线提升高达20.86%，其中0.6B模型在AIME 2025上达到26.13%的准确率，验证了智能体推理能力向轻量级模型迁移的有效性。

🔬 方法详解

问题定义：论文旨在解决小语言模型在工具集成推理（TIR）任务中，因长程交互导致的错误级联问题。现有在线策略蒸馏（OPD）方法在面对错误工具调用时，会强制学生模型模仿教师模型，导致偏差累积，使监督信号失效。

核心思路：核心思想是引入“步进式”的蒸馏控制。通过实时评估学生与教师在每一步推理中的偏差，动态调整蒸馏损失的权重，从而在模型表现不佳时降低对教师信号的依赖，避免错误传播。

技术框架：SOD框架包含学生模型生成轨迹、教师模型提供监督、以及步级偏差评估模块。在训练过程中，系统实时计算学生与教师在当前步骤的输出分布差异，并据此对蒸馏损失进行加权处理。

关键创新：最重要的创新在于提出了基于步级偏差的自适应重加权机制。与传统OPD全局统一的蒸馏强度不同，SOD能够识别并抑制那些因错误工具调用导致的“高偏差”步骤，从而保证训练过程的鲁棒性。

关键设计：该方法通过引入一个偏差度量函数来量化学生与教师的差异，并将其映射为蒸馏损失的缩放因子。在对齐良好的状态下，保持高强度的监督；在偏差较大的状态下，自动衰减蒸馏权重，从而实现对教师信号的智能筛选。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SOD在数学、科学及代码基准测试中表现卓越，较次优基线提升最高达20.86%。特别是在极小参数规模下（0.6B），模型在AIME 2025竞赛级数学题中达到26.13%的准确率，证明了该方法在提升轻量级模型智能体推理能力方面的显著优势。

🎯 应用场景

该研究适用于资源受限的边缘计算设备或移动端智能体开发。通过SOD技术，开发者可以将复杂的推理能力（如数学求解、代码生成、API调用）高效迁移至轻量级模型，显著降低推理成本，并提升模型在复杂任务中的自主决策与工具使用能力。

📄 摘要（原文）

Tool-integrated reasoning (TIR) is difficult to scale to small language models due to instability in long-horizon tool interactions and limited model capacity. While reinforcement learning methods like group relative policy optimization provide only sparse outcome-level rewards. Recently, on-policy distillation (OPD) has gained popularity by supplying dense token-level supervision from a teacher on student-generated trajectories. However, our experiments indicate that applying OPD to TIR leads to a critical failure mode: erroneous tool calls tend to cascade across subsequent reasoning steps, progressively amplifying student-teacher divergence and rendering the teacher's token-level supervision increasingly unreliable. To address this, we propose SOD, a step-wise on-policy distillation framework for small language model agents, which adaptively reweights distillation strength at each step based on step-level divergence. Therefore, SOD can attenuate potentially misleading teacher signals in high-divergence regions while preserving dense guidance in well-aligned states. Experiments on challenging math, science, and code benchmarks show that SOD achieves up to 20.86% improvement over the second-best baseline. Notably, our 0.6B student achieves 26.13% on AIME 2025, demonstrating effective transfer of agentic reasoning to lightweight models. Our code is available at https://github.com/YoungZ365/SOD.